ReLaX: Reasoning with Latent Exploration for Large Reasoning Models

Resumen Técnico: ReLaX

1. El Problema: Colapso de Entropía y Exploración Ineficiente

El artículo aborda una limitación crítica en el entrenamiento de Modelos de Razonamiento Grandes (LRMs) mediante Aprendizaje por Refuerzo con Recompensas Verificables (RLVR). Aunque el RLVR ha demostrado ser efectivo para mejorar el razonamiento, tiende a empujar a la política del modelo hacia un determinismo excesivo.

• Colapso de Entropía: Sin intervención explícita, el RL reduce progresivamente la distribución de la política y la entropía, confinando el gradiente de la política en un subespacio estrecho. Esto provoca una convergencia prematura y una exploración ineficiente.
• Limitación de los Métodos Actuales: Las soluciones existentes se centran en aumentar la entropía a nivel de tokens (la probabilidad de generar palabras específicas). Sin embargo, los autores argumentan que esto es un síntoma superficial. El problema real es la pérdida de flexibilidad en las dinámicas latentes internas (los estados ocultos del modelo) que gobiernan la generación de tokens.
• Desafío Multimodal: En Modelos de Lenguaje Multimodal (MLLMs), existe una desalineación entre los cálculos internos cruzados (visión-idioma) y las salidas unimodales (texto), lo que hace que los feedbacks a nivel de token no reflejen con precisión la complejidad del procesamiento subyacente.

2. Metodología: ReLaX y Dinámicas Latentes

La propuesta central es ReLaX (Reasoning with Latent eXploration), un marco que regula la exploración y explotación operando directamente sobre las dinámicas latentes del modelo, en lugar de solo sobre los tokens de salida.

A. Teoría del Operador de Koopman y DSD
Para hacer analíticamente tratables las dinámicas no lineales y de alta dimensión de los LRMs, los autores utilizan la Teoría del Operador de Koopman:

• Linealización: Transforman las dinámicas no lineales de los estados ocultos en una evolución lineal en un espacio de funciones de dimensión infinita.
• Diccionario de Koopman: Utilizan una red neuronal (ResKoopNet) para aprender un diccionario de observables que mapea los estados ocultos a este espacio lineal.
• Dispersión Espectral Dinámica (DSD): Introducen una nueva métrica, el DSD, definida como la varianza de las magnitudes de los autovalores del operador de Koopman aproximado.
  • Un DSD alto indica un espectro dinámico rico y heterogéneo (alta flexibilidad computacional).
  • Un DSD bajo sugiere un comportamiento rígido, repetitivo y degenerado.

B. Funcionamiento de ReLaX
ReLaX integra el DSD en el objetivo de optimización de la política (basado en GRPO - Group Relative Policy Optimization):

1. Regularización basada en DSD: Se añade un término de pérdida ($L_{xp}$) que penaliza la convergencia hacia dinámicas latentes homogéneas.
2. Forma de la Ventaja (Advantage Shaping): Para evitar una exploración desinformada, la regularización se pondera solo en las trayectorias que obtienen recompensas positivas (ventaja truncada). Esto asegura que la flexibilidad se fomente solo en caminos útiles.
3. Regularización KL Adaptativa: Se aplica una penalización KL selectiva solo a las trayectorias con una dispersión espectral excesiva (demasiado ruido), actuando como un mecanismo de estabilización sin frenar la exploración necesaria.

La función objetivo final combina la ganancia de GRPO, la regularización de exploración latente y la penalización KL adaptativa.

3. Contribuciones Clave

1. Métrica DSD: Propuesta de una métrica fundamental que captura la heterogeneidad de las dinámicas de razonamiento latente, ofreciendo una caracterización más profunda de la exploración que las estadísticas superficiales de los tokens.
2. Marco ReLaX: Desarrollo de un nuevo paradigma de entrenamiento que utiliza la dinámica latente para regular el equilibrio exploración-explotación, mitigando la saturación de rendimiento típica del RLVR.
3. Superioridad Empírica: Demostración de que intervenir en el espacio latente es más efectivo que en el espacio de tokens, logrando nuevos estados del arte (SOTA) en múltiples benchmarks.

4. Resultados Experimentales

Los autores evaluaron ReLaX en modelos de 3B y 7B parámetros, tanto en tareas de texto puro (matemáticas) como multimodales (visión-idioma).

• Razonamiento Multimodal (VLMs):

  • En 7 benchmarks multimodales (incluyendo MMMU, MathVista, DynaMath), ReLaX-VL-7B alcanzó un promedio de 53.2, superando al anterior líder (VL-Rethinker-7B con 52.5) y estableciendo un nuevo SOTA.
  • ReLaX-VL-3B superó a varios modelos de 7B anteriores, demostrando una alta eficiencia.
  • Las dinámicas de entrenamiento mostraron que ReLaX mantiene una entropía y un DSD más altos y estables, evitando el colapso rápido observado en GRPO estándar.

• Razonamiento Matemático (LLMs):

  • En benchmarks como MATH500, Minerva, AMC y AIME, ReLaX superó consistentemente a variantes de GRPO y otros algoritmos de RLVR (como DAPO, KL-Cov, FR3E).
  • En el modelo Qwen2.5-7B-Math, ReLaX superó a FR3E (el SOTA anterior) en 6.3 puntos de promedio.

• Análisis Cualitativo:

  • Los casos de estudio revelaron que los modelos entrenados con métodos de entropía de tokens (como R1-zero-Div) a menudo generan verificaciones ilusorias (ej. escribir código Python sin entorno de ejecución).
  • En contraste, ReLaX fomentó una verificación semánticamente significativa (ej. aplicar correctamente la ley de cosenos), demostrando una comprensión más robusta y flexible del razonamiento.

5. Significado e Impacto

El trabajo de ReLaX representa un cambio de paradigma en la mejora de modelos de razonamiento:

• Más allá de los Tokens: Demuestra que la exploración efectiva no debe limitarse a la diversificación de la salida de tokens, sino que debe dirigirse a la flexibilidad de los procesos computacionales internos (espacio latente).
• Solución al Dilema Exploración-Explotación: Proporciona un mecanismo principiado para mantener la diversidad necesaria para el descubrimiento de soluciones complejas sin sacrificar la estabilidad del entrenamiento.
• Escalabilidad: Al operar en el espacio latente, el método es más generalizable a modelos multimodales donde la relación entre entrada visual y salida textual es compleja, ofreciendo una ruta prometedora para avanzar hacia una Inteligencia Artificial General (AGI) con capacidades de razonamiento robustas.

En resumen, ReLaX utiliza la teoría de sistemas dinámicos (Operador de Koopman) para "desbloquear" el potencial de exploración latente de los modelos, resolviendo el problema del colapso de entropía y logrando un rendimiento superior en tareas de razonamiento complejas.