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

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

Motivation and Problem Setting

This work addresses the persistent challenge of hallucination in LLMs fine-tuned via standard SFT, specifically the tendency to produce overconfident, factually unsupported claims. Existing hallucination mitigation strategies—including retrieval augmentation, post-hoc verification, and preference-based alignment—either impose significant training/inference-time complexity or fail to intervene at the crucial stage of SFT where knowledge misalignment is induced. The paper asserts that SFT with hard targets can inadvertently reinforce unsupported content and proposes an alternative approach: injecting structured factuality risk signals directly into the training objective, without modifying inference pipelines or requiring external retrieval during learning.

Method: Model-Aware, Risk-Gated Probability Reallocation (PRISM)

PRISM extends standard SFT by leveraging coarse-grained, instance-level atomic fact and inter-sentence dependency annotations extracted from training data. The central contributions are:

• Fact-Active Span Extraction: Atomic facts align to token-level spans where specific factual commitments occur. Sentences are assigned factuality risk scores and inter-sentence dependency relations, allowing for the propagation of risk across related claims.
• Model-Aware Gating: Probability reallocation is applied only where (i) a token lies within a fact-critical span, (ii) the model's current preference aligns with the reference token, and (iii) the redistribution does not shift the model's preferred continuation away from the reference token.
• Risk-Weighted Complement Loss: The complement objective penalizes high label token probabilities on risky spans. The risk-aware weight modulates the penalty, while model-aware gating ensures that updates do not induce sequence-level preference mismatches.
• Final Optimization: The overall objective combines standard cross-entropy with risk-gated, model-aware penalties, controlled by a single scalar $\lambda$ regulating the strength of factual correction.

The pipeline is summarized in Figure 1.

Figure 1: The PRISM pipeline extracts atomic facts and their dependencies from training data, constructs fact-active spans and risk weights, and injects risk signals into SFT via consistency-gated probability reallocation.

Empirical Evaluation

Datasets and Experimental Setup

The primary training data is lmsys_chat_1m_clean, processed to obtain fact-level supervision (3.29M facts, 867K dependency edges, 3.14M span annotations over 265K instances). Evaluations target both hallucination-sensitive factual benchmarks (HHEM, PopQA, TriviaQA) and broad "common" benchmarks (MMLU, GSM8K, HumanEval, IFEval). Experiments are run on multiple backbones: Qwen3-4B/8B and Llama 3.1-8B.

Main Results

PRISM consistently raises factual aggregate scores relative to vanilla SFT across all backbones, without marked regressions in general capability (Table~\ref{tab:main_results}). Gains on Qwen3-8B and Llama 3.1-8B are substantial, and increases in factual metrics on Qwen3-4B—already near ceiling in baseline—are modest, indicating minimal negative interference. Notably:

• PRISM outperforms or matches stronger baselines that use abstention learning (SEAL), sample filtering (NOVA), or naive factual masking.
• On Qwen3-8B, the Common Avg improvement is primarily due to increased instruction-following scores (IFEval), not at the expense of reasoning or code generation ability.

The performance trade-off induced by different $\lambda$ values is visualized in Figure 2, illustrating the tension between factual improvement and general capability preservation.

Figure 2: Factual and common performance (relative to baseline) for varying auxiliary weights $\lambda$; stronger auxiliary signals enhance factuality but can degrade general ability, highlighting a precise tunable trade-off.

Ablation Studies

• Auxiliary Weight $\lambda$: Moderate $\lambda$ values yield the optimal factuality-capability balance; overly aggressive reallocation can harm generalization.
• Component Disentanglement: Removing model-aware gating or knowledge masking each degrades either factual or common capacity. The concurrency of both mechanisms is required for optimal, robust performance across metrics.

Theoretical and Practical Implications

This work provides an actionable path for practitioners seeking to mitigate knowledge hallucination at the SFT stage with minimal engineering overhead. The model-aware gating ensures that factual updates are only introduced where they do not create sequence-level misalignment, a failure mode observed in unselective interventions. Empirical results and gradient analysis indicate that selective application of the complement loss avoids catastrophic degradation of out-of-domain or reasoning behaviors. The approach generalizes across model families and scales, with minimal hyperparameter complexity (the principal axis being the auxiliary weight $\lambda$).

Further, by decoupling factual guidance from external retrieval or reward modeling, PRISM makes knowledge-sensitive alignment feasible for production deployment, where auditability and stability are essential.

Future Directions

Potential extensions include:

• Adaptive Scheduling: Dynamic adjustment of $\lambda$ by domain, instance, or during curriculum, based on model uncertainty or factual confidence.
• Enhanced Fact Annotation: Improved or automated fact extraction and dependency parsing could further localize and refine risk attribution.
• Generalization to Multiturn Dialog: Expanding fact-structure extraction and dependency modeling beyond single-turn data.

Conclusion

PRISM operationalizes knowledge-sensitive alignment through structured, model-aware SFT interventions, enforcing a principled balance between factual reliability and general instruction-following. Its differentiable risk-gated correction is empirically validated across architectures and domains, and its ablation studies provide actionable guidance for tuning and deployment. This framework advances the quest for interpretable, stable, and robust factuality in large-scale language generation.