Compositional Generalization from Learned Skills via CoT Training: A Theoretical and Structural Analysis for Reasoning

Aquí presento un resumen técnico detallado del artículo "Compositional Generalization from Learned Skills via CoT Training: A Theoretical and Structural Analysis for Reasoning", publicado en ICLR 2026.

1. Planteamiento del Problema

El entrenamiento de Cadenas de Pensamiento (Chain-of-Thought, CoT) ha mejorado significativamente las capacidades de razonamiento de los Modelos de Lenguaje Grande (LLMs). Sin embargo, los mecanismos subyacentes que permiten que este entrenamiento mejore la generalización (especialmente hacia datos fuera de distribución, OOD) siguen siendo poco claros.

El problema central abordado es la brecha de composicionalidad: los modelos a menudo aprenden hechos básicos pero fallan al combinarlos sistemáticamente para resolver problemas nuevos y más complejos. La pregunta clave es: ¿Mejora el entrenamiento con CoT la generalización tanto en distribución (ID) como fuera de distribución (OOD), y cuáles son los principios teóricos y estructurales que gobiernan esta capacidad?

2. Metodología

Los autores proponen un marco unificado que combina análisis teórico (basado en teoría de la información) y análisis estructural (interpretabilidad mecanicista) para validar su hipótesis de que el CoT fomenta la generalización composicional.

A. Análisis Teórico (Sección 2)

• Acotación del Error de Generalización: Utilizan límites de generalización basados en la divergencia de la información (KL Divergence, $D_{KL}$).
• Descomposición del Error: Demuestran que el error de generalización se puede descomponer en componentes ID (in-distribution) y OOD (out-of-distribution).
• Hipótesis de Composición:
  • En tareas ID, los patrones de composición coinciden con los de entrenamiento; el error tiende a cero con suficiente entrenamiento, independientemente del uso de CoT.
  • En tareas OOD, los problemas presentan patrones de composición novedosos. Los modelos sin CoT fallan porque no aprenden la distribución de los pasos de razonamiento ($P(C|X)$). En cambio, los modelos con CoT aprenden explícitamente a combinar habilidades más simples (sub-tareas), logrando una generalización casi perfecta.
• Robustez al Ruido: Analizan teóricamente cómo la introducción de ruido en los pasos de razonamiento afecta los límites de error, demostrando que el CoT mantiene la generalización siempre que el ruido esté dentro de un umbral tolerable.

B. Análisis Estructural y Experimental (Sección 3)

• Configuración Controlada: Utilizan un entorno sintético basado en hechos atómicos (tripletas entidad-relación-entidad) y hechos multi-hop (dos o más saltos).
  • Datos ID: Composiciones de relaciones vistas en entrenamiento.
  • Datos OOD: Nuevas composiciones de relaciones no vistas.
• Técnicas de Interpretabilidad:
  • Logit Lens: Para proyectar estados ocultos intermedios al espacio de vocabulario y ver qué información se representa en cada capa.
  • Rastreo Causal (Causal Tracing): Para intervenir en las activaciones de las capas y determinar qué circuitos son esenciales para el razonamiento.
• Experimentos de Ruido: Introducen errores deliberados en los pasos intermedios de los datos de entrenamiento para evaluar la robustez del modelo.
• Validación en Datos Reales: Prueban en conjuntos de datos matemáticos (GSM8K, MetaMathQA) y de conocimiento (Wikidata) utilizando fine-tuning con LoRA.

3. Contribuciones Clave

1. Marco Teórico de Generalización Composicional:

   • Formalizan que el entrenamiento con CoT reduce el error de generalización OOD al enseñar al modelo a modelar explícitamente la distribución de los pasos de razonamiento ($P(C|X)$) y la distribución condicional del resultado dado el razonamiento ($P(Y|X, C)$).
   • Demuestran que sin CoT, el modelo solo aprende un mapeo directo $P(Y|X)$, lo que falla ante composiciones novedosas.

2. Descubrimiento del Circuito de Razonamiento de Dos Etapas:

   • Identifican que el entrenamiento con CoT internaliza el razonamiento en un circuito composicional de dos etapas (o más, según los pasos).
   • Hallazgo Crítico: En modelos entrenados con CoT, los resultados intermedios (ej. la entidad puente en un razonamiento de dos saltos) se resuelven en capas más superficiales (ej. capa 3) en comparación con modelos sin CoT (que requieren capas más profundas, ej. capa 5). Esto libera capas profundas para procesar el siguiente paso de razonamiento, mejorando la eficiencia y la capacidad de generalización.

3. Robustez ante Ruido:

   • Demuestran que el entrenamiento con CoT es robusto ante errores en los pasos intermedios, siempre que la proporción de ruido no sea excesiva. Esto es crucial para el entrenamiento práctico donde la anotación perfecta de largas cadenas de pensamiento es costosa o imposible.

4. Insight Fundamental:

   • El entrenamiento con CoT enseña a los modelos "cómo pensar" (fomentando el razonamiento composicional) en lugar de simplemente "qué pensar" (proporcionando solo la respuesta correcta).

4. Resultados Principales

• Convergencia y Generalización: El entrenamiento con CoT acelera significativamente la convergencia y mejora la generalización tanto en ID como en OOD. Los modelos sin CoT muestran una generalización tardía y no sistemática (basada en memorización de patrones) o fallan completamente en OOD.
• Estructura Interna:
  • Con CoT: Se observa un circuito causal claro donde la información fluye secuencialmente a través de capas específicas para cada paso de razonamiento. La resolución de resultados intermedios ocurre más temprano en la red.
  • Sin CoT: El circuito de generalización solo emerge en escenarios ID y requiere capas más profundas para resolver resultados intermedios, lo que limita la capacidad de manejar múltiples pasos en paralelo o secuencialmente en OOD.
• Impacto del Ruido:
  • Con un ruido bajo ($\xi < 0.2$), el modelo mantiene una alta capacidad de generalización.
  • El ruido en ambos saltos (primer y segundo paso) es mucho más perjudicial que el ruido solo en el segundo paso, indicando que la corrección de los pasos iniciales es crítica.
• Validación Real: En GSM8K y datos de Wikidata, el fine-tuning con CoT supera drásticamente al fine-tuning sin CoT (ej. en LLaMA3.1-8B, la precisión sube de ~22% a ~74% en GSM8K), incluso con datos ruidosos.

5. Significado e Implicaciones

Este trabajo proporciona una explicación fundamental de por qué el entrenamiento con CoT es superior a la simple escala de datos o al prompting:

• Diseño de Estrategias de Entrenamiento: Sugiere que la calidad de los datos de CoT no debe medirse solo por la corrección de la respuesta final, sino por la claridad y descomposición de los pasos de razonamiento.
• Eficiencia de Datos: Permite entender cómo modelos más pequeños pueden aprender a razonar complejamente si se les entrena con la estructura adecuada (CoT), reduciendo la necesidad de modelos masivos para tareas de razonamiento.
• Robustez: Ofrece garantías teóricas sobre la tolerancia al ruido, lo que justifica el uso de datos generados automáticamente o con errores parciales en el entrenamiento de modelos de razonamiento avanzados (como los modelos tipo O1 o DeepSeek-R1).
• Arquitectura: Revela que el CoT modifica la arquitectura interna del Transformer, creando circuitos especializados que permiten una mayor profundidad efectiva y una mejor gestión de tareas secuenciales.

En resumen, el paper establece que el CoT no es solo una herramienta de inferencia, sino un mecanismo de entrenamiento que transforma la arquitectura interna del modelo para permitir la composición sistemática de habilidades, logrando así una generalización robusta en escenarios complejos y novedosos.