ShiQ: Bringing back Bellman to LLMs
Pierre Clavier, Nathan Grinsztajn, Raphael Avalos, Yannis Flet-Berliac, Irem Ergun, Omar D. Domingues, Eugene Tarassov, Olivier Pietquin, Pierre H. Richemond, Florian Strub, Matthieu Geist
TL;DR
ShiQ introduces a Q-learning–style, KL-regularized loss for fine-tuning LLMs by interpreting logits as Q-values within a regularized MDP. It advances token-level off-policy learning through three LLM-specific refinements—easing sampling, reward-shaping–based initialization, and a multi-step propagation—culminating in the Shifted-Q loss. The authors provide theoretical guarantees that the optimizer converges to the optimal policy under suitable data support and show how ShiQ can sample directly from learned logits at inference without extra networks. Empirically, ShiQ demonstrates competitive performance on synthetic tasks and real-world LLM benchmarks, excelling particularly in multi-turn scenarios like BFCL-V3 by exploiting token- and trajectory-level rewards. Overall, ShiQ offers an efficient, theoretically grounded framework for offline RL-style fine-tuning of LLMs with strong practical impact for reward-based alignment and tool-use in multi-turn interactions.
Abstract
The fine-tuning of pre-trained large language models (LLMs) using reinforcement learning (RL) is generally formulated as direct policy optimization. This approach was naturally favored as it efficiently improves a pretrained LLM, seen as an initial policy. Another RL paradigm, Q-learning methods, has received far less attention in the LLM community while demonstrating major success in various non-LLM RL tasks. In particular, Q-learning effectiveness comes from its sample efficiency and ability to learn offline, which is particularly valuable given the high computational cost of sampling with LLMs. However, naively applying a Q-learning-style update to the model's logits is ineffective due to the specificity of LLMs. Our core contribution is to derive theoretically grounded loss functions from Bellman equations to adapt Q-learning methods to LLMs. To do so, we carefully adapt insights from the RL literature to account for LLM-specific characteristics, ensuring that the logits become reliable Q-value estimates. We then use this loss to build a practical algorithm, ShiQ for Shifted-Q, that supports off-policy, token-wise learning while remaining simple to implement. Finally, we evaluate ShiQ on both synthetic data and real-world benchmarks, e.g., UltraFeedback and BFCL-V3, demonstrating its effectiveness in both single-turn and multi-turn LLM settings