PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zitong Huang, Gustavo Lucas Carvalho, Deqing Fu, Robin Jia

arXiv:2607.17524v1cs.CLcs.LG

TL;DR

LLMs still struggle with hallucinations and factuality, while existing training approaches face complexity or limited evidence under distribution shift. TOPL reframes off-policy post-training as token-level correctness prediction and achieves strong out-of-distribution factuality across document summarization, with effective transfer to machine translation.

  • Problem

    LLMs continue to struggle with hallucinations and factuality, motivating training methods that improve faithful generation under distribution shift.

  • Method

    TOPL trains LoRA-adapted language models to classify each response token as correct or incorrect rather than directly optimizing next-token generation.

  • Results

    TOPL achieves strong out-of-distribution factuality across 11 document-summarization datasets, transfers effectively to machine translation, and outperforms sequence-level and token-level supervision methods.

  • Takeaways & Limitations

    TOPL’s fine-grained correctness objective reshapes internal representations, with analyses linking LoRA components to token classification and conditional steering.

  • Takeaways & Limitations

    Cleaner token-level perturbations may be needed to reduce supervision noise, and the causal relationship between representation structure and generalization remains unresolved.

Abstract

from arXiv · show

We propose Token-Level Off-Policy Labeling (TOPL), an off-policy training paradigm that reframes post-training as a token-level correctness prediction task. Our key intuition is that by training the model to distinguish good and bad tokens in a response, we naturally guide the model towards generating good tokens, while avoiding the pitfalls that come with directly training the model to generate off-policy tokens. Experiments on document summarization tasks show that TOPL achieves strong out-of-distribution generalization across 11 datasets against a diverse set of sequence-level and token-level baselines. We further demonstrate that TOPL transfers effectively to machine translation, suggesting that its benefits generalize across different faithful generation tasks. Through ablation studies, we confirm that our token-level learning signal is critical to good performance; sequence-level analogues do not confer similar benefits. Finally, we show that TOPL induces interpretable model updates: the LoRA adapters learned through TOPL function as linear classification heads and steering vectors.

1 Introduction

TOPL reframes off-policy post-training as token-level correctness classification, using LoRA to learn signals that guide faithful generation. It achieves strong out-of-distribution performance across text-generation tasks and motivates this effectiveness through conditional steering vectors.

  • Method: TOPL reframes post-training as token-level correctness prediction instead of directly optimizing next-token prediction.It trains a binary classification objective with Low-Rank Adaptation (LoRA).
  • Method: TOPL replaces the language-modeling head with a binary reward head that predicts whether each response token is factual.The training pipeline processes documents, queries, and responses containing factual and non-factual tokens to obtain token-level hidden representations.
  • Results: Sequence-level analogues do not provide similar benefits, highlighting the importance of fine-grained token-level training signals.TOPL also transfers effectively to machine translation, suggesting benefits across different faithful generation tasks.
  • Interpretability: TOPL’s LoRA components support conditional steering by extracting token-level correctness signals and using them to weight corresponding steering vectors.The analysis repurposes LoRA-A as a classifier of factual and nonfactual tokens and examines how LoRA-B vectors use the extracted signal.
  • Results: TOPL achieves strong out-of-distribution performance relative to token-level and sequence-level baselines across text-generation tasks.The primary evaluation focuses on document summarization with a synthetic FAVA-based training dataset and comparisons including SFT, DPO, Token-Level DPO, TLDR, and Token-Level Unlikelihood Training.

2 Related Work

Related work covers preference-based post-training, contrasting complex on-policy methods with lightweight off-policy approaches, and discusses inference-time model steering through activation-space directions.

  • Reinforcement Learning and Preference Optimization: RLHF is a standard language-model post-training approach using preference or reward signals, with both on-policy and off-policy methods.On-policy methods have shown strong performance in improving generalization but rely on online sampling and repeated rollouts.
  • Reinforcement Learning and Preference Optimization: On-policy methods require online sampling and repeated rollouts, introducing substantial training complexity, whereas off-policy methods such as DPO are more lightweight.DPO is identified as an off-policy method within preference optimization.
  • Model Steering: Model steering controls language-model behavior at inference time by modifying internal representations, often using directions in activation space to shift outputs toward desired behaviors.Representation-based steering identifies specific activation-space directions for behavior control.

3 Method

TOPL reframes post-training as token-level correctness prediction, labeling response tokens as faithful or corrupted rather than directly predicting the next token. It trains a lightweight reward head with token-level supervision and transfers the learned LoRA updates back into the generative model for inference.

  • TOPL formulation: TOPL formulates post-training as binary classification of whether each response token is correct or corrupted.This contrasts with next-token prediction based on one-hot labels for a single correct token.
  • Token-level supervision: A ground-truth summary is treated as a faithful positive response, while controlled insertions, deletions, or substitutions create unfaithful perturbed summaries.FAVA supplies model-generated summaries with fine-grained edits simulating typical hallucinations.
  • Token-level supervision: Each perturbed-summary token receives z_k = 1 when faithful to the original summary and z_k = 0 when part of an edit-introduced hallucinated span.The labels define the good-token and bad-token supervision used by TOPL.
  • Training: TOPL predicts p_k = P(z_k = 1 | D, t≤k) using a lightweight reward head attached at an intermediate layer instead of the decoding head.The model applies binary cross-entropy over response tokens and uses LoRA on selected layers rather than updating all parameters.
  • Low-Rank Model Merging: After training, the reward head and added normalization layer are discarded, and the learned LoRA update is merged into the truncated model for standard generation.Inference uses the original decoding head and ordinary forward propagation.

4 Experiments

TOPL is evaluated primarily on document summarization under distribution shift, where it ranks among the strongest methods and achieves the best average performance across backbone models. Layer-placement and label-sensitivity analyses, together with machine-translation experiments, examine the conditions underlying its gains and cross-task generalization.

  • Experimental setup: TOPL uses Qwen3-8B, Llama-3.1-8B, and Gemma-3-4B backbones with rank-4 LoRA and scaling α = 8.LoRA is inserted across layers 0–29 for Qwen3-8B and layers 0–27 for Llama-3.1-8B and Gemma-3-4B.
  • Document summarization: On 11 AggreFact datasets, TOPL is compared with SFT, DPO, SOPL, TDPO, TLDR, and Unlikelihood training using factuality scores.Models are trained on FAVA and evaluated out of distribution with up to 300 unique documents per dataset; higher scores indicate better factual consistency.
  • Document summarization: Under distribution shift, TOPL achieves the best average performance across backbone models, scoring highest on Qwen and Gemma while remaining competitive with the strongest baseline on Llama.The gains cannot be explained solely by generating shorter or more conservative summaries, and TOPL outperforms other existing token-level baselines overall.
  • Layer ablation: Applying LoRA to middle layers yields the largest factuality improvements, early-layer combinations enhance performance, and final-layer inclusion provides limited or negative gains.This conclusion comes from a sliding-window search over transformer layers.
  • Label sensitivity: Label-distribution skew has minor effects, whereas using only good tokens or dense balanced labels changes performance substantially; random labeling degrades Llama-3.1-8B and Gemma-3-4B.Qwen3-8B instead improves over its base model under random labels, consistent with its reported robustness to spurious supervision.
  • Machine translation: In machine translation across four MLQE-PE language pairs, TOPL achieves the strongest performance among compared methods on out-of-distribution FLORES-Plus benchmarks.Evaluation uses XCOMET on held-out MLQE-PE and FLORES-Plus, with word-level quality annotations converted into token-level supervision.

5 Analysis

The analysis interprets TOPL with LoRA as learning adaptive, token-level steering mechanisms: LoRA-A detects factuality-related concepts, while LoRA-B supplies corresponding steering directions. TOPL yields stronger token separability, coherent factuality steering, and better OOD factuality than SFT.

  • LoRA-based steering interpretation: LoRA-A encodes projections onto learned concept directions, while LoRA-B defines corresponding steering directions for adaptive token-level conditioning.LoRA-A can detect factual or non-factual patterns, activating the associated LoRA-B direction.
  • LoRA-A separability: TOPL exhibits substantially stronger separability between good and bad tokens in the LoRA-A subspace than SFT, especially in middle-to-late layers.Separability is quantified using AUROC across layers with LoRA modules.
  • LoRA-B steering: For TOPL, factuality first increases and then decreases as λ varies from negative to positive values, whereas SFT shows much weaker intervention sensitivity.Negative λ steers toward nonfactual directions, positive λ toward factual representations, and an optimal magnitude balances these effects.
  • OOD generalization: TOPL combines higher AUROC in LoRA-A with stronger factuality on OOD datasets, consistent with stronger separability predicting better generalization under distribution shift.LoRA-A is described as a low-rank projection that conditions hidden states on token-level factuality.

6 Discussion and Conclusion

TOPL reframes off-policy post-training as token-level correctness prediction, improving factuality, translation transfer, and out-of-distribution robustness. Its fine-grained supervision distinguishes it from sequence-level and other token-level methods, while future work targets cleaner perturbations, reasoning, and behavioral control.

  • TOPL reframes language-model alignment as token-level correctness prediction and consistently improves factuality, machine-translation generalization, and out-of-distribution robustness.These gains hold against both sequence-level and token-level supervision methods.
  • TOPL localizes supervision to individual tokens, providing substantially richer information than sequence-level supervision, but its gains cannot be attributed solely to token-level supervision.Comparisons with other token-level methods motivate this distinction.
  • Future work includes cleaner token-level perturbations, extending TOPL to reasoning tasks, and varying token-level labels for behavioral control beyond factuality.Cleaner perturbations may reduce supervision noise and further enhance performance; reasoning applications could capture intermediate-step errors.

A Appendix · A.1 Full-Summary Evaluation

The appendix introduces full-summary evaluation with Bespoke-MiniCheck-7B and finds strong agreement with sentence-level scoring. Under this protocol, TOPL remains strongest across the evaluated backbone models, especially on OOD datasets.

  • A.1 Full-Summary Evaluation: Full-summary evaluation scores each generated summary as one sequence using Bespoke-MiniCheck-7B.This contrasts with the main paper’s sentence-level protocol, which decomposes summaries into individual sentences.
  • A.1 Full-Summary Evaluation: Sentence-level evaluation generally assigns higher scores than full-summary evaluation across 300 samples.The comparison evaluates the same 300 samples under both protocols.
  • A.1 Full-Summary Evaluation: Sentence-level evaluation maintains strong rank correlation with full-summary evaluation, with Spearman = 0.79.The higher sentence-level scores may result from well-formed or factual sentences compensating for less accurate summary parts.
  • A.1 Full-Summary Evaluation: Under in-distribution full-summary evaluation, TOPL achieves the best performance on Qwen3-8B and LLaMA-3.1-8B.The evaluation covers a representative subset of methods.
  • A.1 Full-Summary Evaluation: Under out-of-distribution full-summary evaluation, TOPL consistently achieves the strongest performance across all backbone models.Figure 8 reports factuality scores across backbone models in both ID and OOD settings.
  • A.1 Full-Summary Evaluation: Full-summary evaluation trends closely mirror those observed under sentence-level evaluation.This similarity is reported for the evaluated methods across in-distribution and out-of-distribution datasets.

A.2 Summary Length Analysis

The analysis tests whether TOPL’s factuality gains are merely a consequence of generating shorter or more conservative summaries. Results across models show that summary length alone cannot explain the improvements.

  • Summary Length Analysis: TOPL often produces shorter summaries than SFT on Qwen and Llama, but this pattern is not universal.On Gemma, TOPL generates slightly longer summaries than SFT while achieving higher factuality.
  • Summary Length Analysis: On Gemma, TOPL generates slightly longer summaries than SFT while still achieving higher factuality.
  • Summary Length Analysis: Compared with the base model, TOPL generates slightly more sentences on both Qwen and Gemma yet yields markedly larger factuality improvements.These findings suggest that summary length alone cannot explain TOPL’s observed gains.

A.3 Model Training Setup and Hyperparameters

This section provides the model architectures, training setups, and hyperparameter configurations for TOPL and six comparison methods. These configurations are organized across Tables 3–9 and span different backbone models.

  • Training Setup: Training setups and hyperparameter configurations are provided for TOPL, SFT, SOPL, DPO, TLDR, Token-level Unlikelihood Training, and Token-Level DPO.The configurations are presented alongside model architectures in Tables 3–9.
  • TOPL: Table 3 reports TOPL training hyperparameters across different backbone models.
  • Baselines: Tables 4–9 report training hyperparameters for SFT, SOPL, DPO, TLDR, Token-Level Unlikelihood Training, and Token-Level DPO across different backbone models.

A.4 Detailed Results on In-Distribution and Out-of-Distribution Benchmarks … A.7 Detailed Machine Translation Results

The appendix reports detailed factuality results across ID and OOD benchmarks, label-sensitivity statistics, LoRA analyses for DPO and SOPL, and machine-translation results where TOPL performs best in both settings.

  • A.4 Detailed Results on In-Distribution and Out-of-Distribution Benchmarks: Tables 10 and 11 provide detailed in-distribution and out-of-distribution factuality results across backbone models and 11 AggreFact datasets.Table 10 covers ID factuality scores, while Table 11 reports OOD scores averaged across the 11 datasets.
  • A.5 Label Sensitivity Dataset Statistics: Table 12 reports supervised-token counts and positive factual versus negative hallucinated label proportions for each label-sensitivity configuration.The statistics characterize the datasets used in the label sensitivity analysis.
  • A.6 LoRA-A and LoRA-B Analysis for DPO and SOPL: DPO and SOPL are analyzed through LoRA-A separability heatmaps and LoRA-B steering curves under different scaling factors λ.The analysis complements the TOPL and SFT results with component-level diagnostics.
  • A.6 LoRA-A and LoRA-B Analysis for DPO and SOPL: DPO and SOPL show substantially weaker LoRA-A separability than TOPL across layers and modules, indicating less discriminative structure between factual and non-factual tokens.Both methods nevertheless retain clear LoRA-B steering effects under scaling.
  • A.6 LoRA-A and LoRA-B Analysis for DPO and SOPL: LoRA-B steering is not unique to TOPL; the main distinction is stronger LoRA-A structure, which may determine when learned steering directions are activated.The analysis proposes that effective steering requires both useful directions and appropriate activation conditions.
  • A.6 LoRA-A and LoRA-B Analysis for DPO and SOPL: TOPL’s stronger LoRA-A separability may enable more consistent activation of appropriate steering directions, improving factuality and OOD robustness after merging.DPO and SOPL learn usable LoRA-B directions, but weaker LoRA-A structure may limit their utilization.
  • A.7 Detailed Machine Translation Results: Table 13 reports complete Qwen3-8B machine-translation results under ID and OOD settings, with higher values indicating better translation quality.TOPL achieves the best performance in both settings.
Loading 2607.17524v1…