PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizing Anytime Reasoning via Budget Relative Policy Optimization

Penghui Qi, Zichen Liu, Tianyu Pang, Chao Du, Wee Sun Lee, Min Lin

arXiv:2505.13438v3cs.LGcs.AIcs.CL

TL;DR

Existing RL for LLM reasoning often rewards only the answer after a complete, long reasoning trace, despite the costs of such traces. AnytimeReasoner trains on sampled thinking budgets with verifiable summaries and decoupled policy optimization, and its reported results consistently exceed GRPO across thinking budgets and prior distributions.

  • Problem

    Most existing RL approaches optimize final performance from complete reasoning traces, although long chains of thought are costly in training and deployment.

    Source: §1 · p. 2

  • Method

    AnytimeReasoner trains with sampled thinking budgets, verifies summaries of partial reasoning for dense rewards, decouples thinking and summary policy optimization, and applies BRPO for variance reduction.

    Source: §1 · p. 2, §2.1 · p. 3, §2.2 · p. 4

  • Results

    AnytimeReasoner variants consistently outperform GRPO across varying prior distributions and thinking budgets, with different priors converging to similar performance at large budgets.

    Source: §3.1 · p. 7

  • Takeaways & Limitations

    The framework targets reasoning that can provide a useful answer from incomplete thinking and adapt to different available token budgets.

    Source: §1 · p. 2

    BRPO's variance-reduction rationale assumes a perfect summary policy yields rewards that do not decrease as more thinking tokens are generated.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

Scaling test-time compute is crucial for enhancing the reasoning capabilities of large language models (LLMs). Existing approaches typically employ reinforcement learning (RL) to maximize a verifiable reward obtained at the end of reasoning traces. However, such methods optimize only the final performance under a large and fixed token budget, which hinders efficiency in both training and deployment. In this work, we present a novel framework, AnytimeReasoner, to optimize anytime reasoning performance, which aims to improve token efficiency and the flexibility of reasoning under varying token budget constraints. To achieve this, we truncate the complete thinking process to fit within sampled token budgets from a prior distribution, compelling the model to summarize the optimal answer for each truncated thinking for verification. This introduces verifiable dense rewards into the reasoning process, facilitating more effective credit assignment in RL optimization. We then optimize the thinking and summary policies in a decoupled manner to maximize the cumulative reward. Additionally, we introduce a novel variance reduction technique, Budget Relative Policy Optimization (BRPO), to enhance the robustness and efficiency of the learning process when reinforcing the thinking policy. Empirical results in mathematical reasoning tasks demonstrate that our method consistently outperforms GRPO across all thinking budgets under various prior distributions, enhancing both training and token efficiency.

1 Introduction

AnytimeReasoner targets reasoning that remains useful when thinking is interrupted or constrained by a token budget, addressing the inefficiency of optimizing only complete, long reasoning traces. It combines sampled budgets, verifiable rewards at intermediate budgets, decoupled policy optimization, and BRPO; experiments report gains over GRPO.

1 Introduction

  • Existing RL approaches typically optimize final performance from complete reasoning traces, making long chains of thought costly in training and deployment.Anytime reasoning instead aims to provide the best available solution when computation is interrupted, with quality expected to improve as more resources are allocated.Source: §1 · p. 2
  • Sampling thinking budgets during learning aims to make performance robust to interruptions while encouraging correct answers with fewer tokens.The approach seeks to balance token efficiency with thorough exploration, including at larger budgets.Source: §1 · p. 2
  • At predefined sampled budgets, the model summarizes its answer for verification, adding dense rewards that support richer training signals and credit assignment.BRPO uses rewards from previous budgets in its advantage function, while thinking and summary policies are optimized separately; uniform budget sampling is used for the summary policy.Source: §1 · p. 2
  • AnytimeReasoner consistently surpasses GRPO in final and anytime performance, and still does so when trained only at the maximum token budget.The introduction also reports gains from independently adding decoupled optimization, variance reduction, and budget sampling to GRPO.Source: §1 · p. 2

2 Methodology

AnytimeReasoner trains LLMs to produce verifiable answers from thinking truncated at sampled budgets, and separately optimizes thinking and summary policies. BRPO adds budget-relative variance reduction to the thinking-policy update.

2 Methodology

  • The method generates a comprehensive chain of thought, then a summary scored by a rule-based reward, and decouples their generation with separate token budgets.When thinking is interrupted, an ellipsis and </think> prompt the model to summarize the incomplete reasoning.Source: §2 · pp. 2–3

2.1 Optimizing Anytime Reasoning

  • Anytime training maximizes expected score across budgets sampled from a prior, evaluating only a small predefined support of at most eight budgets in experiments.This objective introduces verifiable dense rewards at each budget, providing more informative credit assignment than a single reward for the complete trace.Source: §2.1 · p. 3
  • The anytime objective is a lower bound on the standard reasoning objective and reduces to it when training only at full reasoning length.The paper states that maximizing anytime performance should therefore also enhance standard-task performance.Source: §2.1 · p. 4

2.2 Budget Relative Policy Optimization

  • BRPO uses correlated prefix rewards as a variance-reduction baseline, discounting scores from earlier budgets and complementing them with group-based estimates when prefixes are short.The prefix baseline relies on the assumption of an optimal summary policy that extracts the best answer from the thinking so far.Source: §2.2 · pp. 4–5
  • BRPO significantly reduces variance versus GRPO, especially for long thinking.Figure 3 reports the comparison for R1-Distill-1.5B with λ = 0.5 and a uniform prior over budgets from 1,000 to 8,000 tokens.Source: Figure 3 · p. 5, §2.2 · p. 5

2.3 Decoupled Optimization for Thinking and Summary

  • Thinking and summary policies use different budget distributions because summary quality on incomplete reasoning is sensitive to the prior; summaries are optimized with grouped samples and GRPO.The summary-policy prior is set to a uniform distribution over the budget support in the described approach.Source: §2.3 · p. 5

3 Experiments

Experiments compare AnytimeReasoner with GRPO across thinking budgets and prior distributions, then ablate dense rewards, decoupled optimization, and variance reduction on mathematical reasoning tasks.

3 Experiments

  • The 1.5B model is fine-tuned for one epoch on 40,315 DeepScaleR math problems and evaluated on five benchmarks across 32 budgets from 0 to 8,000 tokens.Training uses four thinking budgets: 2,000, 4,000, 6,000, and 8,000 tokens.Source: §3 · p. 6

3.1 Main Results

  • All method variants consistently outperform GRPO across varying prior distributions and thinking budgets.With small budgets, the uniform variant excels; at large budgets, variants converge to similar performance, and the base variant also beats GRPO at all budgets.Source: §3.1 · p. 7

3.2 Ablations

  • The ablations assess anytime accuracy, final accuracy at 8,000 tokens, and average thinking length at the maximum budget.The study isolates verifiable dense rewards, decoupled optimization, and variance reduction as three factors.Source: §3.2 · p. 7

3.2.1 Verifiable Dense Rewards

  • Verifiable dense rewards improve anytime and final performance and shorten average thinking length; length-penalty reward shaping sacrifices performance and is unstable.The shorter responses are attributed to encouraging the thinking policy to reach a correct answer quickly.Source: §3.2.1 · pp. 7–8

3.2.2 Decoupled Optimization

  • Decoupled GRPO outperforms vanilla GRPO, especially on AMC, and improves anytime accuracy through a better summary policy.The summary policy is trained under uniform budgets of 2,000, 4,000, 6,000, and 8,000 tokens, while the thinking policy is optimized only at 8,000.Source: §3.2.2 · p. 8

3.2.3 Variance Reduction

  • BRPO variance reduction enhances AIME performance, with further gains after decoupled optimization improves the summary policy.The ablation aligned the BRPO objective with GRPO while holding summary-policy training consistent with GRPO.Source: §3.2.3 · pp. 8–9

3.3 Evaluation on 7B Model

  • On DeepSeek-R1-Distill-Qwen-7B, anytime reasoning improves by up to about 5 absolute points, while standard reasoning beats GRPO for most of training despite high final-accuracy variance.The 7B experiment used two epochs on DeepScaleR with batch size 128.Source: §3.3 · p. 9

4 Related Works

Related work addresses verifiable-reward RL, token-budget efficiency, and test-time compute optimization; AnytimeReasoner differs in its sampled-budget formulation and per-token RL approach.

Reinforcement Learning with Verifiable Rewards

  • RLVR work improves reasoning training but often retains sparse outcome-based rewards, which challenge credit assignment and learning efficiency.The paper contrasts those approaches with its use of cheaply estimated verifiable dense rewards.Source: §4 · p. 9

Token Budget Efficiency of Reasoning Models

  • Prior token-efficiency work uses prompting, adaptive sampling, or reward shaping, while AnytimeReasoner instead targets interruptible reasoning without explicitly enforcing conciseness.The paper states that its objective naturally encourages efficient reasoning, as demonstrated empirically.Source: §4 · pp. 9–10

Connection to MRT

  • Unlike MRT’s oracle-regret objective and greedy episode-level progress, this work samples thinking budgets from a prior and optimizes per-token long-term returns with standard RL techniques.The paper reports 32.7% accuracy on AIME 2024 versus MRT’s reported 30.3%.Source: §4 · p. 10

5 Conclusion

The conclusion presents AnytimeReasoner as a way to exploit the ease of verifying answers extracted from partial reasoning traces.

5 Conclusion

  • Preset thinking budgets yield verifiable dense rewards, which support RL credit assignment and a variance-reduction baseline that improves training stability and efficiency.The framework is reported to achieve superior performance on both standard and anytime reasoning tasks.Source: §5 · p. 10

A Implementation Details

The implementation describes truncating reasoning and prompting the model to produce an answer for verification.

A Implementation Details

  • The implementation details focus on how the reasoning process is truncated and how the model is prompted to output an answer.Source: A Implementation Details · p. 14
  • Thinking ends at the special token or budget limit; a Final Answer prompt then elicits a summary, with ellipses and an end token marking out-of-budget cases.Inserted prompt tokens are ignored when calculating training loss.Source: Figure 9 · p. 14

B Tree-like Generation and Training

AnytimeReasoner uses a tree-like structure to generate answers from thinking traces truncated at predefined budgets. Its training attention mask is designed to reduce duplicated computation for these tree-shaped responses.

B Tree-like Generation and Training

  • The method replaces sequential question-response generation and training with a tree-like structure.Figure 10 depicts the tree-like generation and training structure.Source: B Tree-like Generation and Training · p. 14, Figure 10 · p. 15
  • Generation samples a complete thinking process, splits it at predefined token budgets, and prompts for an answer after appending an end-of-think token to each partial trace.Prefix caching in vLLM reuses computations during generation.Source: B Tree-like Generation and Training · p. 14
  • Training appends summaries after the thinking process and uses a FlexAttention block mask from recorded connection positions to avoid duplicated computation and 2D-mask memory issues.The 1D connection-position tensor is converted into a block mask.Source: B Tree-like Generation and Training · p. 15

C Relation Between Standard and Anytime Reasoning

This section proves an inequality and states a nonnegative-reward assumption. It also distinguishes final accuracy at the maximum budget from anytime accuracy evaluated across multiple budgets.

C Relation Between Standard and Anytime Reasoning

D.1 Main Results

The main-results section reports training curves and evaluates model accuracy at training step 600, distinguishing maximum-budget final accuracy from anytime accuracy.

D.1 Main Results

  • For experiments in Section 3.1, the paper presents AnytimeReasoner training curves in Figure 11 and evaluates model performance at training step 600, reporting final accuracy in Table 1 and anytime accuracy in Table 2.Source: §D.1 · p. 16
  • Figure 11 presents the training curves for the main results.Source: Figure 11 · p. 16, §D.1 · p. 16

D.2 Evaluation on 16k Context Length

The 16k-context evaluation fine-tunes DeepSeek-R1-Distill-Qwen-1.5B using a uniform prior over eight budgets and reports stronger performance in both anytime and standard reasoning.

D.2 Evaluation on 16k Context Length

  • The approach consistently achieves stronger performance in both anytime and standard reasoning with a maximum context length of 16k.Figure 12 presents these results for DeepSeek-R1-Distill-Qwen-1.5B.Source: §D.2 · p. 16, Figure 12 · p. 17
  • The evaluation uses a uniform prior over eight thinking budgets from 2k through 16k tokens.Source: §D.2 · p. 16
Loading 2505.13438v3…