Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget
  Allocation
Large Language Models (LLMs) can self-improve through reinforcement learning,
where they generate trajectories to explore and discover better solutions.
However, this exploration process is computationally expensive, often forcing
current methods to assign limited exploration budgets to each task. This
uniform allocation creates problematic edge cases: easy tasks consistently
succeed while difficult tasks consistently fail, both producing zero gradients
during training updates for the widely used Group Relative Policy Optimization
(GRPO). We address this problem from the lens of exploration budget allocation.
Viewing each task's exploration as an "item" with a distinct "value" and
"cost", we establish a connection to the classical knapsack problem. This
formulation allows us to derive an optimal assignment rule that adaptively
distributes resources based on the model's current learning status. When
applied to GRPO, our method increases the effective ratio of non-zero policy
gradients by 20-40% during training. Acting as a computational "free lunch",
our approach could reallocate exploration budgets from tasks where learning is
saturated to those where it is most impactful. This enables significantly
larger budgets (e.g., 93 rollouts) for especially challenging problems, which
would be computationally prohibitive under a uniform allocation. These
improvements translate to meaningful gains on mathematical reasoning
benchmarks, with average improvements of 2-4 points and peak gains of 9 points
on specific tasks. Notably, achieving comparable performance with traditional
homogeneous allocation would require about 2x the computational resources.
68de03246024653e8a3ed14f
Ziniu Li
68de03246024653e8a3ed150
Congliang Chen
68de03246024653e8a3ed151
Tianyun Yang
68de03246024653e8a3ed152
Tian Ding
68de03246024653e8a3ed153
Ruoyu Sun
68de03246024653e8a3ed154
Ge Zhang
68de03246024653e8a3ed155
Wenhao Huang
68de03246024653e8a3ed156
Zhi-Quan Luo
Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget
  Allocation
Large Language Models (LLMs) can self-improve through reinforcement learning,
where they generate trajectories to explore and discover better solutions.
However, this exploration process is computationally expensive, often forcing
current methods to assign limited exploration budgets to each task. This
uniform allocation creates problematic edge cases: easy tasks consistently
succeed while difficult tasks consistently fail, both producing zero gradients
during training updates for the widely used Group Relative Policy Optimization
(GRPO). We address this problem from the lens of exploration budget allocation.
Viewing each task's exploration as an "item" with a distinct "value" and
"cost", we establish a connection to the classical knapsack problem. This
formulation allows us to derive an optimal assignment rule that adaptively
distributes resources based on the model's current learning status. When
applied to GRPO, our method increases the effective ratio of non-zero policy
gradients by 20-40% during training. Acting as a computational "free lunch",
our approach could reallocate exploration budgets from tasks where learning is
saturated to those where it is most impactful. This enables significantly
larger budgets (e.g., 93 rollouts) for especially challenging problems, which
would be computationally prohibitive under a uniform allocation. These
improvements translate to meaningful gains on mathematical reasoning
benchmarks, with average improvements of 2-4 points and peak gains of 9 points
on specific tasks. Notably, achieving comparable performance with traditional
homogeneous allocation would require about 2x the computational resources.
An adaptive exploration budget allocation method for reinforcement learning in Large Language Models improves training efficiency and performance on mathematical reasoning benchmarks.
reinforcement learning trajectories Group Relative Policy Optimization (GRPO) exploration budget allocation knapsack problem policy gradients computational resources mathematical reasoning benchmarks