Knapsack RL: Desbloqueando la exploración de LLM mediante la optimización de la asignación de presupuesto

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

September 30, 2025

Autores: Ziniu Li, Congliang Chen, Tianyun Yang, Tian Ding, Ruoyu Sun, Ge Zhang, Wenhao Huang, Zhi-Quan Luo

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) pueden automejorarse mediante el aprendizaje por refuerzo, donde generan trayectorias para explorar y descubrir soluciones mejores. Sin embargo, este proceso de exploración es computacionalmente costoso, lo que a menudo obliga a los métodos actuales a asignar presupuestos de exploración limitados a cada tarea. Esta asignación uniforme crea casos problemáticos: las tareas fáciles tienen éxito consistentemente, mientras que las tareas difíciles fracasan sistemáticamente, ambos casos produciendo gradientes nulos durante las actualizaciones de entrenamiento para la Optimización de Política Relativa de Grupo (GRPO, por sus siglas en inglés), ampliamente utilizada. Abordamos este problema desde la perspectiva de la asignación del presupuesto de exploración. Al considerar la exploración de cada tarea como un "ítem" con un "valor" y un "costo" distintivos, establecemos una conexión con el clásico problema de la mochila. Esta formulación nos permite derivar una regla de asignación óptima que distribuye los recursos de manera adaptativa según el estado actual de aprendizaje del modelo. Cuando se aplica a GRPO, nuestro método aumenta la proporción efectiva de gradientes de política no nulos en un 20-40% durante el entrenamiento. Funcionando como un "almuerzo gratis" computacional, nuestro enfoque permite reasignar los presupuestos de exploración desde tareas donde el aprendizaje está saturado hacia aquellas donde es más impactante. Esto habilita presupuestos significativamente mayores (por ejemplo, 93 rollouts) para problemas especialmente desafiantes, lo que sería computacionalmente prohibitivo bajo una asignación uniforme. Estas mejoras se traducen en avances significativos en benchmarks de razonamiento matemático, con mejoras promedio de 2-4 puntos y ganancias máximas de 9 puntos en tareas específicas. Cabe destacar que lograr un rendimiento comparable con la asignación homogénea tradicional requeriría aproximadamente el doble de recursos computacionales.

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.