Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Chenning Xu, Mao Zheng, Mingyang Song

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 \textbf{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.

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

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 \textbf{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.

Paper Structure

This paper contains 35 sections, 15 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Hallucination Mitigation
  4. Hallucination Mitigation during SFT
  5. Methodology
  6. Setup and Notation
  7. Fact-Aware Signals from SFT Data
  8. Fact-active mask.
  9. Risk weights with dependency propagation.
  10. Training Objective
  11. Standard SFT loss.
  12. Risk-gated probability mass reallocation.
  13. Final objective.
  14. Gradient Analysis
  15. Training Procedure
  16. ...and 20 more sections

Figures (2)

  • 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.