Pequeño controlador RL, gran modelo de lenguaje: muestreo adaptativo guiado por RL para escalado en tiempo de prueba

Resumen

El escalado en tiempo de prueba mejora el rendimiento de razonamiento de los modelos de lenguaje grandes, pero incurre en un costo sustancial tanto en cómputo total como en latencia. Los métodos existentes de muestreo adaptativo mitigan parcialmente este problema al decidir dinámicamente cuándo dejar de muestrear, aunque típicamente se basan en reglas heurísticas o en supuestos sobre la distribución. En este trabajo, formulamos el muestreo adaptativo como un proceso de decisión de Markov (MDP). Entrenamos un controlador de muestreo ligero mediante aprendizaje por refuerzo (RL) para equilibrar conjuntamente la corrección de la respuesta, la latencia y el costo computacional. En cada ronda, el controlador decide si detener el muestreo o adquirir muestras adicionales. Nuestro método es ligero y solo depende de estadísticas de las respuestas finales, pudiendo ser entrenado e implementado en CPU. Además, mostramos que el marco resultante admite una interpretación como la relajación lagrangiana de un problema de optimización con restricciones explícitas de presupuesto. Los experimentos frente a líneas base sólidas como ASC y ESC demuestran que nuestro método logra mejores compensaciones entre la corrección de la respuesta, las rondas de muestreo y el total de muestras requeridas.

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.