JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Resumen Técnico: JEPA-Reasoner: Desacoplamiento del Razonamiento Latente de la Generación de Tokens

Planteamiento del Problema
Los Grandes Modelos de Lenguaje (LLMs) actuales operan bajo un paradigma autorregresivo unificado donde el razonamiento de alto nivel y la generación de tokens de bajo nivel están inextricablemente acoplados. En este marco, el modelo predice el siguiente token basándose en un historial de tokens discretos previamente muestreados. Este acoplamiento crea un cuello de botella fundamental: cualquier error localizado en la selección de tokens (por ejemplo, un desliz gramatical o una elección de palabra subóptima) contamina la ventana de contexto, corrompiendo directamente la trayectoria de todos los pasos de razonamiento subsiguientes. Esto conduce a errores acumulativos en tareas de largo horizonte. Además, los intentos existentes de razonamiento en el espacio latente (por ejemplo, COCONUT, modelos de Profundidad Recurrente) permanecen restringidos por este acoplamiento, intercalando a menudo vectores latentes con tokens o dependiendo de un enmascaramiento causal que impide que los tokens tempranos sean informados por estados de razonamiento futuros. Estos modelos también tienen dificultades para representar la incertidumbre, forzando a menudo un colapso hacia una única trayectoria de razonamiento debido a la naturaleza discreta del acoplamiento a nivel de token.

Metodología
Los autores proponen JEPA-Reasoner, un nuevo paradigma arquitectónico que desacopla completamente el razonamiento en el espacio latente de la generación de tokens. La arquitectura consta de dos módulos distintos:

1. JEPA-Reasoner (El Razonador): Este módulo opera enteramente dentro de un espacio latente continuo y normalizado. Utiliza una Arquitectura Predictiva de Embebido Conjunto (JEPA) adaptada para la generación autorregresiva en el espacio latente.

   • Mecanismo: Toma un input, lo codifica y utiliza bloques Transformer modificados para predecir la siguiente matriz latente que representa los pasos de razonamiento subsiguientes.
   • Restricción: Crucialmente, el Razonador no proyecta salidas hacia probabilidades de vocabulario. En su lugar, la matriz latente generada se normaliza (mediante una capa híbrida de normalización RMS y L2) a una hiperesfera unitaria y se recircula como entrada para el siguiente paso.
   • Entrenamiento: Se entrena mediante Entrenamiento Autosupervisado (SST) para predecir representaciones latentes proporcionadas por un codificador objetivo (cuyos pesos se actualizan mediante Media Móvil Exponencial). El objetivo es una pérdida de similitud de coseno escalada, que fomplica al modelo a aprender diferencias angulares y dinámicas de características consistentes en lugar de la predicción de tokens discretos.

2. Talker (El Generador): Este es un módulo separado, basado en Transformer y ligero, entrenado de forma independiente para reconstruir texto legible por humanos.

   • Mecanismo: Recibe la cadena completa de razonamiento latente generada por el JEPA-Reasoner y reconstruye la secuencia de tokens.
   • Variantes: El artículo introduce el Mono-Talker (solo decodificadores, para tareas sin contexto) y el Dual-Talker (codificadores y decodificadores, para tareas con conciencia de contexto).
   • Dependencia: El Talker es entrenado únicamente para la reconstrucción basada en la salida del Reasoner; no puede realizar predicciones de forma independiente.

Contribuciones Clave
El artículo identifica tres propiedades emergentes de esta arquitectura desacoplada:

1. Contención de Errores Matemáticos: Los autores proporcionan una factorización probabilística que muestra que, debido a que la cadena de razonamiento ($R$) se genera independientemente de la muestra de tokens ($X$), los errores a nivel de token no tienen una vía matemática para influir en la lógica del Reasoner ($\frac{\partial r_{t+1}}{\partial \hat{x}_t} = 0$). Esto aísla el motor de razonamiento de los fallos de muestreo discreto.
2. Guía Continua: El generador (Talker) tiene acceso a toda la trayectoria de razonamiento sin pérdida antes de generar cualquier token, a diferencia de los modelos acoplados donde la generación se intercala con el razonamiento.
3. Representación de la Incertidumbre: El Reasoner puede producir "vectores latentes mixtos" que codifican múltiples caminos de razonamiento simultáneamente. Estos vectores actúan como combinaciones lineales de embeddings de vocabulario, permitiendo al modelo mantener múltiples hipótesis sin la sobrecarga de la búsqueda de haz (beam search) o múltiples decodificaciones.

Resultados Experimentales
Los autores realizaron experimentos controlados en tareas de lenguaje sintético y natural:

• Tareas Sintéticas (Búsqueda en Árbol y CFG):
  • En una tarea de búsqueda en un árbol binario, el modelo alcanzó una precisión del 99.87%. El análisis PCA confirmó que los vectores latentes predichos formaban una nube continua dentro del espacio ocupado por los embeddings de vocabulario discretos, validando la hipótesis del "estado latente mixto".
  • En tareas de generación de Gramática Independiente del Contexto (CFG), el JEPA-Reasoner demostró una robustez superior tanto al ruido de tokens de entrada (reemplazo del 0–30%) como al ruido Gaussiano en el espacio latente, en comparación con un baseline estándar de Transformer y el modelo acoplado COCONUT.
• Tareas de Lenguaje Natural (GSM8K):
  • Un JEPA-Reasoner de 0.9B de parámetros (combinado con un Talker de 198M) fue entrenado con los mismos datos y la misma configuración de hiperparámetros que un baseline de Transformer acoplado de 0.9B.
  • Rendimiento: El JEPA-Reasoner alcanzó una precisión del 51.9% en el benchmark GSM8K de 8 disparos (8-shot), lo que representa una mejora del 149.5% sobre el baseline (20.8%).
  • Escalabilidad: Mientras que el baseline mostró una "meseta de razonamiento" (mejora del 0.1% de 5 a 8 disparos), el JEPA-Reasoner utilizó ejemplos adicionales para lograr una mejora de rendimiento del 131.7%, indicando un escalado efectivo en contexto.
  • Comparación: El modelo de 0.9B superó a un modelo de CoT de 1B (Llama 3.2) y a un modelo de Profundidad Recurrente de 3.5B (Huginn-0125) en el mismo benchmark.

Significancia y Reivindicaciones
El artículo posiciona al JEPA-Reasoner no meramente como un aumento de rendimiento mediante el escalado, sino como una investigación arquitectónica sobre los beneficios del desacoplamiento. Los autores afirman que la transición de un paradigma acoplado a uno desacoplado produce ganancias significativas en robustez y profundidad de razonamiento sin requerir un aprendizaje por refuerzo sofisticado o recetas de entrenamiento industriales.

El trabajo sugiere que la "meseta de razonamiento" observada en modelos pequeños basados en tokens puede superarse moviendo el proceso de razonamiento hacia un espacio latente continuo y abstracto. Al aislar el motor de razonamiento del proceso de muestreo de tokens discretos, la arquitectura proporciona una base más robusta para el razonamiento complejo, permitiendo la contención de errores y la representación de la incertidumbre que es difícil de lograr en los modelos autorregresivos estándar. Los autores concluyen que este enfoque desacoplado representa un nuevo camino a seguir para el diseño de modelos fundacionales.