From $r$ to $Q^*$: Your Language Model is Secretly a Q-Function (2404.12358v2)

Abstract: Reinforcement Learning From Human Feedback (RLHF) has been critical to the success of the latest generation of generative AI models. In response to the complex nature of the classical RLHF pipeline, direct alignment algorithms such as Direct Preference Optimization (DPO) have emerged as an alternative approach. Although DPO solves the same objective as the standard RLHF setup, there is a mismatch between the two approaches. Standard RLHF deploys reinforcement learning in a specific token-level MDP, while DPO is derived as a bandit problem in which the whole response of the model is treated as a single arm. In this work we rectify this difference. We theoretically show that we can derive DPO in the token-level MDP as a general inverse Q-learning algorithm, which satisfies the Bellman equation. Using our theoretical results, we provide three concrete empirical insights. First, we show that because of its token level interpretation, DPO is able to perform some type of credit assignment. Next, we prove that under the token level formulation, classical search-based algorithms, such as MCTS, which have recently been applied to the language generation space, are equivalent to likelihood-based search on a DPO policy. Empirically we show that a simple beam search yields meaningful improvement over the base DPO policy. Finally, we show how the choice of reference policy causes implicit rewards to decline during training. We conclude by discussing applications of our work, including information elicitation in multi-turn dialogue, reasoning, agentic applications and end-to-end training of multi-model systems.

• The paper introduces a theoretical framework where DPO acts as an inverse Q-learning algorithm, linking token-level MDPs with reinforcement learning principles.
• It demonstrates empirically that DPO-trained models achieve effective token-level credit assignment, improving performance metrics in beam search experiments.
• The paper discusses practical implications for enhancing multi-turn dialogues, end-to-end generative systems, and autonomous agent behaviors in AI.

From $r$ to $Q^*$: Your LLM is Secretly a Q-Function

Introduction to RLHF and DPO

Reinforcement Learning from Human Feedback (RLHF) plays an essential role in aligning LLMs with human intent. Traditional RLHF methods leverage reinforcement learning (RL) frameworks like PPO to fine-tune models based on reward signals derived from human feedback. Direct Preference Optimization (DPO) emerges as an alternative that simplifies this setup by aligning models directly through preference data without an intermediate reward function. This paper introduces theoretical insights that connect DPO with token-level Markov Decision Processes (MDPs) in LLMs, proposing that DPO operates as an inverse Q-learning algorithm, satisfying the Bellman equations.

DPO and Token-Level MDP

In traditional RLHF, LLMs model token sequences as trajectories in an MDP where states are sequences of tokens, actions are vocabulary entries, and rewards are derived from human feedback. However, classical RLHF applies these rewards sparsely at terminal states, driving the overall optimization with policy gradient techniques. Contrastingly, DPO frames the problem in a contextual bandit setting—a paradigm where sequences are treated as single decisions rather than token steps. The novel derivation connects DPO's bandit formulation to the token-level MDP, implying that DPO implicitly learns a per-token reward, forming a Q-function over tokens (Figure 1).

Figure 1: Credit assignment in DPO based on answer-level feedback. Each token is colored corresponding to the DPO implicit reward as expressed in the provided equations.

Empirical Insights and Theoretical Validation

The authors empirically demonstrate that DPO-trained models can manifest token-level credit assignment akin to what RLHF might achieve with dense rewards. This insight is crucial as it forms the theoretical foundation for interpreting the learning dynamics of LLMs under the DPO framework. Furthermore, the paper explores the equivalence between traditional search-based algorithms and likelihood-based search performed over a DPO policy, providing an empirical basis through beam search experiments that highlight meaningful improvements in model performance (Figure 2).

Figure 2: Model performance using beam search, illustrating win rates and verbosity issues beyond five beams.

Performance Degradation and Implicit Rewards

The phenomenon of decreasing likelihoods observed in DPO training, counterintuitive under the assumption of likelihood maximization, is explained through the lens of maximum entropy RL. The paper elucidates that the implicit rewards modeled by DPO diminish over time when SFT precedes DPO—an expected behavior given the entropy-regularized objectives employed in DPO. Figure 3 captures the evolution of implicit rewards during training, affirming this behavior under different initialization conditions.

Figure 3: Evolution of implicit rewards for DPO and CPL during training, indicating reward dynamics under various starting conditions.

Practical Implications and Future Directions

The findings suggest several practical applications and future research avenues, including:

• Reasoning and Multi-turn Dialogues: Given DPO's capacity for per-token reward modeling, the approach could be extended to multi-turn dialogue systems, enhancing conversational alignment significantly better than current bandit evaluations.
• End-to-End Generative Systems: DPO provides a cohesive framework to train prompt generators and conditioning models jointly, optimizing whole multi-model systems based on direct feedback (Figure 4).
• Autonomous Agent Behavior: The ability to learn implicitly via token-level feedback opens possibilities to leverage DPO's strengths in agentic LLMs, promoting behaviors optimized for task-specific objectives obtained from preferences.

  Figure 4: End-to-end generative AI workflow illustration, highlighting interactions between user prompts, refined descriptions, and image generation models.

Conclusion

This paper bridges the conceptual gap between DPO as a bandit-based method and reinforcement algorithms traditionally employed in RLHF pipelines. By framing DPO as a solution within token-level MDPs, the research asserts that LLMs can, and do, embed optimal Q-functions through preference-driven learning. The theoretical advances extend DPO's applicability to nuanced AI systems, prompting its adoption in innovative directions such as integrated speech and vision models, thus fostering broader impacts across AI domains.