Token-Level Off-Policy Learning for Faithful Generation Under Distribution Shift
2026-07-20 • Computation and Language
Computation and LanguageMachine Learning
AI summaryⓘ
The authors propose Token-Level Off-Policy Labeling (TOPL), a new training method where a model learns to tell which tokens (words or pieces of words) in a response are good or bad. This helps the model get better at generating good tokens without the difficulties of directly training on off-policy tokens. They tested TOPL on summarization and machine translation tasks, finding it works well across different datasets and tasks. Their studies show that focusing on token-level feedback is important for the method's success, and that the changes TOPL makes to the model are understandable.
off-policy trainingtoken-level predictionpost-trainingdocument summarizationmachine translationout-of-distribution generalizationablation studyLoRA adapterslinear classificationsteering vectors
Authors
Zitong Huang, Gustavo Lucas Carvalho, Deqing Fu, Robin Jia
Abstract
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.