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

2026-07-19

AI Summary

This work addresses the issue of unfaithful generation by large language models under distribution shift by proposing TOPL, a token-level off-policy learning framework that reframes post-training as a token-wise correctness prediction task. By introducing interpretable token-level reward signals, TOPL guides the model to distinguish high-quality from low-quality tokens in its outputs, thereby avoiding the pitfalls of directly imitating off-policy samples. This approach enables the first interpretable, token-granularity model updates, where the learned LoRA adapter simultaneously functions as both a classification head and a guidance vector. Evaluated across 11 summarization datasets, TOPL significantly outperforms existing baselines and demonstrates successful transfer to machine translation, confirming its generality and superiority across diverse faithfulness-oriented generation tasks.

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.

Problem

• off-policy learning
• faithful generation
• distribution shift
• token-level learning
• out-of-distribution generalization

Innovation

• Token-Level Learning
• Off-Policy Training
• Faithful Generation
• Distribution Shift
• LoRA Adapters