Latent Reasoning in TRMs is Secretly a Policy Improvement Operator

Resumen Técnico: El razonamiento latente en los TRM es secretamente un operador de mejora de políticas

1. Planteamiento del problema

Modelos pequeños recientes que utilizan recursión latente, como los Modelos Recursivos Diminutos (TRMs) y los Modelos de Razonamiento Jerárquico (HRMs), han demostrado resultados prometedores en tareas de razonamiento complejo como ARC-AGI. La teoría prevaleciente sugiere que estos modelos logran mejoras de rendimiento al aumentar la profundidad efectiva de la red, permitiéndoles emular modelos más grandes con menos parámetros. Sin embargo, las observaciones empíricas revelan una discrepancia: las capas añadidas recursivamente suelen tener un rendimiento inferior en comparación con los modelos de un solo paso de profundidad feed-forward equivalente. Esto indica que en las arquitecturas de bucle, no cada paso recursivo contribuye significativamente a la solución, lo que resulta en "cómputo muerto".

Las preguntas fundamentales abordadas en este trabajo son:

• ¿Cuándo y por qué el razonamiento latente mejora el rendimiento?
• ¿Cuándo resulta en un cómputo redundante?
• ¿Cómo se puede mejorar el entrenamiento de estos modelos para asegurar que cada paso sea efectivo?

Los métodos de entrenamiento actuales para los TRMs aplican la supervisión principalmente al resultado final. Esto deja los pasos de refinamiento tempranos con una asignación de crédito débil, lo que hace que la optimización sea frágil, particularmente en regímenes de pocos datos.

2. Metodología

Los autores proponen un reencuadre teórico del razonamiento latente y una metodología de entrenamiento correspondiente.

Formalización Teórica: El razonamiento latente como mejora de política
El artículo formaliza el razonamiento recursivo latente como un operador de mejora de política.

• Política de Referencia ($\hat{\pi}_t$): La distribución de salida del modelo antes de un paso de razonamiento latente (el "primer intento").
• Política Mejorada ($\pi^+_t$): La distribución de salida después de una actualización de razonamiento latente.
• Señal de Ventaja: La diferencia entre los log-probabilidades de la política mejorada y la política de referencia actúa como una señal de ventaja implícita ($A_t(s, a) = \log \pi^+_t(a|s) - \log \hat{\pi}_t(a|s)$).
• Mecanismo: Los autores demuestran que un solo paso de TRM implementa una actualización de política que repondera la política de referencia basándose en una condición de optimalidad. Específicamente, el paso de razonamiento latente se comporta como una política condicionada al evento de que la respuesta se mejore. Esto conecta el algoritmo del espacio latente con el marco bien conocido de la mejora de políticas en Aprendizaje por Refuerzo (RL) y la Guía Libre de Clasificador (CFG) en modelos de difusión.

Esquema de Entrenamiento Propuesto: Supervisión de Mejora Profunda (DIS)
Basándose en la idea de la mejora de política, los autores introducen la Supervisión de Mejora Profunda (DIS) para eliminar los pasos de cómputo muerto.

• Objetivos paso a paso: En lugar de supervisar solo la salida final, DIS proporciona objetivos intermedios $\{y^\dagger_s\}$ para cada paso de recursión. Estos objetivos se generan mediante un esquema de corrupción discreta monotónica de la salida real $y^*$.
• Objetivo de Entrenamiento: El modelo se entrena para mapear el objetivo anterior $y^\dagger_{s-1}$ (a través de la política de referencia) hacia el siguiente objetivo menos corrupto $y^\dagger_s$ (a través de la política mejorada).
• Función de Pérdida: El entrenamiento minimiza una pérdida de entropía cruzada entre los logits post-razonamiento y el objetivo específico del paso $y^\dagger_s$. Esto obliga a que la ventaja de la razón de logaritmos en cada paso sea positiva para el token correcto, asegurando que cada paso de recursión realice una mejora dirigida en lugar de un cómputo redundante.

Simplificaciones Arquitectónicas
El método propuesto simplifica la arquitectura TRM al:

• Eliminar la necesidad de un mecanismo de parada aprendido (Tiempo de Computación Adaptativa).
• Reducir el número de pasos de supervisión y de pasos de razonamiento latente.
• Utilizar un número fijo de pasos de recursión ($N_{sup}$) determinado por el esquema de corrupción.

3. Contribuciones Clave

1. Perspectiva Teórica: El artículo establece que el razonamiento latente en los TRM es matemáticamente equivalente a un operador de mejora de política. Establece un puente entre los algoritmos de razonamiento en el espacio latente y los marcos de RL, explicando el comportamiento del modelo a través de la lente de la maximización de la ventaja.
2. Supervisión de Mejora Profunda (DIS): Un esquema de entrenamiento novedoso que suministra objetivos intermedios paso a paso. Esto transforma el difícil problema de la asignación de crédito a largo plazo en una tarea de aprendizaje supervisado tratable donde cada paso se entrena explícitamente para mejorar sobre el estado anterior.
3. Ganancias de Eficiencia: El método logra un rendimiento competitivo mientras reduce significativamente la carga computacional. Evita el entrenamiento de módulos de parada y reduce el número total de pasadas hacia adelante en un factor de 18× en comparación con las configuraciones estándar de TRM, manteniendo o mejorando la precisión.

4. Resultados Experimentales

Los autores evaluaron su método en el problema de las N-Reinas y en los benchmarks ARC-AGI (ARC-AGI-1 y ARC-AGI-2).

• N-Queens: Utilizando un modelo compacto (0.8M de parámetros), el método DIS logró una precisión de 0.69, igualando el rendimiento de un TRM estándar con una cabeza de parada y 16 pasos de supervisión, pero utilizando significativamente menos pasos de inferencia ($T=1, n=2$ vs. $T=3, n=6$).
• ARC-AGI:
  • Modelo Compacto (0.8M de parámetros): DIS logró un 24.0% de precisión en ARC-AGI-1, superando al TRM-compact original (12.0%) y a la mayoría de los LLMs de código abierto existentes sin conocimiento externo.
  • Modelo Medio (7M de parámetros): DIS logró un 41.3% en ARC-AGI-1 y un 6.0% en ARC-AGI-2, superando al TRM-medium estándar (40.4% / 3.3%) e igualando el rendimiento del TRM original (40.4% / 3.3%) mientras utiliza menos pasos de razonamiento latente.
• Reducción de Cómputo: El método redujo el número total de pasadas hacia adelante en un factor de 18 en comparación con la configuración base de TRM, demostrando que la supervisión explícita de los pasos intermedios elimina la necesidad de excesivos bucles de razonamiento latente.

5. Significado y Reivindicaciones

El artículo afirma que, al ver el razonamiento latente como un proceso de mejora de política, uno puede explicar por qué los modelos recursivos estándar a veces fallan (debido al "cómputo muerto" donde los pasos no mejoran el margen de ventaja) y cómo solucionarlo.

La importancia principal reside en el marco de Supervisión de Mejora Profunda, el cual:

• Proporciona una regla de entrenamiento constructiva para asegurar que cada paso de recursión realice una mejora no trivial.
• Simplifica el entrenamiento de modelos recursivos al eliminar la necesidad de mecanismos de parada complejos y de una asignación de crédito frágil.
• Demuestra que los modelos pequeños y eficientes en parámetros pueden lograr resultados de vanguardia en tareas de razonamiento complejo (como ARC-AGI) cuando el objetivo de entrenamiento está alineado con el mecanismo teórico de la mejora de política.

Los autores concluyen que su enfoque ofrece un camino hacia modelos de razonamiento más eficientes y efectivos al alinear la señal de entrenamiento con la naturaleza algorítmica intrínseca de la recursión latente.