Optimierung der Testzeitberechnung durch Meta-Reinforcement-Fine-Tuning

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

March 10, 2025

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

cs.AI

Zusammenfassung

Das Training von Modellen, um die Rechenleistung zur Testzeit effektiv zu nutzen, ist entscheidend für die Verbesserung der Argumentationsleistung von LLMs. Aktuelle Methoden erreichen dies hauptsächlich durch Feinabstimmung auf Suchspuren oder das Ausführen von RL mit 0/1-Ergebnisbelohnung, aber nutzen diese Ansätze die Rechenleistung zur Testzeit effizient? Würden diese Ansätze weiterhin skalieren, wenn das Budget verbessert wird? In diesem Papier versuchen wir, diese Fragen zu beantworten. Wir formalisieren das Problem der Optimierung der Rechenleistung zur Testzeit als ein Meta-Reinforcement-Learning (RL)-Problem, das eine prinzipielle Perspektive auf die Nutzung der Rechenleistung zur Testzeit bietet. Diese Perspektive ermöglicht es uns, den langen Ausgabestrom des LLMs als mehrere Episoden zu betrachten, die zur Testzeit ausgeführt werden, und führt uns dazu, ein Konzept des kumulativen Bedauerns über Ausgabetoken als Maß für die Wirksamkeit der Rechenleistung zur Testzeit zu verwenden. Ähnlich wie RL-Algorithmen den besten Kompromiss zwischen Exploration und Exploitation während des Trainings finden können, würde die Minimierung des kumulativen Bedauerns auch das beste Gleichgewicht zwischen Exploration und Exploitation im Tokenstrom bieten. Während wir zeigen, dass state-of-the-art Modelle das Bedauern nicht minimieren, kann dies erreicht werden, indem ein dichter Belohnungsbonus in Verbindung mit der 0/1-Ergebnisbelohnung RL maximiert wird. Dieser Bonus ist der „Fortschritt“, der durch jeden nachfolgenden Block im Ausgabestrom erzielt wird, quantifiziert durch die Änderung der Wahrscheinlichkeit des letztendlichen Erfolgs. Mit diesen Erkenntnissen entwickeln wir Meta Reinforcement Fine-Tuning, oder MRT, eine neue Klasse von Feinabstimmungsmethoden zur Optimierung der Rechenleistung zur Testzeit. MRT führt zu einem relativen Leistungsgewinn von 2-3x und einem ungefähren Effizienzgewinn von 1,5x bei mathematischen Argumentationen im Vergleich zu RL mit Ergebnisbelohnung.

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.