A Pre-Training Analogue of Grokking in Language Models: Tracing Delayed Grammatical Generalization

Resumen Técnico: Un Análogo del Grokking en el Preentrenamiento de Modelos de Lenguaje

Planteamiento del Problema

El fenómeno del grokking —donde las redes neuronales transitan de la memorización de datos de entrenamiento a la generalización de ejemplos no vistos tras un periodo prolongado de entrenamiento— ha sido estudiado extensamente en entornos supervisados con divisiones explícitas de entrenamiento/validación y muchos épocas (Power et al., 2022). Sin embargo, esta formulación clásica no se traduce directamente al preentrenamiento de Modelos de Lenguaje Extensos (LLM). En el preentrenamiento, los modelos se optimizan para la predicción del siguiente token sobre un gran corpus no etiquetado (por ejemplo, C4) con una repetición limitada de datos y sin una división natural de entrenamiento/validación supervisada sobre tareas de evaluación posteriores. En consecuencia, no existe un mecanismo estándar para medir si un modelo primero se ajusta a ejemplos específicos y solo después generaliza a variaciones no vistas de la misma tarea durante la fase de preentrenamiento.

Metodología

Para abordar esta brecha, los autores proponen un marco basado en la exposición para construir una división proxy de entrenamiento/validación para estudiar la generalización tardía durante el preentrenamiento.

1. División Proxy Basada en la Exposición:

   • El estudio utiliza BLiMP (Warstadt et al., 2020), un conjunto de datos de pares mínimos que contienen una oración aceptable y una inaceptable dirigidas a fenómenos gramaticales específicos.
   • Para cada par mínimo, los autores extraen una frase crítica: la extensión continua más pequeña que contiene el contraste gramatical y el contexto necesario.
   • División Proxy-Entrenamiento (Expuesta): Ejemplos donde la frase crítica aparece textualmente en los tokens de C4 vistos durante la ventana de preentrenamiento analizada.
   • División Proxy-Validación (No Expuesta): Ejemplos donde la frase crítica no aparece textualmente en la ventana de preentrenamiento.
   • Esta división distingue basándose en la exposición textual, no en si el modelo ha visto la regla gramatical subyacente en otros contextos.

2. Configuración Experimental:

   • Modelos: Modelos de tipo LLaMA de solo decodificador con 35M, 60M y 130M de parámetros.
   • Entrenamiento: Entrenados en C4 inglés durante 2,000 pasos (aprox. 200M de tokens sin relleno) con checkpoints cada 100 pasos.
   • Evaluación: La precisión se mide en pares mínimos de BLiMP. El modelo es correcto si asigna una probabilidad mayor a la oración aceptable.
   • Métrica: El Retraso de Generalización Tardía ($\Delta t$) se define como la diferencia en los pasos de entrenamiento entre cuando la división proxy-entrenamiento y la división proxy-validación alcanzan un umbral de precisión del 80%.

3. Validación y Análisis:

   • Prueba de División Aleatoria Emparejada: Para asegurar que el retraso observado no sea un artefacto del tamaño de la división o de diferencias arbitrarias entre ejemplos, los autores generan 1,000 divisiones aleatorias preservando los tamaños de grupo originales. Comparan el retraso basado en la exposición contra esta distribución nula.
   • Análisis de Representación: Construyen vectores de conceptos gramaticales ($v_i = h_{gram} - h_{ungram}$) para medir la separabilidad de oraciones gramaticales vs. no gramaticales y computan el rango efectivo de estos vectores para analizar la dimensionalidad del espacio de representación.
   • Análisis de Atención: Rastrean los pesos de atención desde el token crítico hacia el token de contexto requerido para resolver la dependencia gramatical, midiendo los cambios en la fuerza y la entropía de la atención.

Resultados Clave

1. Observación de la Generalización Tardía: A través de cinco fenómenos gramaticales (porاً, Concordancia Sujeto-Verbo, Elevación de "There" Existencial), la división proxy-entrenamiento (expuesta) alcanza consistentemente una alta precisión (80%) antes que la división proxy-validación (no expuesta). Esto confirma un análogo del grokking en el preentrenamiento donde los modelos funcionan bien en ejemplos directamente expuestos antes de generalizar a ejemplos estructuralmente similares pero textualmente no expuestos.
2. Validación del Retraso: Las pruebas de permutación de división aleatoria emparejada confirman que este retraso es estadísticamente significativo ($p < 0.05$) y no puede explicarse únicamente por el tamaño de la división. El retraso está específicamente vinculado a la exposición de la frase crítica de forma textual.
3. Cambios en los Vectores de Concepto:
   • Separabilidad: El AUROC de los vectores de conceptos gramaticales generalmente aumenta durante el preentrenamiento, indicando una mejor separabilidad entre oraciones gramaticales y no gramaticales.
   • Dimensionalidad: Después de la transición de la generalización tardía, el rango efectivo de los vectores de conceptos gramaticales aumenta. Esto sugiere que los conceptos gramaticales se distribuyen en un subespacio de mayor dimensionalidad en lugar de colapsar en unas pocas direcciones dominantes.
4. Reorganización de la Atención: La transición hacia la generalización viene acompañada de cambios localizados en la atención. Específicamente, la atención desde el token crítico hacia el token de contexto relevante aumenta en un pequeño número de cabezales de atención específicos. Estos cambios son específicos del fenómeno (por ejemplo, diferentes capas/cabezales para la Concordancia Sujeto-Verbo vs. la Concordancia Determinante-Sustantivo) en lugar de un cambio uniforme en todo el modelo.

Contribuciones Clave

1. Marco de Trabajo: Introducción de un marco basado en la exposición que permite el estudio de la dinámica similar al grokking durante el preentrenamiento de LLMs sin supervisión explícita posterior.
2. Evidencia Empírica: Demostración de que, a través de cinco fenómenos gramaticales, los ejemplos con exposición textual de la frase crítica alcanzan una alta precisión significativamente antes que aquellos sin ella.
3. Robustez: Validación de que el retraso observado no es un artefacto de la división de datos mediante pruebas de división aleatoria emparejada.
4. Perspectiva Mecanística: Evidencia de que la transición conductual hacia la generalización está acompañada de cambios específicos en la geometría del vector de concepto gramatical (aumento de dimensionalidad) y una reorganización localizada de la atención.

Significado y Reivindicaciones

El artículo afirma proporcionar un "análogo del preentrenamiento de la configuración clásica de grokking". Al aprovechar la exposición del corpus como un proxy para los datos de entrenamiento, los autores demuestran que la transición tardía del ajuste a la generalización es observable incluso en el entorno no supervisado de predicción del siguiente token del preentrenamiento.

Los autores enfatizan que sus hallazgos sugieren que la exposición del corpus puede utilizarse para estudiar la generalización tardía de tipo grokking. Los análisis mecánicos revelan que este cambio conductual no es simplemente un cambio en la métrica de rendimiento, sino que tiene su raíz en las representaciones internas del modelo: los conceptos gramaticales se vuelven más separables y ocupan subespacios de mayor dimensionalidad, mientras que los mecanismos de atención se reorganizan para enfocarse más fuertemente en los tokens de contexto relevantes para la dependencia.

El estudio mantiene un alcance modesto, señalando que la división "proxy-validación" representa ejemplos sin exposición textual, no necesariamente ejemplos con una gramática enteramente desconocida. El marco está actualmente limitado a fenómenos gramaticales donde las frases críticas pueden definirse como extensiones sub-sintácticas.