Resumen Técnico: Razonamiento en Memoria (RiM)

Enunciado del Problema

Los métodos actuales para potenciar las capacidades de razonamiento de los Modelos de Lenguaje Grandes (LLM) dependen en gran medida de escalar la computación en tiempo de prueba mediante la generación de tokens intermedios antes de la respuesta final, como en el razonamiento de Cadena de Pensamiento (CoT). Sin embargo, este enfoque acopla la computación interna directamente a la generación de texto autoregresiva. En consecuencia, el modelo se ve obligado a "pensar en voz alta", asignando una parte significativa de su presupuesto computacional a generar lenguaje natural gramaticalmente fluido en lugar de un razonamiento puramente interno.

Aunque los recientes enfoques de "razonamiento latente" intentan eludir las restricciones del lenguaje natural mediante el uso de representaciones continuas, en gran medida preservan el paradigma de generación autoregresiva paso a paso. En estos métodos, los cálculos intermedios aún se externalizan secuencialmente antes de que el siguiente paso pueda condicionarse sobre ellos. Este artículo argumenta que la cognición humana ofrece un principio de diseño superior: la memoria de trabajo. Los humanos pueden retener y manipular información relevante para la tarea internamente sin externalizar cada pensamiento intermedio. Los autores postulan que los LLM pueden entrenarse para utilizar un espacio de trabajo interno similar, desacoplando el razonamiento del cuello de botella secuencial de la generación de tokens.

Metodología: Razonamiento en Memoria (RiM)

Los autores introducen Razonamiento en Memoria (RiM), un método de razonamiento latente que reemplaza la generación autoregresiva de pasos de razonamiento con bloques de memoria fijos.

Mecanismo Central

En lugar de generar una secuencia de tokens de razonamiento, se proporciona al modelo $K$ bloques de memoria fijos, $m_{1:K}$. Cada bloque consiste en una secuencia predefinida de tokens especiales:
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
donde $<b>$ y $</b>$ delimitan los límites del bloque, y los tokens $<m>$ forman el espacio de trabajo latente.

• Entrada Fija: Estos bloques se añaden a la pregunta de entrada como tokens fijos.
• Paso Único de Avance: Dado que los bloques son entradas fijas y no salidas generadas, toda la secuencia aumentada (pregunta + bloques de memoria) se procesa en un solo paso de avance. Esto elimina el cuello de botella de decodificación secuencial asociado con la generación autoregresiva.
• Representaciones Aprendidas: Aunque las identidades y posiciones de los tokens son fijas, sus representaciones contextuales se aprenden para ser dependientes de la tarea, codificando pasos de razonamiento intermedios.

Currículo de Dos Etapas

Para enseñar al modelo a utilizar estos bloques fijos como memoria de trabajo, RiM emplea un currículo de entrenamiento de dos etapas:

Etapa 1: Supervisión de Pasos de Razonamiento

• Objetivo: Fundamentar los bloques de memoria asignándoles un rol computacional.
• Mecanismo: Para una traza de razonamiento con $T$ pasos, el modelo se entrena con $T$ bloques de memoria. Después de cada bloque $m_t$, el modelo debe predecir el siguiente paso de razonamiento explícito $r_{t+1}$.
• Restricción: Una máscara de atención personalizada asegura que la lectura para el paso $t+1$ pueda atender a la pregunta y a los bloques de memoria anteriores, pero no a los pasos de razonamiento explícitos anteriores. Esto obliga al modelo a codificar el cálculo intermedio necesario dentro del espacio de trabajo latente de los bloques de memoria.
• Supervisión: Utiliza un objetivo de verosimilitud negativa ponderada con un programa de recocido para reducir gradualmente la fuerza de la supervisión.

Etapa 2: Refinamiento de la Respuesta Final

• Objetivo: Convertir los bloques de memoria fundamentados en un espacio de trabajo latente para la predicción de la respuesta final.
• Mecanismo: Se descarta la traza de razonamiento explícita. El modelo utiliza un número fijo de $K$ bloques de memoria. Después de cada bloque, el modelo se entrena para predecir la respuesta final directamente.
• Objetivo: Esto desplaza la supervisión de predecir pasos intermedios a refinar progresivamente la respuesta final a medida que están disponibles más bloques de memoria (cálculo latente).
• Inferencia: En tiempo de prueba, el modelo procesa la pregunta y los bloques de memoria en un solo paso, y la lectura del bloque final proporciona la respuesta.

Enmascaramiento de Atención

Una máscara de atención causal personalizada (Figura 2 en el artículo) es crítica para el método. Previene la fuga de información asegurando que:

1. Los bloques de memoria atiendan a la pregunta y a los bloques de memoria anteriores.
2. Las lecturas objetivo (pasos de razonamiento en la Etapa 1, respuesta final en la Etapa 2) atiendan a la pregunta y a los bloques de memoria disponibles, pero no a otras lecturas objetivo.
   Esto permite que todas las objetivos se predigan en un solo paso de avance sin que el modelo "haga trampa" leyendo la verdad fundamental futura.

Contribuciones Clave

1. Razonamiento en Memoria (RiM): Un método novedoso de razonamiento latente que desbloquea la capacidad de memoria de trabajo de los LLM al proporcionar un espacio de trabajo latente dedicado y fijo para el cálculo intermedio, desacoplando este de la generación autoregresiva.
2. Currículo de Dos Etapas: Un enfoque de entrenamiento estructurado que primero fundamenta los bloques de memoria mediante supervisión de pasos explícitos y luego los refina para la predicción de la respuesta final, permitiendo el uso efectivo del espacio latente sin objetivos auxiliares complejos.
3. Eficiencia y Rendimiento: Evidencia empírica que muestra que RiM iguala o supera las mejores líneas base de razonamiento explícito y latente existentes, manteniendo al mismo tiempo la velocidad de inferencia de los modelos de respuesta directa (sin sobrecarga de generación secuencial).

Resultados Experimentales

Los autores evaluaron RiM en los benchmarks de razonamiento matemático GSM8K (en distribución) y GSM-Hard (fuera de distribución) utilizando modelos GPT-2 y Llama-3.2 (1B y 3B parámetros).

• Rendimiento: RiM superó consistentemente a la línea base de razonamiento latente más fuerte (Coconut) y al Ajuste Fino Supervisado (SFT) de respuesta directa.
  • En GSM8K (Llama-3.2-1B), RiM logró una precisión codiciosa del 42.1%, en comparación con el 36.9% de Coconut y el 23.9% de SFT sin CoT.
  • En GSM-Hard, RiM alcanzó el 10.5%, superando a Coconut (8.5%) y a SFT sin CoT (5.3%).
  • RiM también mostró puntuaciones Pass@8 superiores, indicando que las soluciones correctas se distribuyen a lo largo de la trayectoria latente con alta probabilidad.
• Latencia: RiM alcanzó velocidades de inferencia comparables a SFT sin CoT (Tiempo hasta el Primer Token ~16ms para Llama-3.2-1B), mientras que Coconut fue ~7 veces más lento y SFT con CoT fue ~27 veces más lento debido a la generación autoregresiva.
• Análisis de Representación: Las proyecciones PCA de las representaciones de los bloques de memoria mostraron que el modelo aprende a crear espacios de trabajo latentes distintos y dependientes de la muestra. Las representaciones evolucionan sistemáticamente durante el entrenamiento, confirmando que los bloques no son meros marcadores fijos, sino estados computacionales activos.
• Robustez: El rendimiento se mantuvo estable a través de diferentes presupuestos de memoria en tiempo de inferencia (número de bloques), lo que sugiere que el método es robusto ante las restricciones de implementación.

Significado y Afirmaciones

El artículo afirma que los LLM pueden entrenarse para utilizar la memoria de trabajo como un mecanismo efectivo para el razonamiento latente.

El significado principal radica en la desacoplamiento de la computación de la comunicación. Al reemplazar la generación de pensamientos intermedios con bloques de memoria fijos, RiM permite a los modelos realizar cálculos internos densos y paralelizables sin la sobrecarga de generar lenguaje natural o representaciones continuas secuenciales. Esto resulta en un método que es:

1. Computacionalmente Eficiente: Evita el cuello de botella secuencial de la generación autoregresiva, permitiendo inferencia de un solo paso.
2. Efectivo: Logra un rendimiento de vanguardia entre los métodos de razonamiento latente en benchmarks matemáticos.
3. Escalable: Funciona en diferentes familias y tamaños de modelos sin requerir arquitecturas recurrentes complejas ni objetivos de entrenamiento multi-vía.

Los autores concluyen que este enfoque demuestra una vía viable hacia sistemas de razonamiento más eficientes que imitan la eficiencia cognitiva humana al internalizar el proceso de razonamiento en lugar de externalizarlo como texto.