Knapsack RL: Sbloccare l'esplorazione dei modelli linguistici di grandi dimensioni attraverso l'ottimizzazione dell'allocazione del budget

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

September 30, 2025

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

cs.AI

Abstract

I Large Language Model (LLM) possono auto-migliorarsi attraverso l'apprendimento per rinforzo, generando traiettorie per esplorare e scoprire soluzioni migliori. Tuttavia, questo processo di esplorazione è computazionalmente costoso, spesso costringendo i metodi attuali a assegnare budget di esplorazione limitati a ciascun task. Questa allocazione uniforme crea casi problematici: i task facili riescono sistematicamente mentre quelli difficili falliscono costantemente, entrambi producendo gradienti nulli durante gli aggiornamenti di addestramento per il diffuso Group Relative Policy Optimization (GRPO). Affrontiamo questo problema attraverso la lente dell'allocazione del budget di esplorazione. Considerando l'esplorazione di ciascun task come un "oggetto" con un "valore" e un "costo" distinti, stabiliamo una connessione con il classico problema dello zaino. Questa formulazione ci permette di derivare una regola di assegnazione ottimale che distribuisce le risorse in modo adattivo in base allo stato attuale di apprendimento del modello. Quando applicato al GRPO, il nostro metodo aumenta il rapporto effettivo di gradienti di policy non nulli del 20-40% durante l'addestramento. Agendo come un "pranzo gratuito" computazionale, il nostro approccio può riallocare i budget di esplorazione dai task in cui l'apprendimento è saturo a quelli in cui è più impattante. Ciò consente budget significativamente più ampi (ad esempio, 93 rollout) per problemi particolarmente impegnativi, che sarebbero computazionalmente proibitivi con un'allocazione uniforme. Questi miglioramenti si traducono in guadagni significativi sui benchmark di ragionamento matematico, con miglioramenti medi di 2-4 punti e picchi di 9 punti su task specifici. È importante notare che il raggiungimento di prestazioni comparabili con l'allocazione omogenea tradizionale richiederebbe circa il doppio delle risorse computazionali.

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.