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
分类: cs.LG, cs.AI, cs.CL
发布日期: 2025-09-30
💡 一句话要点
Knapsack RL：通过优化预算分配解锁LLM的探索能力
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 强化学习 大型语言模型 探索策略 预算分配 背包问题
📋 核心要点
- 现有强化学习方法在LLM探索中采用均匀预算分配，导致简单任务饱和、困难任务失败，产生大量零梯度，影响学习效率。
- 论文将探索预算分配问题建模为背包问题，根据任务的价值和成本自适应分配资源，优化非零梯度比例。
- 实验表明，该方法在数学推理任务上显著提升性能，平均提升2-4分，最高提升9分，且计算效率更高。
📝 摘要（中文）
大型语言模型（LLM）可以通过强化学习进行自我改进，在此过程中，它们生成轨迹以探索和发现更好的解决方案。然而，这种探索过程计算成本高昂，通常迫使当前方法为每个任务分配有限的探索预算。这种均匀分配会产生有问题的极端情况：简单的任务总是成功，而困难的任务总是失败，这两种情况都会在使用广泛的Group Relative Policy Optimization（GRPO）的训练更新期间产生零梯度。我们从探索预算分配的角度解决了这个问题。将每个任务的探索视为具有不同“价值”和“成本”的“项目”，我们建立了与经典背包问题的联系。这种公式允许我们推导出一种最佳分配规则，该规则可以根据模型当前的学习状态自适应地分配资源。当应用于GRPO时，我们的方法在训练期间将非零策略梯度的有效比率提高了20-40%。作为一种计算上的“免费午餐”，我们的方法可以将探索预算从学习饱和的任务重新分配到影响最大的任务。这使得能够为特别具有挑战性的问题分配更大的预算（例如，93次rollout），这在均匀分配下在计算上是禁止的。这些改进转化为数学推理基准上的有意义的收益，平均改进2-4分，特定任务的峰值收益为9分。值得注意的是，使用传统的同质分配实现相当的性能将需要大约2倍的计算资源。
🔬 方法详解
问题定义：现有基于强化学习的LLM自提升方法，通常采用均匀的探索预算分配策略。这种策略的痛点在于，对于简单的任务，模型可能已经达到饱和，继续探索带来的收益很小；而对于困难的任务，有限的探索预算又不足以让模型找到有效的解决方案，导致训练过程中产生大量的零梯度，降低了学习效率。
核心思路：论文的核心思路是将探索预算的分配问题，类比于经典的背包问题。每个任务的探索过程被视为一个“物品”，具有特定的“价值”（探索带来的性能提升）和“成本”（计算资源消耗）。目标是在有限的预算下，选择一组“物品”（任务），使得总价值最大化。通过这种方式，可以将更多的预算分配给那些能够带来更大性能提升的任务。
技术框架：整体框架是在现有的强化学习训练流程中，引入一个自适应的预算分配模块。该模块根据模型在每个任务上的学习状态（例如，梯度大小、奖励变化等），动态调整分配给每个任务的探索预算。具体流程如下： 1. LLM生成轨迹进行探索。 2. 根据轨迹计算每个任务的价值和成本。 3. 使用背包问题的求解器，确定最优的预算分配方案。 4. 根据新的预算分配方案，重新进行探索和训练。
关键创新：论文最重要的创新点在于将探索预算分配问题，与经典的背包问题联系起来，并提出了一种基于背包问题的自适应预算分配策略。与传统的均匀分配策略相比，该方法能够更有效地利用计算资源，提高模型的学习效率。
关键设计：论文的关键设计包括： 1. 价值函数的设计：如何准确评估每个任务的探索价值，例如使用梯度范数或奖励变化作为指标。 2. 成本函数的设计：如何衡量每个任务的探索成本，通常与rollout的数量成正比。 3. 背包问题求解器的选择：可以使用动态规划或贪心算法等方法来求解背包问题。 4. 探索预算的上下限设置：为了避免过度探索或探索不足，需要设置合理的预算上下限。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Knapsack RL方法在数学推理基准测试中取得了显著的性能提升。与传统的均匀分配策略相比，平均提升了2-4个点，在特定任务上甚至提升了9个点。更重要的是，为了达到相同的性能水平，均匀分配策略需要大约2倍的计算资源。此外，该方法还能够将非零策略梯度的比例提高20-40%，表明其能够更有效地利用训练数据。
🎯 应用场景
该研究成果可广泛应用于各种需要通过强化学习进行自提升的LLM应用场景，例如数学推理、代码生成、对话系统等。通过更有效地利用计算资源，可以显著提升LLM的性能和效率，降低训练成本，加速LLM的开发和部署。此外，该方法也可以推广到其他机器学习模型的训练中，例如图像识别、自然语言处理等。
📄 摘要（原文）
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.