Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

Resumen Técnico: Muestreo Adaptativo Guiado por RL para el Escalamiento en Tiempo de Inferencia

Planteamiento del Problema
El escalamiento en tiempo de inferencia (test-time scaling) ha surgido como una técnica crítica para mejorar las capacidades de razonamiento de los Grandes Modelos de Lenguaje (LLMs) sin necesidad de entrenamiento adicional. Métodos como la Autoconsistencia (SC), el Árbol de Pensamientos (Tree-of-Thought) y el muestreo Best-of-N mejoran la calidad de las respuestas mediante la asignación de más computación durante la inferencia. Sin embargo, estas ganancias incurren en costos sustanciales tanto en computación total como en latencia. Los métodos de muestreo adaptativo existentes, como la Autoconsistencia Adaptativa (ASC) y la Autoconsistencia de Parada Temprana (ESC), intentan mitigar estos costos decidiendo dinámicamente cuándo detener el muestreo. No obstante, estos enfoques sufren limitaciones significativas: a menudo dependen de heurísticas diseñadas por humanos o suposiciones distribucionales, requieren señales auxiliares (por ejemplo, confianza interna del modelo, estados ocultos o dificultad de la pregunta) que pueden no estar disponibles, y pueden ser invasivos para los flujos de inferencia estándar. Existe la necesidad de un método ligero y basado en principios que equilibre la corrección de la respuesta, la latencia y el costo computacional sin requerir características auxiliares ni intervención en el modelo.

Metodología
Los autores proponen el Muestreo Guiado por RL (RL-Guided Sampling), un marco que formula el muestreo adaptativo como un Proceso de Decisión de Markov (MDP) de horizonte finito.

• Formulación del MDP:

  • Estado ($s_t$): El estado se construye únicamente a partir de las estadísticas del conjunto actual de respuestas muestreadas ($D_t$). Incluye los conteos de las $K$ clases de respuestas más frecuentes, el número total de muestras generadas y la entropía de estos conteos. Crucialmente, el estado no depende de información específica de la consulta ni de señales internas del modelo.
  • Acción ($a_t$): En cada ronda, el controlador selecciona una acción de un conjunto discreto: ya sea detener el muestreo (acción 0) o generar un número específico de respuestas candidatas adicionales en paralelo (acciones $k_1, k_2, \dots$).
  • Función de Recompensa: La recompensa equilibra la calidad final de la respuesta frente a los costos de inferencia.
    • Penalización por paso: Se incurre en una penalización por cada ronda de muestreo adicional ($\lambda_{lat}$) y por cada nueva muestra generada ($\lambda_{comp}$).
    • Recompensa Terminal: Al finalizar (ya sea porque el controlador detiene el proceso o alcanza el presupuesto máximo), el modelo recibe una recompensa basada en la corrección de la respuesta de voto mayoritario. La corrección objetivo ($y^\star$) se define como el voto mayoritario que se obtendría si el muestreo continuara hasta el presupuesto máximo $N$, desacoplando la recompensa de las etiquetas de verdad fundamental (ground-truth).
  • Optimización: El controlador se entrena utilizando Optimización de Política Próxima (PPO) para maximizar la recompensa acumulada esperada. Los autores señalan que este objetivo puede interpretarse como la relajación de Lagrange de un problema de optimización restringida donde el objetivo es maximizar la precisión sujeta a presupuestos de latencia y computación.

• Arquitectura y Despliegue:
  El controlador es un Perceptrón de Capas Múltiples (MLP) ligero de cuatro capas. Requiere únicamente las respuestas finales generadas para funcionar, lo que lo hace no invasivo y desplegable en CPUs. No requiere el ajuste fino (fine-tuning) del LLM subyacente.

Contribuciones Clave

1. Formulación del MDP para el Muestreo Adaptativo: El artículo introduce una novedosa formulación del muestreo adaptativo como un MDP, alejándose de las reglas heurísticas hacia una política aprendida que navega explícitamente el compromiso entre rendimiento y costo.
2. Diseño Ligero y No Invasivo: El método se basa puramente en las estadísticas del conjunto de respuestas muestreadas, requiriendo ninguna señal auxiliar (como puntuaciones de confianza) y ninguna modificación en el proceso de razonamiento del LLM. El controlador es lo suficientemente pequeño como para ser entrenado y desplegado en CPUs.
3. Interpretación Teórica: Los autores proporcionan una conexión teórica entre su objetivo de RL y la optimización restringida, mostrando que la función de recompensa ponderada actúa como una relajación de Lagrange para maximizar la precisión bajo estrictas restricciones de recursos.
4. Generalización: El marco demuestra una fuerte transferibilidad, donde las políticas entrenadas en un conjunto de datos o modelo (por ejemplo, Qwen-3-0.6B) pueden aplicarse a diferentes benchmarks e incluso a modelos diferentes y más capaces (por ejemplo, GPT-4.1-nano) con una degradación mínima del rendimiento.

Resultos Experimentales
El método fue evaluado en tres benchmarks de razonamiento matemático (AIME24, AIME25, HMT 2025) utilizando varios muestreadores de LLM (familia Qwen-3 y GPT-4.1-nano).

• Rendimiento vs. Baselines: Comparado con baselines fuertes como ASC y ESC, el Muestreo Guiado por RL logra mejores compromisos:
  • vs. ASC: Reduce las rondas de muestreo aproximadamente 3–4$\times$ y las muestras totales en un ~30% manteniendo una precisión comparable. Esto aborda el cuello de botella de alta latencia del muestreo secuencial de ASC.
  • vs. ESC: Reduce los requisitos de muestras totales en un ~33% y las rondas de muestreo en un ~10%, logrando una mayor precisión al evitar una parada temprana excesivamente agresiva.
  • vs. SC: Aunque todos los métodos adaptativos reducen el conteo total de muestras en comparación con SC de presupuesto fijo, el Muestreo Guiado por RL logra esto con significativamente menos rondas que ASC y una mejor eficiencia que ESC.
• Comportamiento de Escalamiento: Las curvas de escalamiento indican que el Muestreo Guiado por RL supera consistentemente a los baselines en varios ajustes de parámetros, ofreciendo mejor precisión para el mismo número o menos muestras y rondas.
• Asignación de Recursos: El análisis muestra que la política aprendida asigna intuitivamente más muestras a las consultas con mayor entropía de respuesta (diversidad), alineándose con la necesidad de más muestras para converger en un voto mayoritario fiable.
• Estudios de Ablación: Se encontró que el objetivo de "Mayoría Corriente" (voto mayoritario al presupuesto máximo $N$) es superior al uso de etiquetas de verdad fundamental (Real Label), el cual degradó el rendimiento al introducir ruido y prevenir el aprendizaje de criterios de parada universales.

Significancia y Reivindicaciones
El artículo afirma que el Muestreo Guiado por RL ofrece una alternativa fundamentada y ligera a los métodos de muestreo adaptativo existentes. Al enmarcar el problema como un MDP y entrenar un controlador pequeño mediante RL, el método equilibra eficazmente la corrección de la respuesta, la latencia y el costo de computación sin requerir modificaciones invasivas en el LLM o acceso a los estados internos del modelo. Los autores enfatizan que el enfoque es particularmente valioso para el despliegue práctico, ya que puede entrenarse en modelos de código abierto económicos y transferirse para reducir los costos de inferencia de sistemas más fuertes, propietarios o de código cerrado. El trabajo representa un paso inicial hacia el muestreo adaptativo basado en RL, sugiriendo direcciones futuras para incorporar señales de estado más ricas y métricas de costo del mundo real.