Resumen Técnico: MUX - Razonamiento Continuo mediante Tokens Multiplexados

1. Planteamiento del Problema

Los modelos de lenguaje autorregresivos modernos resuelven problemas complejos (matemáticas, programación, sentido común) articulando pasos de razonamiento intermedios en lenguaje natural (Cadena de Pensamiento o Chain-of-Thought). Aunque es efectivo, este enfoque discreto es computacionalmente ineficiente. Cada paso de razonamiento transmite solo un subtoken, lo que crea un estricto cuello de botella de ancho de banda. Además, muchos tokens de razonamiento son redundantes, reflejando patrones de comunicación humana en lugar de una computación optimizada.

El razonamiento latente continuo ofrece una alternativa donde los modelos predicen vectores continuos en lugar de subtokens, lo que teóricamente permite un mayor ancho de banda al codificar múltiples subtokens simultáneamente mediante superposición. Sin embargo, los métodos existentes de razonamiento latente enfrentan obstáculos significativos:

• Los métodos globales (que supervisan solo la respuesta final o el punto final de la trayectoria) a menudo conducen a "comportamientos de atajo" (shortcut behaviors) donde los tokens latentes se convierten en marcadores de posición sin información.
• Los métodos locales (que supervisan cada token latente contra un intervalo de razonamiento discreto) a menudo requieren arquitecturas auxiliares complejas (por ejemplo, decodificadores autorregresivos) o fallan en mantener la diversidad de hipótesis necesaria para tareas de búsqueda.

El desafío central es definir un objetivo de supervisión para el razonamiento latente continuo que sea fiel al razonamiento discreto, evite el colapso latente y permita la exploración paralela sin una complejidad arquitectónica excesiva.

2. Metodología: MUX

Los autores proponen MUX, un método de destilación local que entrena a los modelos de lenguaje para realizar razonamiento latente continuo utilizando tokens multiplexados.

Concepto Central

MUX trata la tarea de destilación local como el aprendizaje de una señal continua de dimensión fija que codifica una señal categórica de longitud variable (un intervalo de pasos de razonamiento discretos). Inspirándose en la multiplexación por división de código de los sistemas de comunicación, MUX representa un intervalo de subtokens discretos como una superposición lineal ponderada de sus codificaciones one-hot en el espacio del vocabulario.

Implementación Técnica

1. Objetivos Multiplexados: Para un intervalo de razonamiento discreto $r_i = (r^1_i, \dots, r^{S_i}_i)$, el objetivo de aprendizaje se construye como:
   $$\text{mux}(r_i) := \sum_{j=1}^{S_i} \alpha^{(i)}_j \cdot \text{onehot}(r^j_i)$$
   donde $\alpha^{(i)}_j$ son pesos posicionales positivos normalizados.
2. Ponderación Posicional: Para asegurar que la superposición sea sin pérdidas (es decir, que el intervalo original pueda recuperarse de forma única), los autores proponen esquemas de ponderación específicos:
   • Geométrico: $w_j = \rho^{j-1}$ (decaimiento exponencial).
   • Sinusoidal: Basado en funciones de seno que alcanzan su pico cerca del final del intervalo.
   • Rotatorio: Basado en componentes de coseno promedio análogos a los incrustaciones de posición rotatorias (rotary position embeddings).
   • Garantía Teórica: Estas ponderaciones satisfacen una "condición de separación de suma de subconjuntos", asegurando que diferentes conjuntos de posiciones produzcan masas totales distintas para cualquier subtoken, evitando la ambigüedad incluso si un subtoken aparece múltiples veces.
3. Objetivo de Entrenamiento: El modelo se entrena para coincidir con estos objetivos multiplexados utilizando una cabeza de softmax lineal y una pérdida de divergencia KL ($L_{local}$). Esto se combina con una pérdida estándar de predicción de respuesta ($L_{answer}$) y una pérdida opcional a nivel de trayectoria ($L_{global}$) para alinear las características ocultas con los modelos de razonamiento discretos.
4. Inferencia: Durante la inferencia, el modelo realiza la autorregresión de los tokens latentes en el espacio oculto. La proyección del vocabulario se utiliza únicamente para la supervisión durante el entrenamiento, no para generar los tokens de salida finales.

3. Contribuciones Clave

(1) Razonamiento Latente mediante Tokens Multiplexados

MUX introduce un marco de destilación local simple donde los tokens latentes representan intervalos de razonamiento discretos mediante superposiciones lineales ponderadas. A diferencia de los métodos locales anteriores, no requiere un decodificador autorregresivo auxiliar ni objetivos de caché comprimidos, apoyándose en su lugar en proyecciones del espacio del vocabulario.

(2) Multiplexación Sin Pérdidas

El artículo proporciona pruebas teóricas de que las ponderaciones posicionales específicas (geométrica, sinusoidal, rotatoria) garantizan una multiplexación sin pérdidas. Esto significa que el intervalo de razonamiento discreto se preserva íntegramamente en la representación continua. Los autores demuestran que esta ausencia de pérdidas previene el colapso latente (donde los tokens se vuelven semánticamente homogeneizados), un modo de fallo común en los métodos de razonamiento latente global.

(3) Búsqueda en Paralelo mediante Superposición

Teóricamente, los autores demuestran que los tokens multiplexados pueden implementar naturalmente una búsqueda en anchura (BFS) paralela. Debido a que los tokens están superpuestos, un solo token latente puede representar y actualizar un conjunto entero de hipótesis (un frente de búsqueda) simultáneamente. Esto permite al modelo explorar múltiples caminos de resolución de problemas en paralelo dentro de un marco de supervisión serial.

(4) Superioridad Empírica

A través de 32 configuraciones de evaluación (que abarcan cuatro modelos de lenguaje, dos corpus de entrenamiento y cuatro conjuntos de prueba), MUX supera a los fuertes modelos de referencia tanto en razonamiento discreto como continuo. Notablemente, MUX a menudo supera la línea base de Cadena de Pensamiento discreta (SFT-CoT) de la cual es destilado, a pesar de utilizar significativamente menos tokens (aproximadamente de 2.4 a 5.9 veces menos).

4. Resultados Experimentales

• Razonamiento Matemático: En GSM8K-AUG, GSM8K-AUG-NL, SVAMP, GSM-Hard y MultiArith, MUX logró el mejor rendimiento de razonamiento latente en los modelos GPT-2, LLaMA 3.2 1B, 3B y 3.1 8B. Frecuentemente superó a SFT-CoT, demostando que el razonamiento superpuesto compacto puede generalizar mejor que las trazas discretas originales.
• Benchmarks de Búsqueda: En las tareas de Suma Mínima No Negativa (MNNS) y Juego de 24, MUX superó significativamente a los modelos de referencia (incluyendo Coconut y CoT2), validando la afirmación teórica de que los tokens multiplexados habilitan capacidades de búsqueda paralela.
• Estudios de Ablación:
  • Pérdida Local: Eliminar la pérdida global ($\gamma=0$) y depender únicamente de la pérdida local multiplexada aún produjo resultados superiores en comparación con SIM-CoT, confirmando la eficacia del objetivo local.
  • Ponderación: Las ponderaciones sin pérdidas (Geométrica, Sinusoidal, Rotatoria) superaron consistentemente a la ponderación uniforme con pérdidas, aunque la brecha fue modesta en conjuntos de datos altamente estructurados.
  • *Fragmentación (Chunking): La fragmentación aleatoria de los intervalos discretos en tokens latentes funcionó mejor, actuando como una aumentación de datos estructurada.

5. Significado y Reivindicaciones

El artículo afirma que la superposición sin pérdidas como objetivo de aprendizaje local constituye una condición suficiente para lograr un razonamiento latente continuo fuerte, eficiente e interpretable.

• Eficiencia: MUX logra un alto rendimiento con menos pasos de razonamiento que el CoT discreto, abordando el cuello de botella del ancho de banda computacional.
• Simplicidad: El método evita la complejidad arquitectónica de los enfoques de destilación local previos (por ejemplo, decodificadores auxiliares) manteniendo o mejorando el rendimiento.
• Interpretabilidad: Los análisis de sondeo confirman que los tokens latentes aprendidos codifican contenido de razonamiento fiel e interpretable, contrarrestando la noción de que el razonamiento latente es simplemente un marcador de posición ininterpretable.
• Paralelismo: Los resultados sugieren que la dinámica de búsqueda paralela puede emerger naturalmente de la supervisión serial cuando se utilizan objetivos multiplexados, ofreciendo una nueva vía para mejorar la planificación y la búsqueda en los modelos de lenguaje.

Los autores concluyen que los objetivos locales simples e interpretables pueden hacer que el razonamiento latente sea más fuerte y fácil de entrenar, cerrando la brecha entre la eficiencia de los espacios continuos y la fidelidad del razonamiento discreto.