Knapsack RL: Раскрытие исследовательского потенциала больших языковых моделей через оптимизацию распределения бюджета

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

September 30, 2025

Авторы: Ziniu Li, Congliang Chen, Tianyun Yang, Tian Ding, Ruoyu Sun, Ge Zhang, Wenhao Huang, Zhi-Quan Luo

cs.AI

Аннотация

Крупные языковые модели (LLM) могут самостоятельно совершенствоваться с помощью обучения с подкреплением, генерируя траектории для исследования и поиска более эффективных решений. Однако этот процесс исследования требует значительных вычислительных ресурсов, что зачастую вынуждает современные методы выделять ограниченные бюджеты на исследование для каждой задачи. Такое равномерное распределение создает проблемные крайние случаи: простые задачи стабильно выполняются успешно, а сложные задачи стабильно терпят неудачу, что в обоих случаях приводит к нулевым градиентам при обновлениях обучения для широко используемой оптимизации групповой относительной политики (GRPO). Мы решаем эту проблему с точки зрения распределения бюджета на исследование. Рассматривая исследование каждой задачи как "элемент" с уникальной "ценностью" и "стоимостью", мы устанавливаем связь с классической задачей о рюкзаке. Такая формулировка позволяет нам вывести оптимальное правило распределения ресурсов, которое адаптивно учитывает текущее состояние обучения модели. При применении к GRPO наш метод увеличивает эффективное соотношение ненулевых градиентов политики на 20-40% в процессе обучения. Выступая в роли вычислительного "бесплатного бонуса", наш подход позволяет перераспределять бюджеты на исследование с задач, где обучение достигло насыщения, на те задачи, где оно наиболее эффективно. Это позволяет выделять значительно большие бюджеты (например, 93 прогона) для особенно сложных проблем, что было бы вычислительно неосуществимо при равномерном распределении. Эти улучшения приводят к значительным успехам на тестах математического мышления, со средним улучшением на 2-4 балла и пиковым увеличением на 9 баллов для конкретных задач. Примечательно, что достижение сопоставимой производительности с традиционным однородным распределением потребовало бы примерно в 2 раза больше вычислительных ресурсов.

English

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.