Token-Level Off-Policy Learning for Faithful Generation Under Distribution Shift

Technical Summary: Token-Level Off-Policy Learning for Faithful Generation

Problem Statement
Large Language Models (LLMs) frequently struggle with hallucinations and factual inconsistencies, particularly when generating content under distribution shifts (Out-of-Distribution or OOD settings). While Reinforcement Learning from Human Feedback (RLHF) and off-policy methods like Direct Preference Optimization (DPO) have improved alignment, they often rely on sequence-level supervision. This coarse granularity limits their ability to pinpoint specific factual errors. Furthermore, on-policy methods (e.g., GRPO) suffer from high computational complexity due to online sampling, while existing off-policy token-level approaches often treat reward modeling as a separate component rather than a unified training objective. The paper identifies a need for a method that retains the simplicity of off-policy training while providing fine-grained, token-level supervision to improve factuality and OOD generalization.

Methodology: Token-Level Off-Policy Labeling (TOPL)
The authors propose Token-Level Off-Policy Labeling (TOPL), a training paradigm that reframes post-training as a binary classification task at the token level rather than a next-token prediction task.

• Data Construction: The method utilizes a dataset (based on FAVA) containing documents, ground-truth summaries, and perturbed ("corrupted") summaries. Corruptions are introduced via token insertions, deletions, or substitutions to simulate hallucinations. Each token in the response is assigned a binary label: $1$ for faithful (correct) and $0$ for non-faithful (corrupted).
• Training Objective: Instead of optimizing a language modeling head for next-token prediction, TOPL attaches a lightweight reward head to an intermediate layer of the pre-trained model. The model is trained to predict the probability that a specific token is factually consistent with the input context.
• Architecture & LoRA: The training is performed using Low-Rank Adaptation (LoRA). The model learns to distinguish good and bad tokens via a binary cross-entropy loss. Crucially, the reward head is discarded after training; the learned LoRA weights are merged into the base model to create a generative model that utilizes the original decoding head.
• Interpretation as Conditional Steering: The authors interpret TOPL through the lens of conditional steering. They hypothesize that the LoRA-A matrix acts as a conditional projection mechanism (classifying token correctness), while the LoRA-B matrix defines steering vectors that shift hidden states toward or away from factual behavior based on the classification signal.

Key Contributions

1. Novel Training Paradigm: Introduction of TOPL, which unifies token-level reward learning and policy optimization into a single off-policy framework, avoiding the complexity of online sampling.
2. Mechanistic Insight: An analysis demonstrating that TOPL induces interpretable model updates where LoRA components function as linear classification heads (LoRA-A) and steering vectors (LoRA-B).
3. Empirical Validation: Comprehensive evaluation showing that TOPL outperforms both sequence-level baselines (SFT, DPO, SOPL) and existing token-level baselines (TDPO, TLDR, Unlikelihood Training) in OOD settings.

Experimental Results

• Document Summarization: Evaluated on 11 OOD datasets from AggreFact using Qwen3-8B, Llama-3.1-8B, and Gemma-3-4B. TOPL achieved the highest average factuality scores across all backbone models.
  • Ablation on Granularity: Sequence-level analogues (SOPL) failed to match TOPL's performance, confirming that fine-grained token-level signals are critical.
  • Layer Selection: Applying LoRA to middle layers yielded the best results; updating final layers provided limited gains or degraded performance, likely due to disruption of pre-trained generative behaviors.
  • Label Sensitivity: The method relies on contrastive supervision (presence of both good and bad tokens). While the ratio of good/bad tokens was not the dominant factor, training with only one type of token led to performance degradation.
• Machine Translation: TOPL was successfully transferred to machine translation tasks (MLQE-PE and FLORES-Plus), achieving state-of-the-art OOD performance, suggesting the benefits generalize across faithful generation tasks.
• Representation Analysis:
  • Separability: TOPL-trained models exhibited significantly higher separability (AUROC) between factual and non-factual tokens in the LoRA-A subspace compared to SFT.
  • Steering: Post-hoc interventions using LoRA-B vectors showed that TOPL models respond coherently to steering strength, whereas SFT models showed weak sensitivity.
  • Correlation: Higher LoRA-A separability correlated strongly with better OOD generalization scores.

Significance and Claims
The paper claims that TOPL offers a robust solution for improving model factuality under distribution shifts by shifting the training objective from "what token comes next" to "is this token correct." The authors argue that this approach avoids the pitfalls of directly training models to generate off-policy tokens. Instead, by training the model to distinguish correctness, the method naturally guides the generation process toward faithful outputs.

The significance of the work lies in its demonstration that:

1. Token-level supervision is superior to sequence-level supervision for OOD generalization in faithful generation.
2. The specific objective matters: Simply having token-level labels is insufficient; the binary classification objective used in TOPL creates a specific internal representation structure (separable LoRA-A and steerable LoRA-B) that enhances robustness.
3. Interpretability: The method provides a clear mechanistic explanation for its success, linking the learned LoRA updates to conditional steering mechanisms that can be analyzed and potentially controlled.

The authors remain modest regarding causality, noting that while they observe a strong correlation between LoRA-A separability and OOD performance, the strict causal relationship between representation structure and generalization remains an open question for future work.