PRISM: Probability Reallocation with In-Span Masking for Knowledge-Sensitive Alignment

Chenning Xu; Mao Zheng; Mingyang Song

cs.CLApr 2, 2026

PRISM: Probability Reallocation with In-Span Masking for Knowledge-Sensitive Alignment

Authors: Chenning Xu, Mao Zheng, Mingyang Song

Organizations: Large Language Model Department, Tencent, China

Abstract

Supervised fine-tuning (SFT) with token-level hard labels can amplify overconfident imitation of factually unsupported targets, causing hallucinations that propagate in multi-sentence generation. We study an augmented SFT setting in which training instances include coarse sentence-level factuality risk labels and inter-sentence dependency annotations, providing structured signals about where factual commitments are weakly supported. We propose PRISM, a differentiable risk-gated framework that modifies learning only at fact-critical positions. PRISM augments standard SFT with a lightweight, model-aware probability reallocation objective that penalizes high-confidence predictions on risky target tokens, with its scope controlled by span-level risk weights and model-aware gating. Experiments on hallucination-sensitive factual benchmarks and general evaluations show that PRISM improves factual aggregates across backbones while maintaining a competitive overall capability profile. Ablations further show that the auxiliary signal is most effective when used conservatively, and that knowledge masking and model-aware reallocation play complementary roles in balancing factual correction and capability preservation.

Figures & tables

Figure 1: Averaged factual and common performance under different auxiliary weights

\lambda

, shown as

\Delta

relative to the SFT baseline. The figure illustrates a trade-off between factual improvement and preservation of general capability.

Figure 2: Overall pipeline of our risk-gated knowledge-aware SFT (PRISM). Starting from instruction–response pairs, we extract atomic facts and fact relations, derive token-level fact spans and risk weights, and inject these signals into SFT via a model-aware, risk-gated probability reallocation objective.

Model	Size	Method	Factual				Common
Model	Size	Method	HHEM $\uparrow$	PopQA $\uparrow$	TriviaQA $\uparrow$	Avg $\uparrow$	MMLU $\uparrow$	GSM8K $\uparrow$	HumanEval $\uparrow$	IFEval $\uparrow$	Avg $\uparrow$
Qwen3	4B	SFT	94.53	58.54	68.95	74.01	76.19	85.14	79.27	60.81	75.35
		SEAL	92.25	57.76	68.87	72.96	76.01	87.64	79.88	60.81	76.09
		NOVA	94.83	58.89	69.34	74.35	75.28	81.58	81.71	57.86	74.11
		PRISM	95.03	59.54	68.90	74.49	76.93	86.20	81.71	63.40	77.06
Qwen3	8B	SFT	94.83	54.90	69.22	72.98	78.21	88.02	83.54	45.66	73.86

Table 1: Main results. Factual Avg is the mean of (HHEM, PopQA, TriviaQA); Common Avg is the mean of (MMLU, GSM8K, HumanEval, IFEval). Best and second-best averages within each backbone are bold and underlined , respectively.

Variant	Factual				Common
Variant	HHEM $\uparrow$	PopQA $\uparrow$	TriviaQA $\uparrow$	Avg. $\uparrow$	MMLU $\uparrow$	GSM8K $\uparrow$	HumanEval $\uparrow$	IFEval $\uparrow$	Avg. $\uparrow$
SFT ( $\lambda{=}0$ )	94.53	58.54	68.95	74.01	76.19	85.14	79.27	60.81	75.35
$\lambda{=}0.01$	92.64	59.76	68.55	73.65	75.11	85.44	82.93	53.42	74.23
$\lambda{=}0.1$	95.03	59.54	68.90	74.49	76.93	86.20	81.71	63.40	77.06
$\lambda{=}0.5$	94.54	58.83	68.06	73.81	74.86	86.81	79.88	51.76	73.33
$\lambda{=}1.0$	94.54	60.33	67.44	74.10	73.12	86.05	81.71	54.53	73.85

Table 2: Effect of the auxiliary weight

\lambda

on Qwen3-4B. The averaged results suggest a trade-off between factual correction and general capability preservation as the auxiliary signal becomes stronger.

Variant	Factual				Common
Variant	HHEM $\uparrow$	PopQA $\uparrow$	TriviaQA $\uparrow$	Avg. $\uparrow$	MMLU $\uparrow$	GSM8K $\uparrow$	HumanEval $\uparrow$	IFEval $\uparrow$	Avg. $\uparrow$
Ours	95.03	59.54	68.90	74.49	76.93	86.20	81.71	63.40	77.06
w/o model-aware reallocation	94.83	58.54	68.14	73.84	75.54	85.22	77.44	61.74	74.99
w/o knowledge mask	94.04	59.97	68.97	74.33	74.68	84.38	80.49	52.31	72.97

Table 3: Component analysis on Qwen3-4B. “Ours” uses both model-aware reallocation and knowledge masking. “w/o model-aware reallocation” keeps factual masking but removes preference-conditioned reallocation, while “w/o knowledge mask” keeps model-aware updates without masking. The combined design yields the best overall common average while remaining competitive on factual benchmarks.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Training instances
Quantity	Count
Verified instances retained for training	265,568
Fact structure
Fact items (atomic facts extracted)	3,291,756
Fact relations (directed support/dependency edges)	867,202
User-input-supported facts (relations grounded in user input)	324,310

Appendix

Table 4: Summary statistics of the fact-aware preprocessing applied to lmsys_chat_1m_clean .

Explore similar work

Stick to What You Know: A Study of Knowledge-Aligned Supervised Fine-Tuning
Aug 31, 2026Arthur Becker, Jakob Kemmler, David Thulke +3Supervised Fine-Tuning LLM Hallucination Mitigation
AppTek GmbH, Aachen, Germany · Machine Learning and Human Language Technology, RWTH Aachen University, Germany
Why Fine-Tuning Encourages Hallucinations and How to Fix It
Apr 16, 2026Guy Kaplan, Zorik Gekhman, Zhen Zhu +5Supervised Fine-Tuning Continual Learning for LLMs
♡Hebrew University of Jerusalem · ♣Technion – Israel Institute of Technology · ♢University of Illinois Urbana-Champaign +1
Knowing but Not Saying: Preventing Factual Access Failures in LLM SFT via Recall-Anchored Distillation
Aug 21, 2026Haodong Chen, Yadong Wang, Shengtao Wen +2Factual Knowledge in Language Models Domain Adaptation for LLMs
MIIT Key Laboratory of Pattern Analysis and Machine Intelligence, College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics

Aug 31, 2026J/K move · Enter open · S save

Stick to What You Know: A Study of Knowledge-Aligned Supervised Fine-Tuning

Arthur Becker, Jakob Kemmler, David Thulke, Christine Schäfer+2

AppTek GmbH, Aachen, Germany · Machine Learning and Human Language Technology, RWTH Aachen University, Germany

Supervised Fine-Tuning LLM Hallucination Mitigation Factual Knowledge in Language Models

Open paper AskPDF arXiv

Supervised fine-tuning (SFT) trains a base language model to imitate target responses, and these targets may require knowledge the base model has not robustly internalized. We study this as a source of hallucinations and frame a group of mitigation methods as knowledge-aligned SFT: constraining SFT training targets to the base model's parametric knowledge. Under a unified setup, we compare existing generation-based and estimation-based knowledge-alignment methods and introduce two new variants: Evidence Rewrite, which verifies base-model generations using external evidence, and Recall Rewrite, which retains claims only when they can be consistently recalled by the base model. Experiments with Qwen 3 4B and OLMo 3 7B show that knowledge-aligned SFT can reduce factual hallucinations on WildHalu and Biography while largely preserving general capabilities. Recall Rewrite yields the strongest factuality gains and improves refusal behavior on UnknownBench. It thereby confirms that SFT targets beyond the base model's knowledge drive hallucination behavior.

Abstract

Explore similar work

Stick to What You Know: A Study of Knowledge-Aligned Supervised Fine-Tuning

Why Fine-Tuning Encourages Hallucinations and How to Fix It

Knowing but Not Saying: Preventing Factual Access Failures in LLM SFT via Recall-Anchored Distillation