Il Meglio di N Mondi: Allineare l'Apprendimento per Rinforzo con il Campionamento Best-of-N tramite Ottimizzazione max@k

The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation

October 27, 2025

Autori: Farid Bagirov, Mikhail Arkhipov, Ksenia Sycheva, Evgeniy Glukhov, Egor Bogomolov

cs.AI

Abstract

L'applicazione del Reinforcement Learning con Ricompense Verificabili (RLVR) ai domini matematici e di programmazione ha dimostrato significativi miglioramenti nelle capacità di ragionamento e problem solving dei Large Language Model. Nonostante il suo successo nella risoluzione di problemi a generazione singola, il processo di fine-tuning con apprendimento per rinforzo può compromettere la capacità di esplorazione del modello, come evidenziato dalla ridotta diversità delle generazioni e dal conseguente deterioramento delle prestazioni durante il campionamento Best-of-N per valori elevati di N. In questo lavoro, ci concentriamo sull'ottimizzazione della metrica max@k, una generalizzazione continua di pass@k. Deriviamo una stima del gradiente on-policy non distorta per l'ottimizzazione diretta di questa metrica. Inoltre, estendiamo le nostre derivazioni agli aggiornamenti off-policy, un elemento comune negli algoritmi RLVR moderni, che consente una migliore efficienza campionaria. Empiricamente, dimostriamo che il nostro obiettivo ottimizza efficacemente la metrica max@k in scenari off-policy, allineando il modello con la strategia di inferenza Best-of-N.

English

The application of Reinforcement Learning with Verifiable Rewards (RLVR) to mathematical and coding domains has demonstrated significant improvements in the reasoning and problem-solving abilities of Large Language Models. Despite its success in single generation problem solving, the reinforcement learning fine-tuning process may harm the model's exploration ability, as reflected in decreased diversity of generations and a resulting degradation of performance during Best-of-N sampling for large N values. In this work, we focus on optimizing the max@k metric, a continuous generalization of pass@k. We derive an unbiased on-policy gradient estimate for direct optimization of this metric. Furthermore, we extend our derivations to the off-policy updates, a common element in modern RLVR algorithms, that allows better sample efficiency. Empirically, we show that our objective effectively optimizes max@k metric in off-policy scenarios, aligning the model with the Best-of-N inference strategy.