GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

Aquí presento un resumen técnico detallado del artículo "GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler", estructurado según los puntos solicitados.

1. El Problema: Limitaciones de la Escalabilidad en Tiempo de Inferencia (ITS)

El artículo aborda un desafío crítico en los modelos de razonamiento latente (donde el razonamiento ocurre en espacios de estados continuos en lugar de generar tokens de texto intermedios).

• Contexto: La Escalabilidad en Tiempo de Inferencia (ITS) mejora el rendimiento de los LLMs generando múltiples trayectorias de razonamiento y seleccionando la mejor. En modelos discretos (basados en tokens), esto se logra mediante distribuciones de probabilidad explícitas (ej. temperature scaling).
• Limitación Actual: En modelos de razonamiento latente, la ausencia de distribuciones de probabilidad explícitas obliga a los métodos existentes a depender de perturbaciones heurísticas (como dropout o ruido gaussiano fijo) para introducir estocasticidad y explorar el espacio de razonamiento.
• Hallazgo Clave: Los autores demuestran que aumentar la magnitud de la perturbación no garantiza una mejor calidad de muestreo.
  • Las perturbaciones fuertes a menudo inducen grandes desplazamientos de distribución (distribution shifts) sin mejorar la probabilidad de encontrar la respuesta correcta.
  • Esto lleva a dos modos de fallo: sub-exploración (perturbaciones demasiado débiles que no se alejan de la trayectoria determinista) y sobre-exploración (perturbaciones fuertes que dispersan las trayectorias hacia regiones no útiles, degradando el rendimiento).
  • La estocasticidad actual carece de una distribución de muestreo condicional explícita, lo que hace que la exploración sea difícil de controlar u optimizar.

2. Metodología: Gaussian Thought Sampler (GTS)

Para resolver esto, los autores proponen GTS, un módulo ligero que reformula la exploración latente como un problema de muestreo condicional explícito.

• Concepto Central: En lugar de inyectar ruido fijo, GTS aprende una distribución gaussiana condicional sobre las perturbaciones de los estados latentes.
• Arquitectura:
  • Se añade un módulo pequeño (MLP) a un modelo base de razonamiento latente congelado (frozen backbone).
  • En cada paso de razonamiento $k$, dado el estado latente determinista $h_k^{det}$, GTS predice los parámetros de una distribución gaussiana: media $\mu_\phi(c_k)$ y desviación estándar $\sigma_\phi(c_k)$, donde $c_k$ es el contexto (el estado oculto).
  • La perturbación $z_k$ se muestrea como: $z_k = \mu_\phi(c_k) + \sigma_\phi(c_k) \odot \epsilon_k$, donde $\epsilon_k \sim \mathcal{N}(0, I)$.
  • El nuevo estado es $\tilde{h}_k = h_k^{det} + z_k$.
• Optimización (Aprendizaje):
  • Dado que la calidad de la trayectoria depende de la respuesta final (una señal no diferenciable respecto a los parámetros del muestreador), GTS se entrena mediante Aprendizaje por Refuerzo (RL).
  • Se utiliza una variante de GRPO (Group Relative Policy Optimization).
  • Recompensa: Se define una recompensa basada en la corrección de la respuesta y un término de "moldeado" (shaping) basado en la confianza (probabilidad log-normalizada). Esto fomenta trayectorias correctas y de alta confianza, desalentando las incorrectas de alta confianza.
  • Se mantiene un muestreador de referencia (reference sampler) para regularizar el KL y evitar el colapso de la política.

3. Contribuciones Clave

1. Identificación de una limitación fundamental: Se demuestra que la escalabilidad efectiva en razonamiento latente requiere controlar la calidad del muestreo, no simplemente inyectar más estocasticidad. Las perturbaciones heurísticas fijas son sensibles al modelo y la tarea, llevando a sub o sobre-exploración.
2. Propuesta de GTS: Un muestreador gaussiano ligero y aprendible que transforma la perturbación heurística en una política de muestreo probabilística explícita y optimizable.
3. Validación Empírica: Demostración de que un enfoque de muestreo basado en principios (aprendido) supera consistentemente a las perturbaciones basadas en dropout y ruido gaussiano estándar en múltiples arquitecturas y benchmarks.

4. Resultados Experimentales

Los experimentos se realizaron en dos arquitecturas de razonamiento latente (COCONUT y CODI) utilizando el conjunto de datos GSM8K y otros benchmarks fuera de distribución (MultiArith, SVAMP, GSM8K-Hard).

• Rendimiento de Escalado (Pass@N):
  • GTS muestra un comportamiento de escalado superior a medida que aumenta el presupuesto de muestreo ($N$).
  • En COCONUT, GTS alcanza un pass@N significativamente mayor que los baselines (ej. ~55% vs ~35-40% en presupuestos grandes).
  • En CODI, aunque las ganancias son moderadas debido a la mayor complejidad del espacio latente, GTS sigue siendo el método más robusto y consistente.
• Calidad de Muestreo (Análisis Diagnóstico):
  • Se introdujeron métricas como Sampling Gain (SG) (mejora en la log-odds de la respuesta correcta) y Jensen-Shannon Divergence (JS) (desplazamiento de distribución).
  • Hallazgo: Las perturbaciones fuertes (ej. dropout $p=0.5$) aumentan drásticamente el JS (desplazamiento) pero reducen el SG (calidad).
  • GTS logra un SG positivo alto manteniendo un desplazamiento de distribución controlado, evitando los extremos de sub/sobre-exploración.
• Generalización: GTS mantiene su superioridad en benchmarks fuera de distribución, mientras que los métodos heurísticos muestran inconsistencia (funcionan bien en un dataset y mal en otro).
• Análisis de SNR (Relación Señal-Ruido): GTS adapta dinámicamente el equilibrio entre la dirección determinista y la exploración estocástica en cada paso de razonamiento, alineándose con la dinámica interna del modelo base (ej. patrones de "sierra" en COCONUT vs. ajuste frontal en CODI).

5. Significado e Impacto

• Cambio de Paradigma: El trabajo sugiere que la exploración en espacios latentes continuos no debe tratarse como un proceso estocástico "a ciegas" (ruido fijo), sino como un problema de política de muestreo aprendible.
• Eficiencia Computacional: GTS es un módulo ligero (añade solo ~1-2% de parámetros al modelo base) que no requiere reentrenar el modelo principal, solo optimizar el muestreador.
• Control de Exploración: Proporciona un mecanismo para controlar explícitamente la diversidad de las trayectorias de razonamiento, asegurando que la exploración se dirija hacia regiones del espacio latente que apoyan la toma de decisiones correcta.
• Futuro: Establece una base para futuras investigaciones en la escalabilidad en tiempo de inferencia para modelos de razonamiento continuo, alejándose de las heurísticas manuales hacia políticas de muestreo optimizadas por recompensa.

En resumen, el paper demuestra que la escalabilidad efectiva en razonamiento latente depende de un muestreo mejor controlado y optimizable, y no simplemente de la amplificación de la estocasticidad, proponiendo GTS como la solución práctica para lograrlo.