Knapsack RL : Optimiser l'allocation budgétaire pour débloquer l'exploration des LLM

papers.abstract

Les grands modèles de langage (LLMs) peuvent s'auto-améliorer grâce à l'apprentissage par renforcement, où ils génèrent des trajectoires pour explorer et découvrir de meilleures solutions. Cependant, ce processus d'exploration est coûteux en calcul, contraignant souvent les méthodes actuelles à allouer des budgets d'exploration limités à chaque tâche. Cette allocation uniforme crée des cas problématiques : les tâches faciles réussissent systématiquement tandis que les tâches difficiles échouent tout aussi systématiquement, produisant tous deux des gradients nuls lors des mises à jour de l'entraînement pour l'Optimisation de Politique Relative par Groupe (GRPO) largement utilisée. Nous abordons ce problème sous l'angle de l'allocation du budget d'exploration. En considérant l'exploration de chaque tâche comme un "objet" ayant une "valeur" et un "coût" distincts, nous établissons un lien avec le problème classique du sac à dos. Cette formulation nous permet de dériver une règle d'affectation optimale qui répartit les ressources de manière adaptative en fonction de l'état actuel d'apprentissage du modèle. Appliquée à la GRPO, notre méthode augmente le ratio effectif de gradients de politique non nuls de 20 à 40 % pendant l'entraînement. Agissant comme un "repas gratuit" computationnel, notre approche permet de réallouer les budgets d'exploration des tâches où l'apprentissage est saturé vers celles où il est le plus impactant. Cela permet d'attribuer des budgets significativement plus importants (par exemple, 93 déploiements) à des problèmes particulièrement complexes, ce qui serait prohibitif en termes de calcul avec une allocation uniforme. Ces améliorations se traduisent par des gains significatifs sur des benchmarks de raisonnement mathématique, avec des améliorations moyennes de 2 à 4 points et des gains maximaux de 9 points sur des tâches spécifiques. Il est à noter qu'atteindre des performances comparables avec une allocation homogène traditionnelle nécessiterait environ deux fois plus de ressources computationnelles.

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.

Knapsack RL : Optimiser l'allocation budgétaire pour débloquer l'exploration des LLM

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

papers.abstract

Support