Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation (2509.25849v1)

Abstract: 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.

• The paper introduces a knapsack-based RL method that dynamically allocates exploration budgets to tasks based on their difficulty.
• It employs a dynamic programming approach to optimize rollout allocations, increasing effective gradient ratios by 20–40% during training.
• The method reduces wasted compute on both easy and hard tasks, achieving up to 9-point performance gains over standard GRPO.

Knapsack RL: Optimizing Exploration Budget Allocation for LLM Reinforcement Learning

Introduction and Motivation

The paper "Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation" (2509.25849) addresses a critical inefficiency in reinforcement learning (RL) for LLMs: the uniform allocation of exploration budgets across tasks of varying difficulty. In standard RL fine-tuning pipelines, each prompt receives a fixed number of rollouts, regardless of its inherent challenge. This leads to two major issues: (1) easy tasks are oversampled, yielding all-success outcomes and thus zero gradients, and (2) hard tasks are undersampled, resulting in all-failure outcomes and again zero gradients. Both cases waste computational resources and stall learning, particularly in algorithms like Group Relative Policy Optimization (GRPO).

The authors propose a principled solution by formulating exploration budget allocation as a knapsack optimization problem. Each task is modeled as an item with a value (learning potential) and a cost (computational effort), and the global objective is to maximize total learning value under a fixed compute budget. This approach, termed Knapsack RL, adaptively reallocates exploration resources to maximize the effective gradient signal during RL training.

Figure 1: Illustration of the knapsack-based framework for allocating exploration budgets among tasks, modeling each as an item with learning value and computational cost.

Theoretical Analysis of Exploration Budget Requirements

The paper provides a rigorous analysis of the exploration budget required to obtain non-zero gradients in GRPO. For a prompt with success rate $p$, the probability of observing both a success and a failure in $N$ rollouts is $1 - p^N - (1-p)^N$. The expected number of rollouts until both outcomes are observed is $1/p + 1/(1-p) - 1$. This analysis reveals that for hard tasks ($p \ll 1$), the required number of rollouts to obtain a learning signal can be orders of magnitude larger than the uniform budget typically used in practice.

Figure 2: Exploration budget required to ensure non-zero gradients as a function of success rate, highlighting the inefficiency of uniform allocation for hard or easy tasks.

Empirical results confirm that with standard budgets (e.g., $N=8$), a large fraction of prompts yield zero gradients, especially at the early and late stages of training. This leads to a low effective gradient ratio, with over 40% of samples failing to contribute to model updates.

Knapsack-Based Budget Allocation Framework

The core contribution is the formalization of exploration budget allocation as a discrete knapsack problem. For each prompt $x_i$, the value of assigning $N_i$ rollouts is defined as:

$$\operatorname{Value}(N_i, p_i) = \operatorname{ProbNonZeroGradient}(N_i, p_i) \times \operatorname{InfoGain}(p_i)$$

where $$\operatorname{ProbNonZeroGradient}(N_i, p_i) = 1 - p_i^{N_i} - (1-p_i)^{N_i}$$, and $\operatorname{InfoGain}(p_i) \approx p_i(1-p_i)^2$ (from a first-order Taylor expansion of policy improvement). The optimization seeks to maximize the sum of values across all prompts, subject to a total rollout budget constraint.

The allocation is solved efficiently via dynamic programming, and the framework is compatible with existing RL pipelines, requiring only minor modifications to the sampling and batching logic.

Implementation and Practical Considerations

Success rates $p_i$ are estimated online using statistics from previous epochs. To prevent degenerate allocations (e.g., zero budget to unsolved or always-solved prompts), the framework enforces lower and upper bounds on $N_i$ and introduces a fallback strategy to ensure hard tasks are not starved of exploration. Rollout balancing is addressed by randomizing job dispatch to parallel workers, maintaining high GPU utilization.

A pseudocode implementation is provided in the appendix, demonstrating the minimal changes required to integrate knapsack-based allocation into standard RL training loops.

Empirical Results and Analysis

Extensive experiments are conducted on mathematical reasoning benchmarks (AIME, AMC, MATH, MINERVA, OLYMPIAD, GPQA) using Qwen and DeepSeek model families (1B–7B parameters). The main findings are:

• Knapsack-GRPO consistently outperforms standard GRPO, with average improvements of 2–4 points and peak gains of 9 points on specific tasks.
• The effective gradient ratio is increased by 20–40% during training, directly translating to more reliable policy improvements.
• Comparable performance to uniform allocation requires approximately 2x the computational resources, demonstrating the efficiency of the knapsack approach.
• The method dynamically assigns up to 93 rollouts to hard tasks, which is infeasible under uniform allocation.

  Figure 3: Learning dynamics of randomly selected prompts, showing that Knapsack-GRPO enables more effective learning on hard tasks compared to GRPO.

  

  Figure 4: Evaluation performance of DPSK-R1-Distill-1.5B across training iterations, illustrating rapid and stable improvements with Knapsack-GRPO.

  

  Figure 5: Evaluation performance of Qwen2.5-Math-7B across training iterations, further confirming the benefits of adaptive budget allocation.

Ablation studies confirm the necessity of the fallback strategy and the robustness of the method to hyperparameter choices for lower and upper bounds on $N_i$.

Comparison with Dynamic Sampling and Related Work

The paper contrasts knapsack-based allocation with dynamic sampling (as in DAPO), which filters prompts to maximize the effective gradient ratio. While dynamic sampling improves sample efficiency, it operates by selecting which prompts to train on, whereas Knapsack RL optimizes how much exploration to allocate to each prompt. The two approaches are complementary and can be combined, but Knapsack RL provides superior compute efficiency when measured by total exploration budget.

Figure 6: Performance of Qwen2.5-Math-7B relative to the number of exploration iterations, demonstrating efficient conversion of compute into performance gains with Knapsack RL.

Figure 7: Performance as a function of the number of LLM gradient updates, validating that effective gradients from knapsack-based allocation yield greater improvements per update.

Implications and Future Directions

The knapsack-based framework provides a principled and practical solution to the exploration–computation dilemma in LLM RL. By maximizing the effective use of fixed computational resources, it enables more rapid and robust self-improvement, particularly on challenging tasks that are critical for advancing reasoning capabilities.

The approach is modular and can be extended to other RL algorithms by adapting the value function and non-zero gradient probability. Future research directions include:

• Incorporating richer value functions that account for response length, multi-turn interactions, or more accurate policy improvement estimates.
• Integrating advanced exploration strategies, such as tree-based search or experience replay, within the knapsack allocation framework.
• Applying the method to broader domains beyond mathematical reasoning, including code generation and agentic tasks.

Conclusion

Knapsack RL introduces a theoretically grounded and empirically validated method for optimizing exploration budget allocation in LLM RL. By modeling the allocation as a knapsack problem, the framework adaptively prioritizes tasks with the highest learning potential, leading to substantial improvements in sample efficiency and downstream performance. The method is readily deployable in existing RL pipelines and opens new avenues for scaling LLM self-improvement under realistic compute constraints.