Table of Contents
Fetching ...

Trajectory Bellman Residual Minimization: A Simple Value-Based Method for LLM Reasoning

Yurun Yuan, Fan Chen, Zeyu Jia, Alexander Rakhlin, Tengyang Xie

TL;DR

This paper introduces Trajectory Bellman Residual Minimization (TBRM), a simple off-policy value-based method for improving LLM reasoning by treating model logits as $Q$-values and minimizing a single trajectory-level Bellman residual. It removes the need for critics, importance sampling, or clipping and achieves convergence to a near-optimal KL-regularized policy under arbitrary off-policy data, backed by a refined change-of-trajectory-measure analysis. Empirically, TBRM matches or surpasses policy-based baselines (PPO and GRPO) on six mathematical benchmarks, delivering stronger performance with lower training time and memory, especially at smaller model scales and with fewer rollouts. The results suggest value-based RL can be a principled and efficient alternative for enhancing reasoning in LLMs, with broad implications for scalable, data-efficient RL in language tasks.

Abstract

Policy-based methods currently dominate reinforcement learning (RL) pipelines for large language model (LLM) reasoning, leaving value-based approaches largely unexplored. We revisit the classical paradigm of Bellman Residual Minimization and introduce Trajectory Bellman Residual Minimization (TBRM), an algorithm that naturally adapts this idea to LLMs, yielding a simple yet effective off-policy algorithm that optimizes a single trajectory-level Bellman objective using the model's own logits as $Q$-values. TBRM removes the need for critics, importance-sampling ratios, or clipping, and operates with only one rollout per prompt. We prove convergence to the near-optimal KL-regularized policy from arbitrary off-policy data via an improved change-of-trajectory-measure analysis. Experiments on standard mathematical-reasoning benchmarks show that TBRM consistently outperforms policy-based baselines, like PPO and GRPO, with comparable or lower computational and memory overhead. Our results indicate that value-based RL might be a principled and efficient alternative for enhancing reasoning capabilities in LLMs.

Trajectory Bellman Residual Minimization: A Simple Value-Based Method for LLM Reasoning

TL;DR

This paper introduces Trajectory Bellman Residual Minimization (TBRM), a simple off-policy value-based method for improving LLM reasoning by treating model logits as -values and minimizing a single trajectory-level Bellman residual. It removes the need for critics, importance sampling, or clipping and achieves convergence to a near-optimal KL-regularized policy under arbitrary off-policy data, backed by a refined change-of-trajectory-measure analysis. Empirically, TBRM matches or surpasses policy-based baselines (PPO and GRPO) on six mathematical benchmarks, delivering stronger performance with lower training time and memory, especially at smaller model scales and with fewer rollouts. The results suggest value-based RL can be a principled and efficient alternative for enhancing reasoning in LLMs, with broad implications for scalable, data-efficient RL in language tasks.

Abstract

Policy-based methods currently dominate reinforcement learning (RL) pipelines for large language model (LLM) reasoning, leaving value-based approaches largely unexplored. We revisit the classical paradigm of Bellman Residual Minimization and introduce Trajectory Bellman Residual Minimization (TBRM), an algorithm that naturally adapts this idea to LLMs, yielding a simple yet effective off-policy algorithm that optimizes a single trajectory-level Bellman objective using the model's own logits as -values. TBRM removes the need for critics, importance-sampling ratios, or clipping, and operates with only one rollout per prompt. We prove convergence to the near-optimal KL-regularized policy from arbitrary off-policy data via an improved change-of-trajectory-measure analysis. Experiments on standard mathematical-reasoning benchmarks show that TBRM consistently outperforms policy-based baselines, like PPO and GRPO, with comparable or lower computational and memory overhead. Our results indicate that value-based RL might be a principled and efficient alternative for enhancing reasoning capabilities in LLMs.

Paper Structure

This paper contains 57 sections, 7 theorems, 80 equations, 5 figures, 5 tables, 1 algorithm.

Key Result

Lemma 1

[lemma]lem: change-of-trajectory-measure Given an MDP $M = (H, \mathcal{S}, \mathcal{A}, \mathcal{P}, r, \rho)$ and a policy $\pi$, where the supremum is over all measurable functions $f: \mathcal{S}\times\mathcal{A}\to \mathbb R$.

Figures (5)

  • Figure 1: \ref{['fig:average_accuracy']} shows average learning accuracy across benchmarks for PPO, GRPO, and TBRM. Per-benchmark results can be found in \ref{['appendix:perbench_accuracy']}. \ref{['fig:response_length']} demonstrates model's ability to engage in extended reasoning at test time with R1 template.
  • Figure 2: \ref{['fig:reward']} demonstrates the increment of rewards during training. \ref{['fig:train_eff']} is a comparison of maximal GPU memory consumption and per-step time cost across different methods with base model Qwen2.5-Math-7B. Time cost is segmented into key partitions, with time partition labels defined in \ref{['tab:train-eff-label-desc']} (\ref{['appendix:train-eff']}).
  • Figure 3: Learning accuracy across benchmarks with model Qwen2.5-Math-1.5B.
  • Figure 4: Learning accuracy across benchmarks with model Qwen2.5-Math-7B.
  • Figure 5: Training reward with classical BRM on model Qwen2.5-Math-1.5B.

Theorems & Definitions (11)

  • Lemma 1: Change-of-Trajectory-Measure Lemma
  • Theorem 2
  • Lemma 3: Soft Performance Difference Lemma via Reference Function
  • proof : Proof of \ref{['lem:soft_pd']}
  • Corollary 4
  • proof : Proof of \ref{['cor:all-pol-sub-opt']}
  • Lemma 5
  • proof : Proof of \ref{['lem: change-of-trajectory-measure']}
  • Corollary 6
  • proof : Proof of \ref{['cor:cotm-cor']}
  • ...and 1 more