DUET: Optimize Token-Budget Allocation for Reinforcement Learning with Verifiable Rewards

• URL: http://arxiv.org/abs/2605.08441v1
• Date: Fri, 08 May 2026 20:03:19 GMT
• Title: DUET: Optimize Token-Budget Allocation for Reinforcement Learning with Verifiable Rewards
• Abstract summary: Reinforcement learning with verifiable rewards generates hundreds of thousands of tokens per training step.<n>We show that jointly tuning both decisions under a shared compute budget improves both reasoning quality and wall-clock training time.
• Score: 37.28110997883518
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Reinforcement learning with verifiable rewards (RLVR) generates hundreds of thousands of tokens per training step, with rollout generation dominating the computational cost. The overall token budget can be controlled along two main dimensions: (i) deciding which prompts to allocate rollouts to, and (ii) deciding how long each rollout should be. Prior work has generally controlled only one of these dimensions at a time. We show that jointly tuning both decisions under a shared compute budget improves both reasoning quality and wall-clock training time. We instantiate this view as \textbf{DU}al-controlled tok\textbf{E}n alloca\textbf{T}ion (DUET), a computationally efficient layer over GRPO that uses a lightweight pre-rollout surrogate of prompt informativeness to set how many rollouts each prompt receives, and a marker-gated abort rule with importance reweighting to set when to stop them. On Qwen3-1.7B trained on MATH, DUET outperforms full-budget GRPO and the other three budget-aware baseline methods. DUET's advantage further generalizes to other benchmarks across math and coding, and is on par with the best baseline on the scientific Q\&A domain, while also achieving a $1.62\times$ wall-clock speedup. More notably, using only 50\% of the token budget, DUET still outperforms all baseline methods at their full budget, achieving an even higher $2.51\times$ speedup over full-budget GRPO. We verify the high performance of DUET on other backbone LLMs, including Qwen3-4B and Llama-3.2-3B-Instruct. Notably, the gap between DUET and the strongest baseline \emph{widens} as the budget tightens, contrary to the usual pattern in which efficient methods trade off quality as compute decreases. More broadly, these results suggest that DUET budget-aware control strategies are valuable not only for accelerating training, but also for improving the quality of the learning signal.

Related papers

• One Policy, Any Budget: Internalizing Budget-Aware Search via Reinforcement Learning [41.55260785497756]
  We propose AnySearch, a framework that enables a single policy to perform budget-aware search under any budget constraint.<n>In the first phase, we train the agent with explicit budget state injection and structured reasoning prompts.<n>In the second phase, the scaffold is removed and the agent learns to operate autonomously under adaptively sampled budget constraints.
  arXiv  Detail & Related papers  (2026-09-01T07:12:32Z)
• Learning as Reasoning Unfolds: Progressive Rollout Allocation for Efficient Reinforcement Learning [23.987185167160117]
  VarIance Guided Online Rollout allocation (VIGOR)<n>VarIance Guided Online Rollout allocation (VIGOR)<n>VarIance Guided Online Rollout allocation (VIGOR)
  arXiv  Detail & Related papers  (2026-07-24T06:04:04Z)
• TurnOPD: Making On-Policy Distillation Turn-Aware for Efficient Long-Horizon Agent Training [52.6661804966136]
  On-policy distillation (OPD) trains a student policy by matching a stronger teacher on the student's own trajectories.<n>We propose TurnOPD, a turn-level budgeting strategy for efficient on-policy distillation of long-horizon agents.
  arXiv  Detail & Related papers  (2026-07-07T03:56:35Z)
• Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs [1.7499351967216341]
  Budget-Adaptive Curriculum Reasoning (BCAE) is a unified framework that jointly optimize reasoning quality and token efficiency.<n>BCAE includes a emphbudget-conditioned unified policy, a emphcurriculum-aware budget scheduler, and a emphtruncation-aware dense reward mechanism.
  arXiv  Detail & Related papers  (2026-03-29T18:31:09Z)
• Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory [56.0946692457838]
  BudgetMem is a runtime agent memory framework for explicit, query-aware performance-cost control.<n>A lightweight router performs budget-tier routing across modules to balance task performance and memory construction cost.<n>Across LoCoMo, LongMemEval, and HotpotQA, BudgetMem surpasses strong baselines when performance is prioritized.
  arXiv  Detail & Related papers  (2026-02-05T18:57:09Z)
• A Relative-Budget Theory for Reinforcement Learning with Verifiable Rewards in Large Language Model Reasoning [48.70183357021465]
  Reinforcement learning (RL) is a dominant paradigm for improving the reasoning abilities of large language models.<n>We propose a emphrelative-budget theory explaining this variation through a single quantity called relative budget $:= H/mathbbE[T]$.<n>We show that $$ determines sample efficiency by controlling reward variance and the likelihood of informative trajectories.
  arXiv  Detail & Related papers  (2026-02-02T01:31:52Z)
• ADAPT: Learning Task Mixtures for Budget-Constrained Instruction Tuning [1.0039548765955955]
  ADAPT is a meta-learning algorithm that emphlearns task sampling proportions under an explicit token budget for instruction tuning.<n>We conduct evaluations on 11 out-of-domain benchmarks spanning reasoning, reading comprehension, code generation, and instruction following.
  arXiv  Detail & Related papers  (2025-12-04T08:17:05Z)
• e1: Learning Adaptive Control of Reasoning Effort [88.51897900019485]
  Increasing the thinking budget of AI models can significantly improve accuracy, but not all questions warrant the same amount of reasoning.<n>Users may prefer to allocate different amounts of reasoning effort depending on how they value output quality versus latency and cost.<n>We propose Adaptive Effort Control, a self-adaptive reinforcement learning method that trains models to use a user-specified fraction of tokens.
  arXiv  Detail & Related papers  (2025-10-30T23:12:21Z)
• Train Long, Think Short: Curriculum Learning for Efficient Reasoning [51.506559652495476]
  We propose a curriculum learning strategy for length-controlled reasoning.<n>Our method starts with generous token budgets and gradually tightens them over training.<n>Experiments on GSM8K, MATH500, SVAMP, College Math, and GSM+ demonstrate that curriculum-based training consistently outperforms fixed-budget baselines.
  arXiv  Detail & Related papers  (2025-08-12T13:48:03Z)
• Steering LLM Thinking with Budget Guidance [48.65894557568655]
  Budget guidance is a method for steering the reasoning process of LLMs toward a target budget without requiring any fine-tuning.<n>Our approach introduces a lightweight predictor that models a Gamma distribution over the remaining thinking length.<n>This signal is then used to guide generation in a soft, token-level manner, ensuring that the overall reasoning trace adheres to the specified thinking budget.
  arXiv  Detail & Related papers  (2025-06-16T17:57:05Z)
• Sparsity Forcing: Reinforcing Token Sparsity of MLLMs [40.93786579652003]
  We explicitly reinforce token sparsity in well-posed multimodal large language models (MLLMs) through a simple RL-based post-training framework named textitSparsity Forcing.<n>Our method explores the efficiency-accuracy trade-off by running multiple rollouts with different token budgets, where both efficiency (token reduction ratio) and performance (answer correctness) are formulated as joint rewards.
  arXiv  Detail & Related papers  (2025-04-23T01:45:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.