Understanding the Staged Dynamics of Transformers in Learning Latent Structure

Aquí presento un resumen técnico detallado del artículo "Understanding the Staged Dynamics of Transformers in Learning Latent Structure" en español.

1. Problema y Contexto

Aunque se sabe que los modelos transformadores pueden descubrir estructuras latentes (como jerarquías sintácticas o grafos causales) a partir del contexto, la dinámica temporal de cómo adquieren estos componentes sigue siendo un "caja negra". Existe la hipótesis de que los modelos simplemente "remezclan" datos de entrenamiento en lugar de generalizar verdaderamente. La dificultad radica en que las estructuras latentes en el lenguaje natural son demasiado complejas para aislar y estudiar el progreso del aprendizaje paso a paso.

El objetivo de este trabajo es investigar cómo un transformador aprende estructuras latentes en un entorno controlado y descomponible, permitiendo atribuir las fases de aprendizaje (estancamientos y saltos de rendimiento) a la adquisición de sub-habilidades específicas.

2. Metodología

Entorno Experimental: Alchemy

Los autores utilizan el benchmark Alchemy (Wang et al., 2021), un entorno de meta-aprendizaje por refuerzo compuesto.

• Estructura: El espacio de estados se modela como un grafo cúbico con 8 vértices (estados de "piedra"). Cada piedra tiene 4 características (color, tamaño, redondez, recompensa).
• Transiciones: Se aplican "pociones" que cambian las características de la piedra, moviéndola a otro estado según reglas latentes invariantes. Las pociones vienen en pares complementarios con efectos inversos.
• Tarea: El modelo recibe un conjunto de ejemplos de soporte (transiciones observadas) y debe predecir el estado final de una consulta (query) basándose en la inferencia de la estructura química subyacente.

Configuración del Modelo

• Arquitectura: Un transformador decodificador único pequeño (~2 millones de parámetros, 4 capas ocultas).
• Entrenamiento: Se entrena durante 1000 épocas con un split 90/10 (entrenamiento/validación), asegurando que las químicas de validación no se vean durante el entrenamiento.

Diseños de Experimentos

El estudio se divide en tres formulaciones de tareas para analizar diferentes aspectos de la composición y descomposición:

1. Descubrimiento de Pares de Pociones Ocultos: Se retiene información sobre un par de pociones en los datos de soporte. El modelo debe inferir sus efectos para conectar dos mitades desconectadas del grafo.
2. Composición: Se proporcionan todas las transiciones de 1 paso (hop). El modelo debe componerlas para resolver consultas de múltiples pasos (2, 3, 4 o 5 hops).
3. Descomposición: Se proporcionan secuencias de múltiples pasos en el soporte. El modelo debe descomponerlas para inferir el efecto de un solo paso (1 hop) en la consulta.

Análisis de Factores

En lugar de medir solo la precisión global, los autores factorizan la tarea en eventos interpretables (sub-habilidades) para descomponer la curva de aprendizaje:

• $P(A)$: Predicción dentro del conjunto de soporte (restringir el espacio de búsqueda).
• $P(B|A)$: Identificar la mitad correcta del grafo o el conjunto de piedras alcanzables.
• $P(C|A \cap B)$: Coincidencia exacta con el objetivo.
  La precisión total es el producto de estas probabilidades condicionales.

Intervenciones Causales

Se utiliza el congelamiento de capas (layer freezing) en momentos específicos del entrenamiento para identificar "ventanas de plasticidad". Si congelar una capa en un momento $t$ retrasa o impide la finalización de una etapa, se concluye que esa capa es crítica para el aprendizaje de esa etapa en ese intervalo temporal.

3. Contribuciones Clave

1. Dinámicas de Aprendizaje por Etapas: Demostraron que el aprendizaje no es continuo, sino que ocurre en etapas discretas caracterizadas por mesetas seguidas de saltos abruptos en el rendimiento.
2. Descomposición de Habilidades: Identificaron que cada salto de rendimiento corresponde a la adquisición de una sub-habilidad específica (ej. primero restringir el espacio, luego identificar el sub-grafo, luego el objetivo exacto).
3. Asimetría Composición-Descomposición:
   • Composición: El modelo es invariante a la complejidad de la tarea (longitud de la secuencia). Aprende a componer pasos múltiples casi al mismo tiempo, independientemente de si son 2 o 5 pasos.
   • Descomposición: El modelo muestra una convergencia retardada a medida que aumenta la longitud de la secuencia de soporte. Las etapas de aprendizaje se retrasan significativamente.
4. Ventanas de Plasticidad Específicas por Capa: Mediante intervenciones causales, identificaron que cada capa del transformador tiene una ventana temporal específica durante la cual debe permanecer entrenable para permitir la adquisición de las siguientes etapas. Congelar una capa fuera de esta ventana no tiene efecto, pero hacerlo dentro la detiene.

4. Resultados Principales

En la Tarea de Pociones Ocultas (Descubrimiento Latente)

• Etapa 1: El modelo rápidamente aprende a restringir sus predicciones al conjunto de soporte ($P(A) \to 1$).
• Etapa 2 (Curiosa): El modelo aprende a identificar el objetivo exacto dentro de la mitad correcta antes de aprender a identificar la mitad correcta en sí misma. Esto se debe a un sesgo de adyacencia basado en recompensas: el modelo explota la estructura de recompensas (ej. +15 está adyacente solo a +1) para adivinar vecinos sin entender la estructura completa del grafo.
• Etapa 3: Finalmente, el modelo supera el sesgo y aprende a identificar la mitad correcta del grafo ($P(B|A)$), completando la tarea.

En Composición vs. Descomposición

• Composición: No hay diferencia notable en la velocidad de convergencia entre consultas de 2, 3, 4 o 5 hops. El modelo aprende a componer reglas fundamentales de manera robusta.
• Descomposición: A medida que aumenta la longitud de la secuencia de soporte (de 2 a 5 hops), la convergencia se retrasa drásticamente. La etapa crítica que se retrasa es la identificación del conjunto de piedras alcanzables ($P(B|A)$). El modelo se estanca en un nivel de rendimiento aleatorio durante mucho tiempo antes de hacer el salto final.

Intervenciones de Congelamiento

• Se descubrieron ventanas de plasticidad temporales. Para que el modelo complete las etapas posteriores (como la coincidencia exacta), las capas de transformador deben permanecer entrenables durante un período específico después de completar la etapa anterior.
• Congelar las capas de transformador demasiado pronto causa retrasos masivos o falta de convergencia.
• La capa de embeddings tiene una ventana de plasticidad mucho más corta en comparación con las capas de transformador.

5. Significado e Impacto

Este trabajo ofrece una visión detallada y mecanicista de cómo los transformadores internalizan estructuras lógicas y latentes:

• Validación de Generalización: Al demostrar que el aprendizaje ocurre en etapas alineadas con la estructura lógica del problema (y no solo con patrones superficiales), se refuta la idea de que el modelo solo "remezcla" datos.
• Comprensión de la Complejidad: Revela una asimetría fundamental: los transformadores son excelentes componiendo reglas simples (hacia adelante), pero luchan significativamente al descomponer secuencias complejas para inferir reglas atómicas (hacia atrás/inverso).
• Implicaciones para la Arquitectura: La identificación de ventanas de plasticidad sugiere que el aprendizaje profundo no es un proceso uniforme; diferentes capas y diferentes momentos del entrenamiento son críticos para diferentes sub-tareas. Esto podría guiar futuras estrategias de entrenamiento, como el curriculum learning o el ajuste fino selectivo.

En resumen, el paper establece un marco metodológico riguroso para estudiar la evolución de capacidades en modelos de lenguaje, demostrando que el aprendizaje de estructuras latentes es un proceso secuencial, dependiente de la arquitectura y sensible a la complejidad de la tarea de manera asimétrica.