Knapsack RL: A Computational 'Free Lunch' for Training Smarter Language Models

Large Language Models (LLMs) have shown an extraordinary capacity to improve themselves through reinforcement learning (RL). By generating solutions, receiving feedback, and adjusting their strategy, they can learn to tackle complex problems such as advanced mathematical reasoning. This process hinges on a critical step: exploration—trying many different approaches, or “rollouts,” to discover what works.

The catch? Exploration is computationally expensive. Generating thousands of possible solutions for thousands of different problems consumes massive amounts of GPU time. To keep costs manageable, current methods typically assign a small, fixed exploration budget to every problem—often 8 possible solutions per task.

This one-size-fits-all approach hides a serious flaw:

• For an easy problem, 8 attempts may be overkill, wasting compute.
• For a hard problem, 8 attempts may be hopelessly insufficient, leading to repeated failure.

In both cases, the model can end up always succeeding or always failing on a given task. This produces a “zero gradient”—no learning signal—wasting precious computational resources.

A recent paper, “Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation”, tackles this problem head-on. The researchers propose: treat exploration budget allocation as a classic optimization problem. Frame each task as an item with a cost (computation) and value (learning potential). Then, use the well-known knapsack problem to intelligently distribute a fixed budget, giving more resources to the tasks that need them most.

Figure 1 Illustration of our framework for allocating exploration budgets among tasks. Tasks are modeled as items with learning value and computational cost, and allocation is optimized via knapsack.

---

Background: The Zero-Gradient Problem in Policy Optimization

To see the significance of this idea, let’s look at RL fine-tuning for LLMs. The goal is to train model parameters \(\theta\) to generate responses \(y\) for prompt \(x\) that maximize a reward:

\[ \max_{\theta} \ \mathbb{E}_{y \sim \pi_{\theta}(\cdot|x)}\,[\,r(x,y)\,] \]

In mathematical reasoning tasks, the reward is often binary:

\[ r(x,y) = \mathbb{I}(\text{answer is correct}) \]

A popular RL algorithm is Group Relative Policy Optimization (GRPO). For a batch of \(M\) prompts, the model generates \(N\) rollouts for each prompt. The gradient update is:

\[ g(\theta) = \sum_{i=1}^{M} \sum_{j=1}^{N} \nabla_{\theta} \log \pi_{\theta}(y_{ij}|x_i) \cdot (\,r(x_i,y_{ij}) - b_i\,) \cdot c_i \]

Here:

• \(b_i\) — average reward for prompt \(i\) (baseline).
• \(r - b_i\) — “advantage” for the rollout.

If a response’s reward is above average, the model increases its probability; if below average, it decreases it.

---

Critical issue: If all \(N\) rollouts for a prompt have the same reward (all successes or all failures):

1. All successes: Rewards = 1, baseline \(b_i\) = 1, advantage = 0 → gradient = 0.
2. All failures: Rewards = 0, baseline \(b_i\) = 0, advantage = 0 → gradient = 0.

In both cases, the model spends computation but learns nothing.

---

Measuring the Inefficiency of Uniform Budgets

The authors introduce effective-gradient-ratio: the fraction of generated responses actually contributing a non-zero gradient:

\[ \text{effective-gradient-ratio} = \frac{1}{M N} \sum_{i=1}^{M} \sum_{j=1}^N \mathbb{I}(g_{i,j} \neq 0) \]

Tracking this during training for a 7B-parameter math model (uniform budget \(N=8\)) revealed worrying trends.

Figure 2 Effective gradient ratio and zero-gradient breakdown during training. Early training: high all-failure rate; late training: high all-success rate.

Observations:

• Early Training: Model fails most tasks; all-failure prompts dominate.
• Mid Training: Mix of successes and failures yields highest efficiency.
• Late Training: Easy tasks dominate all-success prompts; hard tasks persist as all-failure; efficiency drops sharply.

---

How Many Rollouts Do You Really Need?

Probability of a non-zero gradient for prompt success rate \(p_i\):

\[ \text{Prob} = 1 - p_i^{N} - (1 - p_i)^{N} \]

For:

• \(p = 0.5\): Average 3 rollouts.
• \(p = 0.01\): Average 100 rollouts; 229 rollouts for 90% certainty.

Figure 3 Rollouts needed for non-zero gradient, grouped by success rate.

Uniform \(N=8\) budgets are only effective for \(p \in [0.1, 0.9]\).

This is the computation–exploration dilemma:

• Raise \(N\) uniformly → huge cost.
• Filter easy/hard tasks → lose training opportunities on challenging problems.

---

The Knapsack RL Solution

Instead of uniform budgets, reallocate fixed total compute more intelligently.

RL Exploration Setting	Knapsack Analogy
Task with exploration budget	Item in knapsack
Number of rollouts	Item weight
Learning benefit from budget	Item value
Total available computation	Knapsack capacity
Available GPUs	Knapsack physical container

Table 1 Connection between RL exploration and the knapsack problem.

---

Defining Task “Value”

The value of allocating \(N_i\) rollouts to task \(i\):

\[ \text{Value}(N_i, p_i) = \text{ProbNonZeroGradient}(N_i, p_i) \times \text{InfoGain}(p_i) \]

• \(\text{ProbNonZeroGradient}\): \(1 - p_i^{N_i} - (1-p_i)^{N_i}\) — likelihood of mixed successes/failures.
• \(\text{InfoGain}(p_i)\): Estimated learning potential, approximated as \(p_i(1-p_i)^2\).

Implications:

• Low value for very easy (\(p_i \to 1\)) or very hard (\(p_i \to 0\)) tasks.
• Peak value near \(p_i = 1/3\) — “learning sweet spot.”

Figure 4 Relationship between success rate, budget, and learning value.

---

Implementation Highlights

• Estimating success rates: Use prior step’s observed rates as estimates.
• Fallback strategy: Reallocate freed budget from easy tasks to extremely hard ones (\(p=0\)) to keep them in play.
• Rollout balancing: Distribute allocated rollouts evenly among workers to avoid idle GPUs.

---

Experiments: Proof of Concept

Setup: Compare Knapsack-GRPO to baseline GRPO across multiple LLMs for math-intensive benchmarks (AIME, AMC, MATH, MINERVA, OLYMPIAD) and GPQA.

Model	AIME	AMC	MATH	MINERVA	OLYMPIAD	GPQA	Avg
DPSK-R1-Distill-1.5B	25.3	62.1	81.4	25.8	41.7	39.1	42.9
+ GRPO	27.6	71.1	84.0	27.6	46.4	36.7	45.9
+ Knapsack-GRPO	34.0	75.1	86.7	28.5	49.7	40.3	49.7
Qwen3-4B-Base	6.6	29.9	48.0	19.4	23.1	26.4	22.9
+ GRPO	20.7	56.9	80.6	31.9	44.9	46.6	43.2
+ Knapsack-GRPO	20.8	66.0	81.0	35.7	46.2	45.5	45.1
Qwen2.5-Math-7B	12.3	41.0	61.2	11.8	26.1	22.0	26.7
+ GRPO	23.9	70.6	81.7	33.6	41.9	40.8	45.2
+ Knapsack-GRPO	24.3	77.4	83.9	34.5	44.1	43.8	47.5

Table 2 Performance comparison across models and benchmarks. Knapsack-GRPO consistently leads.

---

Understanding the Gains

Dynamic allocation: While average \(N=8\), hard tasks sometimes receive up to 93 rollouts.

Figure 5 Distribution of budgets during training — long tail for challenging tasks.

Gradient efficiency boost: Effective gradient ratio improves by 20–40%.

Figure 6 Knapsack-GRPO maintains high efficiency over training.

Better learning dynamics: More transitions from hard to easy tasks.

Figure 7 Prompt transition matrices — larger proportion reaching extremely-easy.

Figure 8 Final distribution of task difficulty — Knapsack-GRPO yields more mastered tasks.

---

The 2× Efficiency “Free Lunch”

Figure 9 Performance vs total budget. Knapsack-GRPO matches GRPO with half the compute.

---

Conclusion and Future Directions

The Knapsack RL framework delivers a clear message: optimize exploration allocation to match task difficulty, and you unlock both performance and efficiency — without extra compute.

Future opportunities:

• Richer value functions: Model learning potential more precisely.
• Multi-dimensional costs: Include text length or conversation depth.
• Advanced exploration integration: Combine with techniques like tree search (AlphaGo).

By treating computation as a resource to optimize rather than just spend, Knapsack RL offers a sustainable path to training smarter, more capable language models.