Petit contrôleur RL, grand modèle de langage : échantillonnage adaptatif guidé par RL pour la mise à l'échelle au moment du test

Résumé

La mise à l'échelle au moment du test améliore la performance de raisonnement des grands modèles de langage, mais engendre un coût substantiel en termes de calcul total et de latence. Les méthodes existantes d'échantillonnage adaptatif atténuent partiellement ce problème en décidant dynamiquement quand arrêter l'échantillonnage, mais elles reposent généralement sur des règles heuristiques ou des hypothèses de distribution. Dans ce travail, nous formulons l'échantillonnage adaptatif comme un processus de décision markovien (PDM). Nous entraînons un contrôleur d'échantillonnage léger par apprentissage par renforcement (RL) pour équilibrer conjointement l'exactitude des réponses, la latence et le coût de calcul. À chaque tour, le contrôleur décide d'arrêter l'échantillonnage ou d'acquérir des échantillons supplémentaires. Notre méthode est légère : elle ne repose que sur les statistiques des réponses finales et peut être entraînée et déployée sur CPU. Nous montrons en outre que le cadre résultant admet une interprétation comme la relaxation lagrangienne d'un problème d'optimisation sous contraintes avec des contraintes budgétaires explicites. Des expériences menées face à des références solides telles que ASC et ESC montrent que notre méthode atteint des compromis améliorés entre l'exactitude des réponses, le nombre de tours d'échantillonnage et le nombre total d'échantillons requis.

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.