Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

Resumen Técnico: Regímenes de Decaimiento de Pesos en Transformers de Comprensión

Enunciado del Problema
La comprensión (grokking) se refiere al fenómeno donde los transformers entrenados en tareas algorítmicas (específicamente aritmética modular) exhiben una transición abrupta de la memorización a la generalización mucho después de que la pérdida de entrenamiento se satura. Mientras que estudios mecanísticos han identificado la formación de circuitos internos (por ejemplo, representaciones de Fourier) y cabezas de inducción como impulsores de esta transición, la literatura existente carece de un "diagnóstico en línea económico" para rastrear estas dinámicas en tiempo real. Además, el papel del decaimiento de pesos ($\lambda$) como parámetro de control para estos regímenes no ha sido mapeado sistemáticamente a través de escalas de modelos, ni se ha aislado la contribución específica de la coordinación de cabezas de atención de los efectos de la magnitud de los pesos.

Metodología
Los autores entrenan transformers de solo decodificador en tareas de aritmética modular (mod+, mod-, mod$\times$, mod$\div$) utilizando el optimizador AdamW. El estudio emplea un barrido denso de valores de decaimiento de pesos y varía las escalas de los modelos (desde 0.82M hasta 85M parámetros) y las arquitecturas (Transformers, MLPs, LSTMs, Mamba).

Los componentes metodológicos clave incluyen:

1. Parámetros de Orden en Línea: Se calculan dos diagnósticos escalares en cada paso de entrenamiento utilizando únicamente las activaciones de atención:
   • Similitud Coseno Media Pares de Cabezas de Atención ($\bar{s}(t)$): Mide la coordinación (sincronización) entre las cabezas de atención.
   • Desviación Estándar de la Entropía ($\sigma_H(t)$): Mide la dispersión de la entropía de Shannon entre las cabezas, rastreando la diferenciación.
2. Escalado de Tamaño Finito: Los autores ajustan una función logística a la probabilidad de comprensión ($P(\text{grok})$) frente al decaimiento de pesos para estimar un umbral crítico $\lambda_c$. También ajustan una ley de potencia al "tiempo hasta la comprensión" ($t_{\text{grok}}$) para determinar el exponente crítico $\nu$.
3. Intervenciones Causales: Para distinguir entre la estructura del patrón de cabezas y la magnitud de los pesos, los autores realizan experimentos emparejados donde las cabezas de atención se reinician frente a donde las normas de los pesos se recortan en el pico de la fase de diferenciación.
4. Verificación Formal: Las identidades diagnósticas (límites y propiedades algebraicas) se verifican mediante máquina en Lean 4 para garantizar la buena formación matemática.

Contribuciones Clave

1. Umbral Crítico Cuantitativo: El estudio identifica un umbral crítico de decaimiento de pesos agudo para el transformer canónico de 4 capas y 8 cabezas en tareas mod+: $\lambda_c = 0.0158$ (IC 95% [0.0109, 0.0200]).
2. Diagrama de Regímenes: Un mapa de dos ejes ($\lambda$, $N$) revela tres regímenes distintos:
   • Memorización: $\lambda < \lambda_c$ (comprensión cercana a cero).
   • Comprensión de Desarrollo: $\lambda \ge \lambda_c$ (alcanzando ~100% de tasa de comprensión, con el tiempo hasta la comprensión disminuyendo a medida que aumenta $\lambda$).
   • Colapso: $\lambda = 10$ (todas las cabezas convergen a patrones idénticos, $\bar{s}=1.0$).
3. Dinámica de Dos Fases: La trayectoria de entrenamiento se caracteriza por:
   • Fase 1 (Sincronización): Las cabezas convergen ($\bar{s}$ aumenta de ~0.93 a 0.995) coincidiendo con el inicio de la comprensión.
   • Fase 2 (Diferenciación): Las cabezas divergen ($\bar{s}$ desciende a ~0.88, $\sigma_H$ aumenta) mientras la precisión de prueba se mantiene alta.
4. Dimensión por Cabeza como Modulador de Amplitud: La amplitud de la diferenciación ($\sigma_H$) está empíricamente dominada por la dimensión por cabeza ($d/H$), saturándose monótonamente a medida que aumenta $d/H$.
5. Alcance Trans-Arquitectónico: La transición controlada por decaimiento de pesos se observa en arquitecturas sin atención (MLP, LSTM, Mamba), aunque el umbral crítico $\lambda_c$ se desplaza significativamente (abarcando aproximadamente un orden de magnitud) dependiendo de la arquitectura.

Resultados

• Exponente de Escalado: El tiempo hasta la comprensión sigue una ley de potencia $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ con un exponente empírico $\nu = 0.757$ (IC 95% [0.725, 0.799]). Los autores señalan explícitamente que este valor no coincide con exponentes de referencia estándar (por ejemplo, $\nu=1/2$ o Ising 3D $\nu \approx 0.63$) bajo su densidad de cuadrícula actual, posponiendo la identificación de la clase de universalidad para trabajos futuros.
• Fragilidad de Largo Horizonte: En horizontes extendidos (20,000 épocas), se observa un colapso de "anti-comprensión" dependiente de la semilla en algunas trayectorias, caracterizado por un patrón de cinco etapas que termina en decadencia de la precisión. Esto se identifica como una fragilidad y no como un ciclo universal.
• Aislamiento Causal: Reiniciar las cabezas de atención reduce significativamente la amplitud de la Fase 2 ($\sigma_H$) en comparación con los controles, mientras que el recorte de la norma de los pesos no lo hace. Esto aísla el efecto a la estructura del patrón de cabezas y no a la magnitud de los pesos.
• Replicación Multi-Tarea: El patrón de control por decaimiento de pesos se mantiene a través de cuatro operaciones modulares y dos escalas de modelos, confirmando la robustez del mapa de regímenes más allá de la suma.

Significado y Afirmaciones
El artículo afirma proporcionar un marco diagnóstico económico y computable en línea para rastrear las dinámicas de comprensión utilizando únicamente activaciones de atención, complementando diagnósticos más costosos del paisaje de pérdida como el coeficiente de aprendizaje local refinado (rLLC). Establece el decaimiento de pesos como un parámetro de control empírico primario que separa los regímenes de memorización, comprensión de desarrollo y colapso.

Los autores son modestos en su alcance:

• Afirman ninguna clase de universalidad para el exponente observado $\nu \approx 0.76$, señalando que los exponentes de referencia probados caen fuera de sus intervalos de confianza.
• No afirman que el decaimiento de pesos sea necesario para la generalización en todas las arquitecturas, solo que actúa como un parámetro de control dentro de las arquitecturas de atención estándar.
• Posponen explícitamente las afirmaciones sobre el mecanismo causal respecto a por qué ocurre la diferenciación para trabajos de seguimiento, enfocándose en cambio en cuándo y cómo se manifiesta.
• Los hallazgos se limitan a la aritmética modular en modelos de atención de transformers pequeños; las afirmaciones sobre modelos de lenguaje grandes o universalidad más amplia están fuera del alcance.

El trabajo se posiciona como una contribución fenomenológica y diagnóstica, proporcionando un "mapa de regímenes" y "parámetros de orden" que complementan el trabajo teórico concurrente sobre escalado de tamaño finito, contracción de Lyapunov y relajación de vidrio.