Optimización del Cómputo en Tiempo de Prueba mediante Ajuste Fino con Meta-Aprendizaje por Refuerzo

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

Resumen

Entrenar modelos para utilizar eficazmente el cómputo en tiempo de prueba es crucial para mejorar el rendimiento en razonamiento de los LLM. Los métodos actuales lo hacen principalmente mediante ajuste fino en trazas de búsqueda o ejecutando RL con recompensas de resultado 0/1, pero ¿utilizan estos enfoques eficientemente el cómputo en tiempo de prueba? ¿Seguirían escalando estos enfoques a medida que aumenta el presupuesto? En este artículo, intentamos responder estas preguntas. Formalizamos el problema de optimizar el cómputo en tiempo de prueba como un problema de meta-aprendizaje por refuerzo (RL), lo cual proporciona una perspectiva fundamentada sobre cómo gastar el cómputo en tiempo de prueba. Esta perspectiva nos permite ver el flujo extenso de salida del LLM como compuesto por varios episodios ejecutados en tiempo de prueba y nos lleva a utilizar una noción de arrepentimiento acumulado sobre los tokens de salida como una forma de medir la eficacia del cómputo en tiempo de prueba. Similar a cómo los algoritmos de RL pueden equilibrar mejor la exploración y la explotación durante el entrenamiento, minimizar el arrepentimiento acumulado también proporcionaría el mejor equilibrio entre exploración y explotación en el flujo de tokens. Aunque demostramos que los modelos de vanguardia no minimizan el arrepentimiento, es posible lograrlo maximizando una recompensa densa adicional junto con la recompensa de resultado 0/1 en RL. Esta recompensa adicional es el "progreso" realizado por cada bloque subsiguiente en el flujo de salida, cuantificado por el cambio en la probabilidad de éxito eventual. Utilizando estas ideas, desarrollamos Meta Reinforcement Fine-Tuning, o MRT, una nueva clase de métodos de ajuste fino para optimizar el cómputo en tiempo de prueba. MRT conduce a una mejora relativa de 2-3x en el rendimiento y aproximadamente 1.5x en eficiencia de tokens para razonamiento matemático en comparación con RL basado en recompensas 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.