Resumen Técnico: T2MLR: Transformer con Recurrencia de Capas Medias Temporales

Declaración del Problema

Los autores identifican una limitación fundamental en las arquitecturas Transformer estándar con respecto a las capacidades de razonamiento. Si bien los Transformers sobresalen en la recuperación en contexto (in-context retrieval), su mecanismo de decodificación autorregresiva crea un cuello de botella de información. En cada paso de decodificación, las representaciones latentes ricas y de alta dimensión se comprimen repetidamente en un espacio de tokens disperso (vectores one-hot) antes de ser introducidas como entrada al siguiente paso. Este proceso impide que los estados de razonamiento abstracto intermedio persistan a través del tiempo, obligando al modelo a reconstruir la información necesaria de forma indirecta a través de la atención o mediante embeddings de entrada altamente comprimidos.

Los enfoques existentes para el razonamiento latente intentan mitigar esto realizando el razonamiento enteramente en un espacio continuo (a menudo en la capa final) o propagando la incertidumbre a través de combinaciones lineales de los embeddings de los tokens. Sin embargo, estos métodos suelen alimentar las señales recurrentes a través del embedding de entrada, omitiendo las capas medias donde se sabe que ocurre el razonamiento abstracto. Alternativamente, las arquitecturas Transformer de tipo "bucle" (looped) aumentan la profundidad efectiva iterando sobre todas las capas, pero esto incurre en costos de inferencia que escalan linealmente con el número de bucles, lo que las hace poco prácticas para el despliegue.

Metodología

El artículo introduce T2MLR (Transformers con Recurrencia de Capas Medias Temporales), una arquitectura novedosa diseñada para permitir que las representaciones abstractas intermedias persistan y evolucionen a través de los pasos de decodificación sin aumentar la profundidad en el tiempo de inferencia.

Diseño Arquitectónico

T2MLR modifica un decodificador Transformer estándar de $L$ capas introduciendo una vía recurrente con compuerta (gated) que conecta una capa profunda del token anterior con una capa temprana del token actual.

1. Límites de Recurrencia: La arquitectura está parametrizada por dos índices de capa, $\ell_{start}$ y $\ell_{end}$, que definen un bloque "medio" de capas.
2. Caché Recurrente ($R_t$): Se mantiene un vector de caché de tamaño constante para cada paso de decodificación $t$. Este almacena la representación después de la capa $\ell_{end}$ del paso anterior.
3. Módulo de Fusión ($\Phi$): Antes de procesar el token actual en la capa $\ell_{start}$, el modelo fusiona la representación del token actual ($h^{(\ell_{start}-1)}_t$) con la caché recurrente ($R_{t-1}$) del token anterior. La fusión se rige por un mecanismo de compuerta aprendible:
   $$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$$
   Este diseño permite que el modelo aprenda cuándo depender del flujo del token actual frente al estado latente acumulado temporalmente.
4. Actualización de la Caché: Después de pasar por la capa $\ell_{end}$, la caché recurrente se actualiza utilizando la nueva representación y la caché previa mediante una conexión residual y RMSNorm.

Estrategia de Entrenamiento: Paralelismo Temporal

Un desafío clave para las arquitecturas de razonamiento latente es que el entrenamiento estándar de paralelismo de secuencia (paralelismo de tokens) se rompe cuando existen dependencias a través de los pasos temporales. Para abordar esto, los autores proponen un esquema de entrenamiento de paralelismo temporal aproximado utilizando iteraciones de punto fijo de Jacobi:

• En lugar de un procesamiento secuencial de tokens, el modelo realiza $d_{forward}$ iteraciones de refinamiento en paralelo.
• Comienza con una caché inicial derivada de un paso hacia adelante (forward pass) estándar, luego fusiona iterativamente esta caché con la representación del token actual, refina el bloque medio y actualiza la caché.
• Los gradientes se retropropagan a través de estas iteraciones con una profundidad truncada ($d_{backward}$), de forma análoga a la Retropropagación a Través del Tiempo Truncada (TBPTT).
• Este enfoque permite entrenar el modelo con teacher forcing denso y paralelismo de secuencia, una capacidad que los autores afirman es única entre las variantes de cadena de pensamiento (chain-of-thought) continuas.

Contribuciones Clave

1. Innovación Arquitectónica: T2MLR introduce una vía recurrente dirigida específicamente en las capas medias del Transformer, en lugar de en el embedding de entrada o la salida final. Esto preserva la interfaz autorregresiva estándar mientras permite la persistencia del estado latente.
2. Entrenamiento Escalable: Los autores demuestran un método para entrenar estas dependencias recurrentes utilizando paralelismo temporal aproximado, permitiendo el preentrenamiento y el ajuste fino en grandes conjuntos de datos sin sacrificar la escalabilidad.
3. Capacidad de Retrofit (Reajuste): El artículo muestra que T2MLR no requiere un entrenamiento desde cero. La vía recurrente puede insertarse en un modelo preentrenado existente (por ejemplo, SmolLM2-1.7B) y ajustarse para obtener mejoras significativas en el razonamiento.
4. Eficiencia: A diferencia de los Transformers con bucles, T2MLR añade un overhead de inferencia insignificante (a lo sumo un ~8% por token) porque la recurrencia se gestiona mediante una caché de tamaño constante y un único paso hacia adelante durante la inferencia.

Resultados Experimentales

Los autores evalúan T2MLR a través de benchmarks sintéticos, preentrenamiento en lenguaje natural y tareas de razonamiento descendentes.

• Razonamiento Sintético (S5-Retrieval): T2MLR supera significativamente tanto a los Transformers estándar como a los LSTMs en una tarea que requiere seguimiento de estado simultáneo y recuperación en contexto. Mientras que los Transformers fallan en secuencias largas debido a las restricciones de profundidad y los LSTMs fallan en la recuperación debido a la compresión de estado fija, T2MLR mantiene una alta precisión.
• Preentrenamiento (FineWeb-Edu): En modelos de 135M de parámetros, las variantes de T2MLR igualan o superan consistentemente a los baselines de Transformer con la misma cantidad de parámetros en benchmarks de NLP zero-shot (ARC, HellaSwag, etc.). Notablemente, la recurrencia de capa media (iterando sobre ~20% de las capas, por ejemplo, capas 13–18) suele superar a la recurrencia de capa completa.
• Razonamiento Descendente: En tareas como GSM-Aug (matemáticas), HotPotQA (QA multi-salto) y Variable Assignment, T2MLM muestra ganancias consistentes sobre los baselines. Las variantes de recurrencia de capa media superan nuevamente a las de recurrencia de capa completa, sugiriendo que la recurrencia dirigida es más efectiva que el bucle global.
• Escalabilidad: Las ganancias de rendimiento persisten y crecen al escalar a 361M y 1B de parámetros y extender el preentrenamiento a 50B de tokens.
• Retrofit: Aplicar T2MLR a un modelo SmolLM2-1.7B-Instruct preentrenado y realizar un ajuste fino en datos matemáticos mejoró la precisión de GSM8K de 35.8% a 39.9% y de MATH500 de 12.8% a 18.0%, demostando su aplicabilidad práctica sin necesidad de un preentrenamiento completo.

Significado y Reivindicaciones

El artículo postula que el razonamiento latente efectivo en Transformers no requiere iterar sobre todas las capas (lo que aumenta el costo de inferencia) ni confinar la recurrencia al embedding de entrada (lo que evita el cómputo abstracto). En su lugar, los autores afirman que la recurrencia dirigida de capa media es suficiente para desbloquear capacidades de razonamiento significativas.

El significado de T2MLR reside en su capacidad para:

1. Desacoplar la Profundidad de Razonamiento del Costo de Inferencia: Introduce una vía latente temporal que persiste a través de los pasos de tiempo sin la penalización lineal de inferencia asociada con las arquitecturas de bucle.
2. Reducir la Barrera de Adopción: Al demostrar que la arquitectura puede adaptarse (retrofit) a modelos preentrenados actuales, los autores argumentan que se pueden añadir capacidades de razonamiento avanzadas a los modelos actuales a una fracción del costo de entrenar desde cero.
3. Validar Sesgos Inductivos: Los resultados sugieren que el "punto óptimo" para la recurrencia en Transformers son las capas medias, donde ocurre el cómputo abstracto, en lugar de las interfaces de entrada o salida.

Los autores reconocen que T2MLR incurre en un overhead de tiempo de entrenamiento (2–4x el tiempo de reloj dependiendo del tamaño del bloque recurrente) debido a los pasos de aproximación de Jacobi. Lo plantean como un compromiso: T2MLR sacrifica cómputo de entrenamiento para ganar una arquitectura de inferencia de bajo costo y alto rendimiento para tareas de razonamiento.