Kleine RL-controller, groot taalmodel: RL-gestuurde adaptieve sampling voor schaling tijdens testtijd

Samenvatting

Testtijdsschaalvergroting verbetert de redeneerprestatie van grote taalmodellen, maar brengt aanzienlijke kosten met zich mee in zowel totale rekenkracht als latentie. Bestaande adaptieve bemonsteringsmethoden verhelpen dit probleem gedeeltelijk door dynamisch te beslissen wanneer te stoppen met bemonsteren, maar vertrouwen doorgaans op heuristische regels of op verdelingsaannames. In dit werk formuleren we adaptieve bemonstering als een Markov-beslissingsproces (MDP). We trainen een lichtgewicht bemonsteringscontroller met behulp van reinforcement learning (RL) om gezamenlijk de juistheid van antwoorden, latentie en rekenkosten in evenwicht te brengen. In elke ronde beslist de controller of hij stopt met bemonsteren of extra monsters neemt. Onze methode is lichtgewicht en is alleen afhankelijk van statistieken van uiteindelijke antwoorden, en kan worden getraind en ingezet op een CPU. We laten verder zien dat het resulterende raamwerk kan worden geïnterpreteerd als de Lagrangiaanse relaxatie van een beperkt optimalisatieprobleem met expliciete budgetbeperkingen. Experimenten tegen sterke basislijnen zoals ASC en ESC tonen aan dat onze methode verbeterde afwegingen bereikt tussen antwoordjuistheid, bemonsteringsronden en totale benodigde monsters.

English

Test-time scaling improves the reasoning performance of large language models but incurs substantial cost in both total computation and latency. Existing adaptive sampling methods partially mitigate this issue by dynamically deciding when to stop sampling, yet they typically rely on heuristic rules or rely on distribution assumptions. In this work, we formulate adaptive sampling as a Markov decision process (MDP). We train a lightweight sampling controller with reinforcement learning (RL) to jointly balance answer correctness, latency, and computation cost. At each round, the controller decides to stop sampling or to acquire additional samples. Our method is lightweight which only relies on statistics of final answers, and can be trained and deployed on CPU. We further show that the resulting framework admits an interpretation as the Lagrangian relaxation of a constrained optimization problem with explicit budget constraints. Experiments against strong baselines such as ASC and ESC show that our method achieves improved trade-offs among answer correctness, sampling rounds, and total samples required.