PaperlayerView exact revision on arXiv

Source-linked AI summary

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

Ziniu Li, Congliang Chen, Tianyun Yang, Tian Ding, Ruoyu Sun, Ge Zhang, Wenhao Huang, Zhi-Quan Luo

arXiv:2509.25849v1cs.LGcs.AIcs.CL

TL;DR

Uniform exploration budgets can leave easy tasks with only successes and hard tasks with only failures, producing zero gradients in GRPO. The paper allocates exploration through a knapsack formulation and reports higher effective gradient ratios and benchmark performance without additional computation.

  • Problem

    Uniform exploration budgets can leave easy tasks all-success and hard tasks all-failure, producing zero gradients in GRPO while hard tasks may need extensive exploration.

    Source: §1 · p. 2

  • Method

    Knapsack RL treats task-budget pairs as items with learning value and computational cost, allocating a fixed exploration budget using a value estimate based on expected success-probability increase.

    Source: §1 · p. 2, §4.1 · p. 7

  • Results

    20–40% improvement in effective gradient ratios and about 2–4 points in average performance gains are reported across challenging benchmarks and Qwen models sized 1B–7B.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The approach prioritizes harder tasks for exploration and is reported to improve gradients and policy performance at no additional computational cost.

    Source: §7 · p. 12

    About 20% of prompts remain extremely hard after 1,000 training iterations, although 577 recorded at least one positive trajectory during optimization.

    Source: §5.2 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames exploration as a costly but necessary part of RL self-improvement and argues that uniform rollout budgets mismatch task needs. Knapsack RL allocates exploration heterogeneously, with reported gains in effective gradients and benchmark performance.

1 Introduction

  • Exploration requires costly sequential generation, yet uniform budgets can make easy tasks all-success and hard tasks all-failure, leaving GRPO with zero gradients.RL with verifiable rewards uses generated responses and binary outcome feedback to refine model policies.Source: §1 · pp. 1–2
  • The authors frame the core mismatch as hard tasks receiving too little exploration while easy tasks are over-sampled and waste computation.Hard tasks may require more than 100 explorations to find useful trajectories.Source: §1 · p. 2
  • Knapsack RL treats each task at an assigned budget as an item with learning value and computational cost, then maximizes value under a fixed global budget.The framework enables dynamic, heterogeneous exploration budgets within GRPO.Source: §1 · p. 2, Figure 1 · p. 2
  • 20–40% improvement in effective gradient ratios accompanies average gains of about 2–4 points on challenging benchmarks across Qwen models sized 1B–7B.The authors report that comparable improvements with uniform allocation would require nearly 2× the computation.Source: §1 · p. 2

2 Preliminary

The preliminary section describes autoregressive response generation and binary verifiable rewards as the basis for policy-gradient RL. It also outlines homogeneous rollout allocation and GRPO's within-prompt relative feedback.

2 Preliminary

  • Language generation is modeled as autoregressive sampling from a conditional distribution, with RL aiming to maximize expected response performance.The prompt is x, the generated response is y, and θ denotes trainable parameters.Source: §2 · p. 2
  • Verifiable rewards score an answer as correct or incorrect, producing binary feedback for generated responses.The paper describes the reward as 1 for correct and 0 for incorrect responses.Source: §2 · p. 3
  • Classical multi-task policy-gradient training uses homogeneous exploration, sampling the same number of responses for each prompt before evaluating rewards and updating parameters.The described loop samples N responses per prompt, evaluates rewards, computes gradients, and updates the model.Source: §2 · p. 3
  • GRPO computes relative advantages within each prompt's response group, increasing the likelihood of positive responses and decreasing that of negative ones.A small constant prevents division by zero when the reward standard deviation is zero.Source: §2 · p. 3

3 Diagnosing Exploration in Homogeneous Budget Allocation

This section diagnoses how homogeneous exploration can waste samples and leave GRPO without useful gradients, then analyzes the rollout budgets needed to obtain non-zero gradients.

3.1 Motivation

  • When all sampled responses for a prompt have identical rewards, its gradient is zero and it supplies no learning signal.The section tracks the proportion of individual samples with non-zero gradients as the effective-gradient-ratio.Source: §3.1 · pp. 3–4
  • During 0–70 iterations, all-negative rewards peak near 95%; from 70–600, mixed outcomes let the effective-gradient ratio remain above 40%.Mixed successes and failures within a response group provide the outcomes needed for effective gradients.Source: §3.1 · pp. 3–4
  • After 600 iterations, all-positive rewards rise to 40% while all-negative rewards fluctuate around 20%, and the effective-gradient ratio falls to about 20% by iteration 1,000.Overall, the effective-gradient ratio remains below 60%, so more than 40% of sampled data fails to contribute to updates.Source: §3.1 · p. 4

3.2 Theoretical Analysis

  • At p=0.01, non-zero gradients need 100 rollouts on average or 229 for a 90% chance; N=8 covers only p∈[0.1,0.9].For tasks near p=0 or p=1, even N=16 or 32 may be insufficient.Source: Figure 3 · p. 5, §3.2 · p. 5
  • Uniformly increasing the budget can help extreme tasks but very large budgets are often impractical; filtering hard prompts can remove opportunities to learn from failures.These trade-offs motivate a more principled allocation of exploration budgets.Source: §3.2 · pp. 5–6

4 Proposed Approach: Knapsack-based Budget Allocation

The section formulates exploration allocation as a knapsack problem: with compute fixed, task-specific rollout budgets are chosen to maximize learning value. It defines that value using success rates and gradient information, then describes estimation, optimization, and rollout-balancing procedures.

4 Proposed Approach: Knapsack-based Budget Allocation

  • The method maximizes total learning value under a fixed rollout budget by treating each task–budget pair as a knapsack item, with rollout count as weight.The task–budget pair, rather than task difficulty alone, determines an item's value.Source: §4 · p. 6, Table 1 · p. 6

4.1 Formulation of Task Value

  • For GRPO, Value combines the probability of a non-zero gradient with InfoGain, the expected success-probability increase after an update, approximated from current success rate.The post-update success probability is intractable to compute directly; the approximation uses current success rate.Source: §4.1 · p. 7
  • InfoGain peaks at pi = 1/3, favors harder over equally distant easier tasks, and diminishes near pi = 0 or 1; it does not model positive-versus-negative sample balance.The design accounts for whether a non-zero gradient is likely, not the exact balance of positive and negative samples.Source: §4.1 · p. 7
  • For comparable high value, pi = 0.35 needs Ni = 4, pi = 0.19 needs Ni = 16, and pi = 0.52 needs Ni = 8.The contour plot depicts equal-value lines across success rate and exploration budget.Source: §4.1 · p. 7

4.2 Algorithm Implementation

  • Success rates use prior-epoch observations, with homogeneous allocation in epoch one; despite delay and noise, dynamic programming solves the discrete allocation in polynomial time, typically in 1–2 seconds with Numba.The estimated success rates guide the constrained optimization.Source: §4.2 · pp. 7–8
  • For pi = 1, a minimum budget such as 2 maintains coverage; for pi = 0, remaining budget is assigned to extremely hard prompts.The fallback is especially useful later in training, when many prompts have become easy.Source: §4.2 · p. 8
  • Randomly dispatching individual rollout jobs across workers mitigates workload imbalance; for long prompts, the authors propose balanced batches using Karmarkar–Karp.The random-dispatch strategy is presented for settings where prompts are not excessively long.Source: §4.2 · p. 8
  • The method adds negligible computational overhead, requires no additional hyperparameters, and leaves core inference and training components unchanged.The authors also state that the approach does not bias policy gradients.Source: §4.2 · p. 8

5 Experiments

Across tested models, Knapsack-GRPO outperforms GRPO on mathematical reasoning benchmarks. The experiments show how adaptive budgets affect gradient effectiveness and prompt learning, including under different compute budgets.

5.1 Main Results

  • Knapsack-GRPO consistently outperforms GRPO across tested models, including a 3.8-point average gain for DPSK-R1-Distill-1.5B and a 9.1-point AMC gain for Qwen3-4B-Base.Other reported gains include 6.4 points on AIME for DPSK-R1-Distill-1.5B, 5.5 points on GPQA for Qwen3-4B-Instruct, and 6.8 points on AMC for Qwen2.5-Math-7B.Source: Table 2 · p. 9, §5.1 · p. 9

5.2 Understanding Knapsack-based Exploration

  • Knapsack-GRPO dynamically assigns up to 93 exploration rollouts to individual prompts without additional compute, a level impractical under homogeneous allocation.Figure 5 shows the frequency of different exploration budgets allocated to prompts during Qwen2.5-Math-7B training.Source: Figure 5 · p. 9, §5.2 · p. 9
  • Knapsack-based allocation improves the effective gradient ratio by approximately 20–40% across models and avoids uniform allocation’s clear decreasing trend.The authors attribute this to directing exploration toward tasks with mixed successful and failed trajectories.Source: §5.2 · pp. 9–10, Figure 6 · p. 10
  • Extremely-hard prompts remain in that status less often with Knapsack-GRPO than GRPO: 43.4% versus 47.1%, with more transitions to extremely-easy status.The transition matrix covers five prompt categories based on success rate.Source: Figure 7 · p. 10, §5.2 · p. 10
  • About 20% of prompts remain extremely hard after 1,000 iterations, although 577 such Knapsack-GRPO prompts recorded at least one positive trajectory.The authors suggest experience replay as a direction for addressing these samples.Source: §5.2 · p. 11
  • After training, Knapsack-GRPO has 3,596 extremely-hard tasks versus GRPO’s 3,793, and 9,274 extremely-easy tasks versus 8,676.These are counts in the final prompt-status distribution.Source: §5.2 · p. 10, Figure 8 · p. 11

5.3 Experiments with Different Exploration Budgets

  • In the low-budget scenario, Knapsack-GRPO improves Qwen2.5-Math-7B performance from 39.8 to 45.5 and maintains its advantage over GRPO with higher exploration budgets.Source: §5.3 · p. 11

6 Related Work

Related work addresses prompt heterogeneity, prompt selection, and resource allocation from different angles. This paper focuses on allocating exploration budgets across prompts during online RL.

Data heterogeneity

  • Prior work identifies varying reward distributions across prompts as a challenge for stable gradient estimation, while refined baselines target variance during exploitation.The paper distinguishes its focus on exploration: allocating resources to collect informative trajectories.Source: §6 · p. 11

Prompt selection and curriculum learning

  • Prompt-selection and curriculum methods choose which prompts to train on, whereas Knapsack RL varies how much exploration each prompt receives.The paper contrasts selecting a subset of simpler prompts with designing budgets for all prompts, especially harder ones.Source: §6 · pp. 11–12

Resource allocation

  • Unlike prior knapsack-constrained online exploration work typically framed for a single task [3] [7], Knapsack RL allocates a centralized budget across multiple tasks.The formulation explicitly distributes exploration budgets across tasks.Source: §6 · p. 12
  • Related resource allocation for rejection sampling and RAFT focuses on variance reduction [47], while this method targets online RL training signals.Other cited work allocates computation at inference time [41].Source: §6 · p. 12

7 Conclusion

The conclusion frames exploration-budget allocation as a knapsack optimization that prioritizes harder tasks for more effective gradients and policy improvements, without additional computational cost. It identifies three directions for extending the framework.

7 Conclusion

  • Knapsack allocation prioritizes harder tasks to yield more effective gradients and superior policy improvements at no additional computational cost.Each task-budget pair is treated as an item with a cost and value.Source: §7 · p. 12
  • The current budget measures exploration by rollout count; future versions could also account for response token length or interaction turns.Source: §7 · p. 12
  • The task-value model uses a first-order Taylor expansion of one-step policy improvement; alternative value functions could assess learning potential more accurately.Source: §7 · p. 12
  • The study uses simple on-policy rollouts, which leave some tasks unsolved; tree-based exploration and state rollbacks are proposed as directions to explore.The conclusion points to Tree-RL [11] and TreePO [19] as examples.Source: §7 · p. 12

A Proof

The proof analyzes rollout requirements for obtaining non-zero gradients and derives an approximation for task value under specified policy-gradient assumptions. It also reports an empirical check of that approximation.

A Proof

  • The high-probability analysis seeks the minimum rollout count N such that P(g_i ≠ 0) ≥ α, with q^N dominating for sufficiently large N.Here q is defined as max{p_i, 1 − p_i}, and the proof notes q > 1/2.Source: A Proof · p. 16
  • The proof defines N as the first rollout count at which both a success and a failure have appeared, then computes its expectation by conditioning on the first trial.The two cases condition on whether the first trial is a success or failure.Source: A Proof · pp. 16–17
  • The InfoGain derivation assumes a softmax policy and policy-gradient update, with unit learning rate and advantage for simplicity.It uses a first-order Taylor approximation and neglects second-order cross-terms when updates are small.Source: A Proof · p. 17
  • With 100 actions, the empirical comparison found that p(1 − p)^2 closely approximates InfoGain across different success rates.Source: A Proof · p. 18

B Extensions

The extensions discuss adapting the budget-allocation framework to other RL algorithms by calculating each algorithm’s probability of a non-zero gradient. They describe different gradient conditions across algorithms while retaining the same InfoGain formulation.

B Extensions

  • Figure 10 compares exact InfoGain with the approximate formula.Source: Figure 10 · p. 18
  • RLOO’s policy-gradient estimator is equivalent to GRPO up to constants, so fundamental changes may not be needed.Source: B Extensions · p. 18
  • For ReMax, a gradient update occurs when a sampled trajectory differs from the greedy response; with greedy-response probability α, that event has probability 1 − α.Source: B Extensions · p. 18
  • Vanilla REINFORCE’s non-zero-gradient probability accounts for whether at least one trajectory receives positive reward, rather than balancing positive and negative samples.Source: B Extensions · p. 18
  • The framework can adapt to other RL algorithms by recalculating ProbNonZeroGradient while maintaining the InfoGain(pi) = pi(1 − pi)^2 task-value formulation.Source: B Extensions · p. 18

C Experiment Details

The experiment details describe implementation and evaluation settings, alongside visualizations of prompt learning and training curves. The accompanying observations include prompt-level variation, early performance improvement, and some performance degeneration.

C Experiment Details

  • Training evaluations use 16 generated responses and, when benchmarks exceed 100 samples, randomly select 100; final Table 2 evaluations use model-specific sequence limits and may not align perfectly with training curves.The final limits are 4K tokens for Qwen2.5-Math-7B, 8K for Qwen3-4B and Qwen3-4B-Instruct, and 16K for DPSK-R1-Distill-1.5B.Source: C Experiment Details · p. 19
  • The implementation changes budget allocation to knapsack optimization and updates task status from external feedback; advantage values are clipped between -5 and 5 for numerical stability.Source: C Experiment Details · pp. 18–19

D.1 Visualization of Exploration Process

  • Prompt trajectories show that some prompts receive more exploration and are fully learned, while others remain challenging for both Knapsack-GRPO and GRPO.The examples include prompts in the first row, first column and second row, first column, as well as one in the third row, second column.Source: §D.1 · p. 19

D.2 Training curves

  • Knapsack-GRPO shows rapid performance improvement early in training, though a few cases of performance degeneration motivate exploring more stable policy optimization.The training curves are in Figures 12–15; final results are reported in Table 2.Source: §D.2 · p. 19

D.3 Ablation Studies

The ablations show that fallback reallocation is important for stable benchmark performance, while allocation bounds guard against degenerate or computationally costly knapsack solutions.

D.3 Ablation Studies

  • Figure 11 tracks success-rate evolution for randomly selected prompts, comparing GRPO with Knapsack-GRPO.Source: Figure 11 · p. 20
  • Without fallback, a task at 0.9 success received over 50 exploration units while an unsolved task received 2; fallback raised the unsolved task’s allocation to 29.The example covers eight prompts.Source: §D.3 · p. 20
  • Removing fallback led to unstable training, large AMC and OlympiadBench fluctuations, and degraded overall results for Qwen2.5-Math-7B.The authors suggest that neglecting challenging examples weakens the reinforcement signal and harms generalization.Source: §D.3 · p. 20
  • Nup=128 speeds dynamic-programming optimization, while Nlow=2 guards against degenerate allocations when success-rate estimates may be inaccurate.The paper reports that Nup’s specific value does not critically affect performance and says Figure 17 supports the bound choices.Source: §D.3 · p. 20

D.4 Comparing with Dynamic Sampling in DAPO

Dynamic sampling and knapsack-based exploration scale different parts of training and can be combined. The reported comparisons show gains on some benchmarks, with outcomes depending on whether computation or gradient updates are held fixed.

D.4 Comparing with Dynamic Sampling in DAPO

  • Dynamic sampling selects prompts with both positive and negative rewards, whereas knapsack-based exploration scales effective responses; the approaches are parallel and can be combined.Dynamic sampling was introduced in DAPO [48].Source: §D.4 · p. 20
  • 45.2 to 46.2: dynamic sampling improved GRPO scores on AIME and OLYMPIAD; combining it with knapsack raised AMC from 69.8 to 73.0 and yielded 46.5 total performance.The combined result was slightly better than dynamic sampling alone but worse than pure knapsack; the authors partly attribute this to knapsack-GRPO using more gradient iterations.Source: §D.4 · p. 23, Figure 18 · p. 24
  • For the same number of gradient updates, effective gradients from dynamic sampling or knapsack-based exploration lead to greater performance gains.Figure 19 compares performance against gradient-update iterations to assess the value of each update.Source: §D.4 · p. 23, Figure 19 · p. 25

Related papers

Loading 2509.25849v1…