To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning

Resumen Técnico

1. El Problema: Falta de Generalización en Razonamiento Simbólico

Los modelos de lenguaje grandes (LLM) basados en transformadores han demostrado un gran éxito en tareas de razonamiento formal (matemáticas, ciencia). Sin embargo, presentan una fragilidad crítica: fallan al generalizar a problemas que involucran tokens simbólicos no vistos durante el entrenamiento (nombres de variables, símbolos nuevos).

• Fenómeno observado: Cuando se cambia el nombre de las variables o se introducen símbolos desconocidos, el rendimiento del modelo cae drásticamente.
• Causas conocidas: Trabajos anteriores (Boix-Adserà et al., 2024) señalaron la dificultad de copiar o generar tokens no vistos. Anand et al. (2025) propusieron el "olvido activo" (resetear embeddings periódicamente) como solución empírica, pero sin explicar el mecanismo subyacente.
• Brecha de conocimiento: No existía una explicación teórica ni mecánica clara de por qué los modelos fallan al distinguir múltiples variables no vistas, ni por qué las soluciones arquitectónicas simples no bastaban.

2. Metodología y Análisis Teórico

Los autores investigan la capacidad de los transformadores decoder-only para ejecutar algoritmos en contexto (in-context learning) con símbolos arbitrarios nuevos.

A. Descubrimiento Teórico: Colapso de los (Des)Embeddings
El hallazgo central es que, durante el entrenamiento, los pesos de desembedding (unembeddings) de los tokens no vistos tienden a colapsar hacia el mismo vector.

• Mecanismo: Utilizando el descenso de gradiente estocástico (SGD) con regularización $L_2$ (weight decay) y normalización por capas (LayerNorm), se demuestra teóricamente (Lema 4.1) que los pesos de las clases no presentes en los datos se contraen.
• Consecuencia: Si los tokens no vistos tienen embeddings casi idénticos, el modelo no puede distinguir entre diferentes variables desconocidas (ej. $U_1$ vs $U_2$), lo que impide el razonamiento lógico correcto cuando se requiere copiar o manipular múltiples símbolos nuevos.

B. Arquitectura Propuesta: Atención de Copia (Copy Attention)
Para mitigar la dificultad de copiar tokens, los autores proponen una modificación arquitectónica basada en la idea de Boix-Adserà et al., pero mejorada:

• Se añade una cabeza de atención dedicada ("copy attention") que conecta directamente las embeddings de entrada ($X W_E$) con la salida, en lugar de pasar por las capas MLP estándar.
• Esto facilita la copia de tokens (vistos o no) al proporcionar un atajo directo para propagar la información del token de entrada a la salida.

C. Estrategias de Entrenamiento
Para combatir el colapso, se evalúan tres enfoques:

1. Congelar (Freezing): Mantener los embeddings fijos en su inicialización (ortogonal y distinguible).
2. Reinicialización Activa (Active Forgetting): Reiniciar los embeddings periódicamente durante el entrenamiento para evitar que colapsen.
3. Diversidad de Datos: Aumentar la variedad de símbolos en los datos de entrenamiento.

3. Experimentos y Resultados

A. Experimentos Sintéticos (Lógica Proposicional)
Se entrenaron modelos pequeños (60M-75M parámetros) en problemas de lógica proposicional (cláusulas de Horn) con razonamiento paso a paso (Chain-of-Thought).

• Escenarios de prueba:
  1. Todos los tokens vistos (control).
  2. Solo la variable de consulta es no vista.
  3. Todas las variables son no vistas.
• Resultados Clave:
  • Los transformadores estándar fallan en los escenarios 2 y 3, aprendiendo a ignorar o reemplazar los tokens no vistos.
  • La atención de copia mejora el rendimiento, pero falla si todas las variables son no vistas debido al colapso de embeddings.
  • Congelar embeddings o reinicializarlos periódicamente permite una generalización perfecta a tokens no vistos, incluso con arquitecturas estándar.
  • Se confirma que la diversidad de datos es crucial: modelos con vocabularios más grandes (mayor diversidad) mantienen mejor la capacidad de generalización, pero el colapso sigue siendo un obstáculo sin las otras intervenciones.

B. Experimentos en Modelos Grandes (Gemma 3)
Los autores validaron sus hallazgos en la familia de modelos de código abierto Gemma 3 (1B a 27B parámetros), que incluye 99 tokens deliberadamente no utilizados.

• Colapso en Gemma: Se observó que los embeddings de los 99 tokens no utilizados tienen una alta similitud coseno entre sí (colapsados), mucho mayor que la de los tokens utilizados. Este efecto es más pronunciado en modelos más pequeños.
• Impacto en Fine-tuning: Al ajustar (fine-tuning) Gemma 3 1B en problemas de lógica usando estos tokens no vistos, el entrenamiento fue 10 veces más lento (5000 pasos vs 500 pasos) comparado con el uso de palabras comunes (adjetivos, letras). Esto sugiere que el colapso de embeddings actúa como una mala inicialización para tareas de downstream.

4. Contribuciones Clave

1. Explicación Mecanística del Fallo: Identifican y prueban teóricamente que el colapso de los pesos de desembedding de tokens no vistos es la causa principal de la incapacidad de distinguir múltiples variables nuevas.
2. Validación de Soluciones Existentes: Explican por qué estrategias como el "olvido activo" funcionan (previenen el colapso) y por qué las soluciones arquitectónicas puras (solo añadir copia) son insuficientes sin controlar los embeddings.
3. Solución Integral: Demuestran que la combinación de arquitectura de copia + diversidad de datos + congelamiento/reinicialización de embeddings logra una generalización robusta a símbolos no vistos.
4. Evidencia en Modelos Reales: Proporcionan la primera evidencia empírica de este fenómeno de colapso en modelos de lenguaje grandes reales (Gemma 3) y su impacto negativo en la eficiencia del fine-tuning.

5. Significado e Implicaciones

• Para la Investigación de Razonamiento: Este trabajo sugiere que la capacidad de razonamiento simbólico abstracto en LLMs no es solo una cuestión de capacidad de memoria o atención, sino de la estabilidad de la representación vectorial de los símbolos.
• Para el Diseño de Modelos: Destaca la importancia de gestionar los embeddings de tokens no utilizados o raros. El uso de embeddings fijos o estrategias de reinicialización podría ser esencial para aplicaciones que requieren manipulación dinámica de símbolos (ej. programación, lógica formal).
• Para el Fine-tuning: Advierte que utilizar tokens "reservados" o no vistos en modelos preentrenados para tareas de fine-tuning puede ser ineficiente debido a su alta correlación inicial, requiriendo estrategias de entrenamiento específicas para "descolapsar" sus representaciones.

En resumen, el artículo demuestra que la fragilidad de los LLMs en el razonamiento simbólico no es solo un problema de falta de datos, sino un fallo estructural en cómo el modelo aprende a representar (y distinguir) símbolos nuevos, y ofrece soluciones prácticas tanto a nivel arquitectónico como de optimización.