Otimização do Cálculo em Tempo de Teste por meio de Ajuste Fino com Meta Reforço

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

March 10, 2025

Autores: Yuxiao Qu, Matthew Y. R. Yang, Amrith Setlur, Lewis Tunstall, Edward Emanuel Beeching, Ruslan Salakhutdinov, Aviral Kumar

cs.AI

Resumo

Treinar modelos para utilizar efetivamente o poder computacional durante o teste é crucial para melhorar o desempenho de raciocínio dos LLMs (Large Language Models). Os métodos atuais fazem isso principalmente por meio de ajuste fino em traços de busca ou executando RL (Reinforcement Learning) com recompensa de resultado 0/1, mas essas abordagens utilizam eficientemente o poder computacional durante o teste? Elas continuariam a escalar conforme o orçamento melhora? Neste artigo, tentamos responder a essas perguntas. Formalizamos o problema de otimizar o poder computacional durante o teste como um problema de meta-reinforcement learning (RL), o que fornece uma perspectiva fundamentada sobre como gastar esse poder computacional. Essa perspectiva nos permite ver o fluxo longo de saída do LLM como consistindo de vários episódios executados durante o teste e nos leva a usar uma noção de arrependimento cumulativo sobre os tokens de saída como uma forma de medir a eficácia do poder computacional durante o teste. Assim como os algoritmos de RL podem equilibrar melhor a exploração e a exploração durante o treinamento, minimizar o arrependimento cumulativo também forneceria o melhor equilíbrio entre exploração e exploração no fluxo de tokens. Embora mostremos que os modelos de última geração não minimizam o arrependimento, é possível fazê-lo maximizando um bônus de recompensa densa em conjunto com o RL de recompensa de resultado 0/1. Esse bônus é o "progresso" feito por cada bloco subsequente no fluxo de saída, quantificado pela mudança na probabilidade de sucesso eventual. Usando essas ideias, desenvolvemos o Meta Reinforcement Fine-Tuning, ou MRT, uma nova classe de métodos de ajuste fino para otimizar o poder computacional durante o teste. O MRT resulta em um ganho relativo de 2-3x no desempenho e um ganho de aproximadamente 1,5x na eficiência de tokens para raciocínio matemático em comparação com o RL de recompensa de resultado.

English

Training models to effectively use test-time compute is crucial for improving the reasoning performance of LLMs. Current methods mostly do so via fine-tuning on search traces or running RL with 0/1 outcome reward, but do these approaches efficiently utilize test-time compute? Would these approaches continue to scale as the budget improves? In this paper, we try to answer these questions. We formalize the problem of optimizing test-time compute as a meta-reinforcement learning (RL) problem, which provides a principled perspective on spending test-time compute. This perspective enables us to view the long output stream from the LLM as consisting of several episodes run at test time and leads us to use a notion of cumulative regret over output tokens as a way to measure the efficacy of test-time compute. Akin to how RL algorithms can best tradeoff exploration and exploitation over training, minimizing cumulative regret would also provide the best balance between exploration and exploitation in the token stream. While we show that state-of-the-art models do not minimize regret, one can do so by maximizing a dense reward bonus in conjunction with the outcome 0/1 reward RL. This bonus is the ''progress'' made by each subsequent block in the output stream, quantified by the change in the likelihood of eventual success. Using these insights, we develop Meta Reinforcement Fine-Tuning, or MRT, a new class of fine-tuning methods for optimizing test-time compute. MRT leads to a 2-3x relative gain in performance and roughly a 1.5x gain in token efficiency for math reasoning compared to outcome-reward RL.