Lo Mejor de N Mundos: Alineando el Aprendizaje por Refuerzo con el Muestreo Best-of-N mediante Optimización max@k

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

October 27, 2025

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

cs.AI

Resumen

La aplicación del Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) en dominios matemáticos y de programación ha demostrado mejoras significativas en las capacidades de razonamiento y resolución de problemas de los Modelos de Lenguaje a Gran Escala. A pesar de su éxito en la resolución de problemas de generación única, el proceso de ajuste fino mediante aprendizaje por refuerzo puede perjudicar la capacidad de exploración del modelo, lo cual se refleja en una disminución de la diversidad de las generaciones y una consiguiente degradación del rendimiento durante el muestreo Best-of-N para valores grandes de N. En este trabajo, nos centramos en optimizar la métrica max@k, una generalización continua de pass@k. Derivamos una estimación de gradiente insesgada y *on-policy* para la optimización directa de esta métrica. Además, extendemos nuestras derivaciones a las actualizaciones *off-policy*, un elemento común en los algoritmos modernos de RLVR que permite una mejor eficiencia muestral. Empíricamente, demostramos que nuestro objetivo optimiza efectivamente la métrica max@k en escenarios *off-policy*, alineando el modelo con la estrategia de inferencia 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.