O Melhor de N Mundos: Alinhando Aprendizagem por Reforço com Amostragem Best-of-N através da Optimização 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

Resumo

A aplicação de Aprendizagem por Reforço com Recompensas Verificáveis (RLVR) em domínios matemáticos e de programação demonstrou melhorias significativas nas capacidades de raciocínio e resolução de problemas de Modelos de Linguagem de Grande Porte. Apesar do seu sucesso na resolução de problemas de geração única, o processo de afinação por aprendizagem por reforço pode prejudicar a capacidade de exploração do modelo, conforme refletido na diminuição da diversidade das gerações e numa consequente degradação do desempenho durante a amostragem Melhor-de-N para valores grandes de N. Neste trabalho, focamo-nos na otimização da métrica max@k, uma generalização contínua de pass@k. Derivamos uma estimativa de gradiente não enviesada e on-policy para a otimização direta desta métrica. Adicionalmente, estendemos as nossas derivações para atualizações off-policy, um elemento comum nos algoritmos modernos de RLVR, que permite uma melhor eficiência amostral. Empiricamente, mostramos que o nosso objetivo otimiza eficazmente a métrica max@k em cenários off-policy, alinhando o modelo com a estratégia de inferência Melhor-de-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.