Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

📄 arXiv: 2509.25849v1 📥 PDF

作者: 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)

关键词: 大型语言模型 强化学习 探索预算分配 背包问题 自适应学习 数学推理 计算效率

📋 核心要点

  1. 现有强化学习方法在LLM探索中采用均匀预算分配,导致简单任务饱和、困难任务失败,产生大量零梯度。
  2. 论文将探索预算分配问题建模为背包问题,提出自适应分配规则,根据任务学习状态动态调整资源。
  3. 实验表明,该方法显著提升了非零梯度比例,并在数学推理任务上取得2-9分的性能提升,计算效率提升2倍。

📝 摘要(中文)

大型语言模型(LLM)可以通过强化学习进行自我改进,在此过程中,它们生成轨迹来探索和发现更好的解决方案。然而,这种探索过程计算成本高昂,通常迫使当前方法为每个任务分配有限的探索预算。这种均匀分配会产生有问题的极端情况:简单的任务总是成功,而困难的任务总是失败,这两种情况都会在使用广泛的Group Relative Policy Optimization (GRPO)的训练更新期间产生零梯度。我们从探索预算分配的角度解决了这个问题。将每个任务的探索视为具有不同“价值”和“成本”的“项目”,我们建立了与经典背包问题的联系。这种公式允许我们推导出一种最佳分配规则,该规则可以根据模型当前的学习状态自适应地分配资源。当应用于GRPO时,我们的方法在训练期间将非零策略梯度的有效比率提高了20-40%。作为一种计算上的“免费午餐”,我们的方法可以将探索预算从学习饱和的任务重新分配到影响最大的任务。这使得能够为特别具有挑战性的问题提供更大的预算(例如,93次rollout),这在统一分配下在计算上是禁止的。这些改进转化为数学推理基准上的有意义的收益,平均改进2-4分,特定任务的峰值收益为9分。值得注意的是,使用传统的同质分配实现相当的性能将需要大约2倍的计算资源。

🔬 方法详解

问题定义:现有基于强化学习的LLM自提升方法,通常采用均匀的探索预算分配策略。这种策略的痛点在于,对于简单的任务,模型可能已经学习得很好,继续探索带来的收益很小,造成资源浪费;而对于困难的任务,有限的探索预算可能不足以让模型找到有效的解决方案,导致训练停滞,梯度消失。

核心思路:论文将探索预算分配问题类比于经典的背包问题。每个任务的探索过程被视为一个“物品”,具有一定的“价值”(例如,模型性能的提升)和“成本”(例如,计算资源消耗)。目标是在有限的预算下,最大化所有任务的总价值。核心思想是根据模型在每个任务上的学习状态,自适应地分配探索预算,将资源集中在那些最有可能带来收益的任务上。

技术框架:整体框架基于现有的强化学习自提升流程,主要改进在于探索阶段的预算分配策略。具体流程如下:1) LLM生成轨迹进行探索;2) 根据轨迹计算奖励;3) 使用Group Relative Policy Optimization (GRPO)等算法更新策略;4) 根据提出的背包算法,重新分配下一轮的探索预算。关键模块包括:价值评估模块(评估每个任务的探索价值)和预算分配模块(根据价值和成本,分配探索预算)。

关键创新:最重要的创新点在于将探索预算分配问题建模为背包问题,并提出了一种基于任务学习状态的自适应分配规则。与传统的均匀分配策略相比,该方法能够更有效地利用计算资源,提高模型的学习效率。本质区别在于,该方法不再将所有任务同等对待,而是根据其学习难度和潜在收益,进行差异化处理。

关键设计:论文的关键设计包括:1) 如何定义和评估每个任务的探索价值。一种可能的方案是使用模型在任务上的性能提升作为价值的衡量标准。2) 如何设计预算分配算法,以在满足预算约束的前提下,最大化总价值。论文可能采用动态规划或贪心算法等方法来解决背包问题。3) 如何将背包算法与现有的强化学习算法(如GRPO)相结合,以实现端到端的训练。

📊 实验亮点

实验结果表明,该方法在数学推理基准上取得了显著的性能提升,平均提升2-4分,特定任务的峰值收益达到9分。更重要的是,该方法在计算效率方面具有显著优势,达到同等性能所需的计算资源减少了约2倍。此外,该方法能够将非零策略梯度的有效比率提高20-40%,表明其能够更有效地利用训练数据。

🎯 应用场景

该研究成果可广泛应用于各种需要通过强化学习进行优化的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.