Resumen Técnico: Los Transformadores de Predicción de Latente Siguiente Aprenden Modelos del Mundo Compactos

Enunciado del Problema

Las arquitecturas estándar de Transformadores reemplazan la recurrencia con la autoatención y un mecanismo de memoria que escala con la longitud de la secuencia. Si bien esto permite una paralelización y potencia predictiva excepcionales, elimina el sesgo inductivo inherente para comprimir el historial en estados latentes compactos con reglas de transición consistentes. En consecuencia, los Transformadores a menudo aprenden atajos complejos específicos de la tarea que se ajustan bien a los datos de entrenamiento pero generalizan mal, fallando en formar "modelos del mundo" coherentes o estados de creencia necesarios para la planificación y el razonamiento a largo plazo. El objetivo estándar de predicción del siguiente token no garantiza que los estados ocultos internos del modelo evolucionen como estadísticas suficientes para el historial pasado, ni asegura que estos estados puedan predecir recursivamente las observaciones futuras.

Metodología: Predicción de Latente Siguiente (NextLat)

Los autores introducen Predicción de Latente Siguiente (NextLat), un marco que complementa el entrenamiento estándar de siguiente token con predicciones auto-supervisadas en el espacio latente. La idea central es entrenar al Transformador para aprender representaciones latentes (específicamente, los estados ocultos de la capa final $h_t$) que sean predictivos del siguiente estado latente dado el siguiente token.

Fundamento Teórico

El artículo establece que si un modelo optimiza tanto la Consistencia del Siguiente Token (predecir $X_{t+1}$ a partir de $h_t$) como la Consistencia de Transición (predecir $h_{t+1}$ a partir de $h_t$ y $X_{t+1}$), los estados ocultos $h_t$ están teóricamente garantizados a converger hacia estados de creencia. Un estado de creencia se define como una estadística suficiente del historial $X_{1:t}$ para predecir el futuro $X_{t+1:T}$. A diferencia del entrenamiento autorregresivo estándar, que permite búsquedas ad-hoc de tokens pasados sin comprimir información, NextLat impone un sesgo inductivo recurrente donde el estado latente debe contener toda la información necesaria para transicionar al siguiente estado y predecir el siguiente token.

Implementación

NextLat opera añadiendo un objetivo auxiliar ligero al bucle de entrenamiento estándar del Transformador:

1. Modelo de Dinámicas Latentes ($p_\psi$): Un Perceptrón Multicapa (MLP) simple se entrena para predecir el siguiente estado oculto $\hat{h}_{t+1}$ dado el estado oculto actual $h_t$ y el siguiente token $X_{t+1}$.
2. Función de Pérdida: La pérdida total combina:
   • Pérdida de Siguiente Token ($L_{next-token}$): Entropía cruzada estándar para la predicción de tokens.
   • Pérdida de Siguiente Estado Oculto ($L_{next-h}$): Una pérdida Smooth L1 que regresa el estado predicho $\hat{h}_{t+1}$ contra el verdadero siguiente estado $h_{t+1}$ (con gradiente detenido aplicado al objetivo para evitar el colapso representacional).
   • Pérdida de Divergencia KL ($L_{KL}$): Un objetivo complementario que asegura que la distribución de tokens predicha a partir del estado latente predicho coincida con la distribución del estado latente verdadero, actuando como una señal de supervisión suave.
3. Supervisión Multi-paso: El modelo puede entrenarse para desarrollar recursivamente las dinámicas latentes durante $d$ pasos, proporcionando señales de gradiente más ricas para dependencias de largo alcance.

Crucialmente, NextLat deja la arquitectura del Transformador, el procedimiento de inferencia y la eficiencia del entrenamiento paralelo sin cambios. El modelo de dinámicas latentes se utiliza solo durante el entrenamiento para dar forma a las representaciones; durante la inferencia, el modelo decodifica de manera autorregresiva como de costumbre.

Contribuciones Clave

1. Garantía Teórica: El artículo demuestra que optimizar para la consistencia del siguiente token y del siguiente latente moldea teóricamente las representaciones del transformador en estados de creencia, una propiedad no garantizada por la predicción de siguiente token sola.
2. Marco Práctico: NextLat introduce un objetivo auxiliar computacionalmente eficiente que inyecta un sesgo inductivo recurrente sin alterar la estructura de entrenamiento paralelo o de inferencia del Transformador.
3. Decodificación Auto-especulativa de Longitud Variable: Al aprender dinámicas latentes coherentes, NextLat permite la predicción recursiva en el espacio latente. Esto permite la decodificación auto-especulativa de longitud variable, donde el modelo puede redactar tokens más allá del horizonte fijo de entrenamiento, acelerando significativamente la inferencia.

Resultados Experimentales

Los autores evalúan NextLat en cuatro dominios: modelado del mundo, razonamiento, planificación y modelado del lenguaje.

• Modelado del Mundo (Viajes de Taxi en Manhattan): NextLat aprende mapas internos significativamente más consistentes con el verdadero modelo del mundo (topología de las calles) en comparación con las líneas base estándar GPT, MTP y JTP. Logra una mayor compresión de secuencia (0.71 frente a 0.65 para GPT) y un rango latente efectivo mucho menor (52.7 frente a 160.1 para GPT), indicando representaciones más compactas y eficientes.
• Razonamiento (Countdown): En la tarea de razonamiento matemático, NextLat supera sustancialmente a las líneas base (54.8% de precisión en $d=1$ frente a 33.1% para GPT). Demuestra una planificación de anticipación superior, reduciendo los "compromisos lamentables" donde los modelos fuerzan ecuaciones finales inválidas para coincidir con un objetivo.
• Planificación (Gráficos Path-Star): NextLat mantiene tasas de resolución cercanas al 100% en diversas topologías de grafos, mientras que líneas base como BST y JTP fallan en grafos más grandes ($G_{7,7}$). Esto sugiere que NextLat evita el aprendizaje de atajos "Clever Hans" común en la predicción en el espacio de tokens.
• Modelado del Lenguaje (FineWeb-Edu):
  • Precisión en Tareas Derivadas: NextLat preserva mejor la perplejidad del siguiente token que MTP y JTP mientras logra ganancias modestas en precisión cero-shot en tareas derivadas (59.21% promedio frente a 58.82% para GPT).
  • Velocidad de Inferencia: Aprovechando la decodificación auto-especulativa de longitud variable, NextLat logra hasta 3.3× más velocidad de inferencia en comparación con el muestreo autorregresivo estándar. Esta aceleración persiste incluso al redactar más allá del horizonte de entrenamiento, una capacidad que MTP y JTP carecen debido a sus horizontes fijos en el espacio de tokens.

Significado y Afirmaciones

El artículo posiciona a NextLat como un paradigma simple pero efectivo para aprender representaciones compactas, predictivas y generalizables en Transformadores. Su principal significado radica en unificar el sesgo inductivo recurrente de los modelos del mundo compactos con la escalabilidad y el paralelismo de los Transformadores.

Los autores afirman que NextLat ofrece una solución a la naturaleza "miópica" de la predicción de siguiente token, fomentando que los modelos formen estados internos estructurados y predictivos en lugar de depender de correlaciones superficiales de tokens. Al permitir la decodificación especulativa de longitud variable sin cambios arquitectónicos, NextLat proporciona un camino práctico hacia una inferencia más rápida y un mejor razonamiento a largo plazo. El artículo concluye que NextLat es un paso hacia objetivos de entrenamiento que dotan a los modelos autorregresivos de representaciones más simples, compactas y, por lo tanto, más generalizables de sistemas complejos.

Los autores permanecen modestos respecto a las direcciones futuras, señalando que no exploraron arquitecturas complejas de dinámicas latentes (utilizando solo MLPs simples) y que el diseño óptimo del objetivo (por ejemplo, la necesidad de la pérdida KL a escalas mayores) requiere un estudio sistemático adicional. También reconocen que, aunque NextLat mejora las representaciones, la complejidad de circuito subyacente del Transformador de profundidad fija sigue estando acotada por la clase TC0, aunque el modelo de dinámicas latentes co-entrenado (RNN) parece capaz de generalizar más allá de estos límites.