Optimalisatie van Testtijdberekening via Meta Reinforcement Fine-Tuning

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

March 10, 2025

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

cs.AI

Samenvatting

Het trainen van modellen om rekentijd tijdens het testen effectief te benutten, is cruciaal voor het verbeteren van het redeneervermogen van LLM's (Large Language Models). Huidige methoden doen dit voornamelijk via fine-tuning op zoeksporen of door reinforcement learning (RL) uit te voeren met een 0/1 beloning voor het resultaat, maar maken deze benaderingen efficiënt gebruik van de rekentijd tijdens het testen? Zouden deze benaderingen blijven schalen naarmate het budget toeneemt? In dit artikel proberen we deze vragen te beantwoorden. We formaliseren het probleem van het optimaliseren van rekentijd tijdens het testen als een meta-reinforcement learning (RL) probleem, wat een principieel perspectief biedt op het besteden van rekentijd tijdens het testen. Dit perspectief stelt ons in staat om de lange uitvoerstroom van de LLM te zien als bestaande uit meerdere episodes die tijdens het testen worden uitgevoerd, en leidt ons ertoe om een notie van cumulatieve spijt over uitvoertokens te gebruiken als een manier om de effectiviteit van rekentijd tijdens het testen te meten. Net zoals RL-algoritmen de beste afweging kunnen maken tussen exploratie en exploitatie tijdens de training, zou het minimaliseren van cumulatieve spijt ook de beste balans bieden tussen exploratie en exploitatie in de tokenstroom. Hoewel we aantonen dat state-of-the-art modellen geen spijt minimaliseren, kan dit wel worden bereikt door een dichte beloningsbonus te maximaliseren in combinatie met de 0/1 resultaatbeloning van RL. Deze bonus is de "voortgang" die wordt gemaakt door elk volgend blok in de uitvoerstroom, gekwantificeerd door de verandering in de waarschijnlijkheid van uiteindelijk succes. Met deze inzichten ontwikkelen we Meta Reinforcement Fine-Tuning, of MRT, een nieuwe klasse van fine-tuning methoden voor het optimaliseren van rekentijd tijdens het testen. MRT leidt tot een relatieve prestatieverbetering van 2-3x en een ruwweg 1.5x hogere token-efficiëntie voor wiskundig redeneren in vergelijking met RL met resultaatbeloning.

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.