PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rafael Rafailov, Joey Hejna, Ryan Park, Chelsea Finn

arXiv:2404.12358v2cs.LG

TL;DR

The paper addresses the difference between DPO’s response-level framing and the token-level MDP used in language-model RL. It derives DPO as token-level optimal-Q-function learning, then uses that formulation to explain credit assignment, likelihood search, and reward trajectories.

  • Problem

    DPO is commonly framed with whole responses as actions, whereas language-model reinforcement learning uses token-level states and actions.

  • Method

    The paper derives DPO in the token-level MDP and shows that an LLM’s logits represent the optimal Q-function for an implicitly learned token-level reward.

  • Results

    DPO’s token-level interpretation supports per-token credit assignment, links likelihood search to reward-function search, and identifies initial-policy and reference-distribution choices as important for implicit-reward trajectories.

  • Takeaways & Limitations

    The formulation motivates beam-search decoding and future DPO applications in multi-turn conversations, agentic systems, and end-to-end generative-AI pipelines.

  • Takeaways & Limitations

    DPO adjusts the reward or advantage defining an optimal policy rather than training a policy to maximize reward, so its implied value functions need not increase over time.

Abstract

from arXiv · show

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.

1 Introduction

RLHF aligns language models using human-feedback rewards, while DPO offers a simpler direct-alignment alternative but differs fundamentally in how it models generated responses. This work derives DPO in the token-level MDP setting and identifies token-level reward and search implications.

  • RLHF learns reward functions from human-labeled comparisons to align large language models with complex human-intent objectives.
  • Direct alignment methods such as DPO optimize the reward-policy relationship simultaneously instead of separately learning a reward function and applying reinforcement learning.
  • Classical RLHF optimizes token-level value functions with terminal sparse rewards, whereas DPO treats an entire response as one contextual-bandit arm.
  • DPO is derived in the token-level MDP as learning a token-level reward function whose language-model logits define the optimal Q-function.
  • The token-level formulation motivates per-token credit assignment, likelihood search during decoding, and analysis of how reference policies affect implicit rewards during training.

2 Related Work

Prior work addresses preference alignment through classical RLHF, dense rewards, general-MDP direct alignment, value-guided search, and inverse-RL connections. This paper links these directions through a token-level interpretation of DPO.

  • Classical RLHF commonly learns terminal reward functions and optimizes them with policy-gradient methods, which can be unstable and difficult to scale.
  • Prior studies have applied dense rewards and approximate dynamic programming despite direct alignment methods traditionally treating language models as bandits.
  • Contrastive Preference Learning extends direct alignment to general MDPs, whereas this work interprets DPO as optimizing per-token rewards in the token MDP.
  • Earlier work uses best-of-K or tree search with value functions or discriminators, while this paper shows likelihood search produces a similar direct-alignment solution using implicit reward.
  • The paper builds on maximum-entropy reinforcement learning and inverse reinforcement learning mappings between Q-functions and reward functions.
  • Unlike related approaches requiring an additional reinforcement-learning optimization loop, the proposed feedback-learning formulation disposes of that loop for language models.

3 Preliminaries

The paper formalizes language generation as a token-level MDP and contrasts classical RLHF with DPO’s contextual-bandit derivation. Classical methods use token-level optimization around a learned terminal reward, while DPO uses a closed-form policy-reward relationship.

  • 3 Preliminaries: Language prompts and target responses are sequences of tokens, while contextual-bandit notation treats the entire response as one action.
  • 3.1 The Token-level MDP for Large Language Models: The token-level MDP uses generated-token histories as states, vocabulary tokens as actions, deterministic concatenation dynamics, and prompts as initial states.
  • 3.1 The Token-level MDP for Large Language Models: Preference trajectories share an initial prompt and terminate with an EOS token, after which future rewards are zero.
  • 3.2 The Classical RLHF Methods: Classical RLHF first learns a reward from preferred response pairs, then optimizes it with KL-constrained policy-gradient reinforcement learning.
  • 3.2 The Classical RLHF Methods: In token-level PPO, the learned contextual-bandit reward is applied at the final timestep, with reference-policy terms used across token transitions.
  • 3.2 The Classical RLHF Methods: This creates a mismatch because the reward is modeled as a bandit quantity while value estimation and optimization operate per token.
  • 3.3 Direct Preference Optimization: DPO remains in the contextual-bandit setting and derives its loss by substituting a closed-form KL-bandit reward-policy relationship into the preference-modeling loss.

4 Theoretical Insights

The paper derives DPO in the token-level MDP by establishing a bijection between reward functions and optimal Q-functions. This interpretation supports dense reward representation, credit assignment, and preference-preserving reward re-parameterization.

  • 4.1 DPO as a Q-function in the Token Level MDP: A bijection exists between reward functions and corresponding optimal Q-functions in the token MDP under mild assumptions.This relationship is established through the Bellman equations connecting current rewards, reference-policy terms, and future value.
  • 4.1 DPO as a Q-function in the Token Level MDP: DPO can be cast in the token-level MDP, where a language model’s logits represent an optimal Q-function for some reward.The optimal Q-function models expected total future reward, and the policy is obtained by applying a temperature-modulated softmax to the logits.
  • 4.1 DPO as a Q-function in the Token Level MDP: DPO’s token-level loss aligns the policy-induced preference model with empirical preference data by expressing sequence returns through Q-functions and policies.The derivation inverts the Bellman equation, substitutes the log-linearized optimal policy, and minimizes KL divergence between empirical and policy-implied preferences.
  • 4.2 Token-Level DPO Can Parameterize Any Dense Reward Function.: Advantage-based re-parameterization preserves both the optimal policy and the exact preference distribution associated with the original reward.The theorem applies to reward classes consistent with the Plackett-Luce and Bradley-Terry preference models, with a reference policy and β > 0.
  • 4.2 Token-Level DPO Can Parameterize Any Dense Reward Function.: Token-level implicit rewards can identify erroneous statements while leaving other token values largely unchanged, indicating credit assignment from answer-level feedback.In the job-interview-summary example, the model highlights errors in salary range and position level.
  • 4.2 Token-Level DPO Can Parameterize Any Dense Reward Function.: DPO can represent any per-token reward function when preference queries begin at the same state and terminate at a terminal state.The learned reward is restricted to an optimal advantage function, but this restriction does not reduce the representable reward class under the stated conditions.

5 Practical Insights

The empirical analysis examines DPO’s token-level credit assignment, its connection to search-based decoding, and how reference-policy choices shape implicit rewards during training.

  • 5.1 Does DPO Learn Credit Assignment?: DPO assigns per-token rewards that identify erroneous statements while preserving comparable values for correct parts of a response.The analysis computes each token’s DPO reward, equivalently interpreting it as an advantage or coherent reward.
  • 5.2 Connecting Guided Decoding and Search to Likelihood-Based DPO Optimization: Likelihood search on a DPO policy is theoretically equivalent to search over the corresponding learned reward function.Substituting the DPO reward representation into the search objective yields a sum of optimal-policy log likelihoods, up to a fixed initial-state term.
  • 5.2 Connecting Guided Decoding and Search to Likelihood-Based DPO Optimization: A 5-beam search improves win rates by 10-15% over the 1-beam base policy, while larger beam counts degrade performance as answer lengths explode.The comparison uses DPO models trained with three β values against preferred summaries on the test dataset.
  • 5.2 Connecting Guided Decoding and Search to Likelihood-Based DPO Optimization: DPO likelihood search connects to prior search and proxy-guidance methods, including MCTS-style token-MDP search and proposal-based reasoning-chain ranking.The paper relates likelihood search to value-function-guided search and describes V-STaR as using a proposal model with DPO likelihood ranking.
  • 5.4 Likelihoods should decrease when using DPO.: Starting DPO from SFT causes implicit rewards to decline on average, whereas training without SFT leaves DPO rewards positive and increases CPL rewards.The paper attributes the expected decline to positive KL divergence at convergence when the reference policy differs from the optimal policy.
  • 5.4 Likelihoods should decrease when using DPO.: Using a chosen-response SFT reference implies that DPO can decrease chosen-response likelihoods while favoring extrapolated responses, potentially causing over-fitting.The authors distinguish this likelihood behavior from any claim that DPO’s implied value functions must monotonically increase.

6 Discussion

The paper interprets DPO as learning an optimal Q-function represented by an LLM, connecting token-level credit assignment with likelihood-based search and several future applications.

  • DPO is formulated as learning an optimal Q-function represented by a language model.The authors report that this explains training phenomena not accounted for by the original bandit formulation.
  • DPO’s per-token implicit rewards provide qualitative early evidence of credit assignment from feedback data.The authors caution that larger-scale empirical exploration remains necessary.
  • Likelihood search under DPO unifies several proposed language-model search algorithms, with comparable empirical gains from a one-line beam-search change.
  • The framework suggests extending DPO to intermediate reasoning from outcome feedback and multi-turn conversational trees.The discussion identifies multi-turn dialogue as difficult because RLHF is optimized as a single-turn bandit formulation.
  • The derivations suggest applications to agentic LLM exploration and end-to-end optimization of multi-model generative AI systems.Examples include agents receiving feedback on final outputs and systems combining language models with diffusion generators.

A Proof of Lemma 1

Lemma 1 establishes a bijection between reward functions and optimal Q-functions in the deterministic, tree-structured token MDP, while emphasizing that the proof relies on token-MDP structure.

  • For a fixed policy, reward functions correspond bijectively to optimal Q-functions in the deterministic tree-structured LLM MDP.
  • Injectivity follows by tracing differences backward from terminal states to the first state-action pair where two rewards differ.Equal future rewards imply equal dynamic-programming values beyond that point, so differing immediate rewards produce differing Q-functions.
  • The proof recursively uses the Bellman-style relationships among rewards, optimal values, and Q-functions.
  • The argument does not hold for general MDPs because token MDPs cannot return to the same state after actions.
  • Surjectivity follows because, under deterministic dynamics, a reward can be recovered from any optimal Q-function as Q*(s_t, a_t) − V*(s_t+1).

B.2 Theoretical Results for the Diffusion MDP

The diffusion-MDP analysis extends the reward–Q-function bijection under mild assumptions and supports applying the paper’s token-level theoretical results to diffusion generation.

  • Under mild assumptions, reward functions and corresponding optimal Q-functions are bijective in the diffusion MDP.
  • The bijection is proved by backward induction through diffusion steps, starting from zero terminal value and uniquely determining earlier critic values.
  • The result permits direct application of the paper’s Section 4.2 results to the diffusion MDP.
  • Alternative intermediate-diffusion sampling schemes can be viewed as valid off-policy data-generation methods under the Q-learning interpretation of DPO.The authors suggest these choices could yield a family of DPO algorithms, leaving exploration for future work.

C Reddit TL;DR Posts

This section presents Reddit examples used to illustrate DPO’s token-level error identification, alongside an end-to-end generative-AI workflow example.

  • Reddit TL;DR examples: DPO identifies “busted knee” as the incorrect token in a preferred summary and correctly classifies the pair.
  • Reddit TL;DR examples: DPO identifies the incorrect token “iPhone” in a preferred summary and correctly classifies the pair.
  • Reddit TL;DR examples: DPO identifies “250k” and “management position” as incorrect tokens in another preferred summary and correctly classifies the pair.
  • End-to-end generative AI workflow: In the Minotaur example, joint training could optimize prompts for realizability while aligning the generated images.

D End-to-End Training of Generative AI Systems

The paper outlines end-to-end DPO training for a multi-step system combining an autoregressive prompt refiner with an image-generation diffusion model. Human preferences over generated images supervise optimization across the hybrid prompt-refinement and diffusion MDPs, with a real example and extensions discussed.

  • System design: The end-to-end system combines a language-model prompt refiner πθ(z|x) with a diffusion image generator πϕ(y|z).The refiner generates a discrete refined prompt, which conditions the downstream image model.
  • Feedback generation: Users compare two generated images, producing a preference pair that includes both the preferred image and its refined prompt.Two refined prompts are sampled, images are generated from them, and the user selects a preferred image.
  • Optimization: Optimization uses a hybrid MDP that switches from token-level prompt refinement to diffusion denoising when the EOS token is encountered.The authors state that this remains a valid MDP to which their derivations apply.
  • Empirical example: The motivating example shows a prompt refiner describing a Minotaur correctly while the image generator fails to follow the description for the rejected image.The chosen image reflects the prompt, and its refiner produces more descriptive text.
  • Extensions: The framework is presented as extendable to long-form video systems that combine autoregressive generation with diffusion-based frame generation or upscaling.The passage describes this as a straightforward extension of the joint DPO optimization framework.
  • PPO comparison: For PPO, 2 beams increased win rate and decreased sample length, whereas more than 2 beams reduced downstream performance and increased sample length.The authors describe this as over-optimization and report more limited beam-size benefits for the PPO model.

F.1 Samples Across Different Beam Searches

The beam-search samples show how increasing beam count changes summaries across Reddit topics. Larger beams sometimes produce more detailed summaries, but can also alter or omit details and introduce truncated outputs.

  • Illustration example: For the illustration question, the 1-beam summary captures the education-versus-self-teaching question, while 5 and 20 beams add age and slow-improvement details.The 20-beam version also retains the concern about being behind peers.
  • Workplace-syringe example: For the workplace-syringe post, 1 beam gives a short summary, while 5 and 20 beams retain more details about the injury and employer response.The 20-beam output additionally includes an answer fragment.
  • Wheelchair example: For the wheelchair example, 1 beam omits the sidewalk obstruction and limb-loss detail, whereas 5 and 20 beams preserve both.The 5- and 20-beam summaries are nearly identical.
  • Finance example: The finance-related sample includes separate SFT and non-SFT outputs, with the non-SFT result shown as unreadable script.The SFT output foregrounds inherited money, debt to the father, and the mother’s financial behavior.

G Dataset and Hyperparameter Details

The experiments use the Webis TLDR-derived dataset, Pythia 2.8B models, and specified SFT, DPO, generation, and evaluation settings. Around 5% of the data is held out for validation, and held-out evaluations use 256 samples.

  • Dataset: The TL;DR dataset contains 64,832 summary comparisons derived from Webis TLDR, with human feedback collected by OpenAI.Examples come from Reddit posts across several subreddits, and approximately 5% is reserved for validation.
  • Training: Each TL;DR model starts from a pretrained Pythia 2.8B model and uses one-epoch SFT and DPO training with a 0.5 × 10^-6 learning rate.The setup also specifies batch sizes, gradient accumulation, warmup, and four NVIDIA A40 GPUs for SFT.
  • Evaluation setup: Generation uses temperature 1.0 and maximum length 512 unless otherwise specified, while held-out evaluations use 256 samples.Win rates and lengths are computed from the held-out TL;DR set.

H GPT 4 Evaluation

GPT-4 evaluates pairs of summaries using a fixed comparison prompt. The evaluation randomizes the order of the model sample and reference response and requires both a preference explanation and an A/B choice.

  • Bias control: The order of the model sample and reference response is randomized in each evaluation to reduce positional bias.This randomization is part of the GPT-4 evaluation procedure.
  • Evaluation prompt: The GPT-4 evaluation asks which summary better captures the most important points of the forum post.The evaluator must compare the two summaries and explain the preference in one sentence.
  • Output format: The evaluator outputs a one-sentence comparison followed on a new line by only A or B as the preferred summary.The required format is `Comparison: ... Preferred: ...`.
Loading 2404.12358v2…