Technical Summary: Learning to Reason Efficiently with Discounted Reinforcement Learning

Problem Statement

Large Reasoning Models (LRMs) solve complex tasks (e.g., math, code) by emitting intermediate reasoning tokens before a final answer. While Reinforcement Learning (RL) post-training improves accuracy, it often leads to excessively long chains of thought, inflating computational costs, memory usage (via growing KV caches), and inference latency.

The prevailing assumption in the field is that longer responses inherently improve accuracy, suggesting a tradeoff between response length and correctness. This paper challenges that assumption. It posits that up to a regime determined by the model class and problem instance, there is no inherent tradeoff: one can reduce response length to a specific instance-dependent threshold without sacrificing accuracy. The goal is to train LRMs that reach correct answers effectively and efficiently, analogous to finding the shortest successful trajectory in a Stochastic Shortest Path (SSP) problem.

Methodology

Theoretical Framework: Blackwell Optimality

The authors model verifier-based reasoning as a finite-horizon Markov Decision Process (MDP) with a binary terminal reward (1 for correct, 0 for incorrect). To formalize the objective of maximizing accuracy while minimizing path length, they utilize Blackwell optimality.

• Definition: A policy is Blackwell optimal if it is optimal for all discount factors $\gamma$ sufficiently close to 1.
• Mechanism: In goal-reaching MDPs with binary terminal rewards, maximizing the discounted return $J_\gamma(\pi)$ as $\gamma \to 1$ first maximizes the probability of success (accuracy). Among policies that achieve the maximum success probability, the discounting mechanism breaks ties by preferring those with shorter expected trajectory lengths.
• Key Insight: For a finite policy class $\Pi$, there exists a critical discount factor $\gamma_{bw} < 1$ such that for all $\gamma \in (\gamma_{bw}, 1)$, the set of $\gamma$-optimal policies coincides with the set of Blackwell optimal policies. These policies simultaneously maximize accuracy and minimize expected path length among accuracy maximizers.

Practical Training Recipe

Guided by the theoretical results, the authors propose a specific training methodology using Group Relative Policy Optimization (GRPO):

1. Discounted Environment Reward Only: The discount factor $\gamma \in (0, 1)$ is applied exclusively to the extrinsic environment reward (correctness). Intrinsic rewards (formatting, structural shaping) remain undiscounted to ensure the model produces well-structured outputs.
2. Reasoning Token Masking: Discounting is applied only to tokens identified as part of the "reasoning span" (e.g., tokens between <reasoning> and </reasoning> tags). Tokens required for prompt adherence, formatting, and the final answer are not discounted.
3. KL Regularization: A KL penalty is applied against a moving reference policy (updated periodically) to prevent policy collapse (where the model stops reasoning too early) and stabilize learning.
4. Comparable Token Budgets: To ensure fair comparison, the number of rollouts is adjusted so that the total number of tokens processed during training is comparable between discounted and undiscounted methods.

The objective function optimized is:
$$\mathbb{E}_{S_1 \sim \mu, \tau \sim \pi(S_1)} [\gamma^{K(\tau)} r_e(\tau) + r_f(\tau)] - \beta \text{KL}(\pi || \pi')$$
Where $K(\tau)$ is the count of reasoning tokens, $r_e$ is the environment reward, and $r_f$ is the formatting reward.

Key Contributions

1. Theoretical Guarantee of No Tradeoff: The paper proves that within any fixed (possibly restricted) policy class, Blackwell optimal policies maximize undiscounted success probability and, among those maximizers, minimize expected trajectory length. This establishes that accuracy and path length are not inherently traded off up to an instance-dependent threshold.
2. Existence of Blackwell Factor: For finite policy classes, the authors prove the existence of a Blackwell discount factor $\gamma_{bw} < 1$. They provide bounds on how close $\gamma$ must be to 1 to maintain accuracy while shortening responses, clarifying hyperparameter selection for restricted deployment classes.
3. Empirical Validation: Using GRPO with the discounted objective, the authors demonstrate that models can substantially reduce mean response length on math benchmarks (GSM8K, MATH, AMC, AIME, etc.) while matching the Pass@1 accuracy of undiscounted baselines.

Experimental Results

The authors evaluated four instruction-tuned models (Qwen2.5 7B/14B, Llama 3 8B, Phi-4) on various math datasets.

• GSM8K & MATH: Discounted models matched or slightly exceeded the Pass@1 accuracy of undiscounted baselines while significantly reducing token counts.
  • Qwen2.5 7B (GSM8K): Accuracy remained stable (91.06% vs 91.07%), while mean response length dropped from 217.60 to 170.08 tokens (~22% reduction).
  • Llama 3 8B (GSM8K): Accuracy improved slightly (80.87% vs 81.07%), with length dropping from 125.43 to 108.67 tokens (~13% reduction).
• Harder Benchmarks (AMC, AIME, MINERVA, OLYMPIAD): The trend held for larger models (Qwen2.5 14B, Phi-4) trained on DeepScaleR. Discounted models achieved comparable or higher accuracy with significantly shorter traces (e.g., Phi-4 on AMC 2023: 51.00% $\to$ 61.00% accuracy, 1134 $\to$ 716 tokens).
• Discount Factor Sensitivity: Experiments varying $\gamma$ confirmed the theoretical tradeoff: as $\gamma$ moves further from 1, response lengths shorten reliably, but accuracy eventually drops if the discount is too aggressive.

Significance and Claims

The paper claims to provide a principled, instance-aware mechanism for efficient reasoning that does not rely on curated data or prompt engineering. By leveraging Blackwell optimality, the approach reframes the problem of "efficient reasoning" as a standard RL objective with a specific discounting strategy.

The authors assert that their results:

• Refute the inherent length-accuracy tradeoff: They show that shorter reasoning is possible without accuracy loss up to a specific threshold determined by the model's capacity.
• Unify Length Penalties: The work suggests that methods using small per-step negative rewards are equivalent to discounting in this setting, implying that several existing length-penalty methods may recover the same accuracy-then-length ordering when properly implemented.
• Offer a Practical Recipe: The proposed method requires no algorithmic changes to standard policy optimization (like GRPO), only a modification to the reward calculation (discounting reasoning tokens) and the use of a moving reference policy.

The paper concludes by noting that while shorter traces improve efficiency, the potential for compressed reasoning to improve generalization remains an open question for future work.