Kleiner RL-Controller, großes Sprachmodell: RL-gesteuertes adaptives Sampling für die Testzeit-Skalierung

Zusammenfassung

Die Skalierung zur Testzeit verbessert die Denkleistung großer Sprachmodelle, verursacht jedoch erhebliche Kosten sowohl in Bezug auf die Gesamtberechnung als auch auf die Latenz. Bestehende adaptive Sampling-Methoden mildern dieses Problem teilweise, indem sie dynamisch entscheiden, wann das Sampling beendet werden soll, doch stützen sie sich typischerweise auf heuristische Regeln oder auf Verteilungsannahmen. In dieser Arbeit formulieren wir adaptives Sampling als einen Markov-Entscheidungsprozess (MDP). Wir trainieren einen leichten Sampling-Controller mit bestärkendem Lernen (RL), um gleichzeitig Antwortkorrektheit, Latenz und Rechenaufwand auszubalancieren. In jeder Runde entscheidet der Controller, ob das Sampling beendet oder ob zusätzliche Stichproben eingeholt werden sollen. Unsere Methode ist leichtgewichtig, da sie nur auf Statistiken der endgültigen Antworten beruht, und kann auf CPU trainiert und eingesetzt werden. Wir zeigen weiterhin, dass das resultierende Framework als Lagrange-Relaxation eines Optimierungsproblems unter Nebenbedingungen mit expliziten Budgetbeschränkungen interpretiert werden kann. Experimente gegen starke Basislinien wie ASC und ESC zeigen, dass unsere Methode verbesserte Abwägungen zwischen Antwortkorrektheit, Sampling-Runden und der benötigten Gesamtstichprobenzahl erzielt.

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.