Resumen Técnico: ¿Grokking o Fallo? Cómo la Baja Precisión Impulsa Picos de Pérdida tipo Resorte

1. Declaración del Problema

Las redes neuronales profundas entrenadas sin regularización explícita a menudo exhiben un fenómeno conocido como "Grokking", donde los modelos logran una generalización repentina mucho después de alcanzar una precisión de entrenamiento perfecta. Este proceso va frecuentemente acompañado del Mecanismo de Resorte: inestabilidades periódicas caracterizadas por el crecimiento rápido, casi exponencial, de la norma de los parámetros de la última capa, seguido de picos abruptos en la pérdida de entrenamiento.

Aunque la literatura existente atribuye estos picos a dinámicas de optimización intrínsecas —como el Borde de la Estabilidad (EOS) o las interacciones entre las escalas de gradiente y los optimizadores adaptativos—, el mecanismo de disparo preciso permanece poco claro. Este trabajo investiga si estas inestabilidades son inherentes al paisaje de optimización o artefactos de la aritmética de precisión finita.

2. Metodología y Marco Teórico

Los autores proponen que el Mecanismo de Resorte no es un fenómeno de optimización intrínseco, sino un resultado de los límites de precisión de la aritmética de punto flotante durante el cálculo de la pérdida de Entropía Cruzada (CE).

2.1 Colapso de Softmax (SC)

El núcleo del mecanismo reside en el error de absorción inherente a la suma de punto flotante. A medida que avanza el entrenamiento y el modelo entra en un régimen de alta confianza (precisión de entrenamiento perfecta), el logit de la clase correcta ($z_r$) se vuelve significativamente mayor que los logits de las clases incorrectas ($z_k$).

• En el cálculo de Softmax, el denominador implica una operación de log-sum-exp: $Z = \log(\sum \exp(z_k))$.
• Cuando el margen $z_r - \max_{k \neq r} z_k$ supera un umbral determinado por la precisión de la mantisa (por ejemplo, $\approx 16$ para float32), los términos más pequeños son absorbidos por el término más grande debido a la precisión limitada.
• En consecuencia, la probabilidad calculada para la clase correcta se vuelve exactamente 1, y el gradiente con respecto al logit de la clase correcta, $\partial L / \partial z_r$, se redondea a exactamente cero. Este fenómeno se denomina Colapso de Softmax (SC).

2.2 Ruptura de la Restricción de Suma Cero

En aritmética exacta, la suma de los gradientes a través de todas las clases para una sola muestra es cero ($\sum_k \nabla_{W_k} L = 0$), asegurando que la media global de los pesos del clasificador ($W_G$) permanezca estática.

• Bajo SC, el gradiente para la clase correcta desaparece, mientras que los gradientes para las clases incorrectas permanecen no nulos (aunque pequeños).
• Esto rompe la restricción de suma cero, introduciendo una deriva neta en la actualización de la media global del clasificador $W_G$. Específicamente, $W_G$ deriva en la dirección de $-\mu_G$, donde $\mu_G$ es la media global de las características de la capa penúltima.

2.3 Inflación Numérica de Características (NFI)

El trabajo identifica un bucle de retroalimentación positiva denominado Inflación Numérica de Características (NFI):

1. La deriva en $W_G$ altera el flujo de gradiente hacia atrás hasta la capa de características.
2. La media de características $\mu_G$ recibe una deriva en la dirección de $-W_G$.
3. Esto crea una dinámica acoplada donde $W_G$ y $\mu_G$ se vuelven antiparalelos, y sus normas crecen exponencialmente.
4. A medida que estas normas crecen, el margen para algunas muestras (particularmente valores atípicos) eventualmente cae por debajo del umbral de absorción, haciendo que el gradiente de la clase correcta reaparezca abruptamente.

2.4 El Pico de Resorte

El reaparecimiento del gradiente de la clase correcta desencadena el pico de pérdida:

• Antes del pico, el gradiente de la clase correcta es cero y los gradientes de las clases incorrectas son diminutos. El acumulador de segundo momento del optimizador Adam ($v_t$) decae, haciendo que la tasa de aprendizaje efectiva ($\eta / \sqrt{v_t + \epsilon}$) se vuelva extremadamente grande.
• Cuando el gradiente reaparece repentinamente, esta enorme tasa de aprendizaje efectiva resulta en una actualización de parámetros desproporcionadamente grande, impulsando la pérdida de nuevo a niveles de adivinanza aleatoria (el "Resorte").

3. Contribuciones Clave

1. Explicación Teórica del Resorte: El trabajo proporciona una prueba rigurosa de que el Mecanismo de Resorte es un artefacto numérico causado por el cálculo de pérdida CE de precisión finita, vinculándolo específicamente al Colapso de Softmax y la subsiguiente ruptura de la restricción de gradiente de suma cero.
2. Identificación de NFI: Los autores definen la Inflación Numérica de Características (NFI), un bucle de retroalimentación entre la media global del clasificador y la media global de características, que impulsa el crecimiento exponencial de las normas de parámetros observado antes de los picos de pérdida.
3. Distinción de Dinámicas Intrínsecas: El trabajo demuestra que el Resorte es distinto del Borde de la Estabilidad (EOS). Mientras que EOS implica oscilaciones debido a la curvatura del Hessiano, el Resorte ocurre en regímenes donde los valores propios del Hessiano se anulan (régimen de interpolación) y se elimina aumentando la precisión.
4. Intervenciones Prácticas: El trabajo propone y valida métodos para suprimir la NFI, incluyendo:
   • Usar float64 para el cálculo de la pérdida (eliminando errores de absorción).
   • Hacer cumplir la restricción de suma cero mediante proyección de gradientes.
   • Aplicar Normalización por Lotes inmediatamente antes del clasificador final para eliminar el componente de media global.
   • Aumentar el parámetro $\epsilon$ de Adam para limitar la tasa de aprendizaje efectiva.

4. Resultados Experimentales

Los autores validan su teoría a través de múltiples arquitecturas (MLP, VGG11, ResNet18, ViT, Transformers) y conjuntos de datos (Aritmética Modular, CIFAR-10, FineWeb para LLMs).

• Sensibilidad a la Precisión: Cambiar el cálculo de la pérdida a float64 elimina completamente los picos de Resorte, incluso cuando los parámetros del modelo permanecen en float32. Esto confirma que la inestabilidad se origina en el cálculo de la pérdida, no en la dinámica del modelo.
• Dependencia Arquitectónica: El fenómeno se observa en la mayoría de las arquitecturas, pero está notablemente ausente en ResNet18. Los autores atribuyen esto a la capacidad de ResNet18 para reducir la masa de probabilidad residual ($\epsilon$) lo suficientemente rápido como para evitar que el bucle de retroalimentación NFI entre en el régimen de crecimiento exponencial antes de que el modelo se estabilice.
• Mitigación:
  • La Normalización por Capa (LN) aplicada antes del clasificador falla en mitigar el efecto y puede acelerar los picos al restringir las normas de características.
  • La Normalización por Lotes (BN) aplicada inmediatamente antes de la capa lineal final elimina con éxito los picos al centrar las características y eliminar el componente de deriva.
  • El Suavizado de Etiquetas elimina los picos inducidos por precisión, pero introduce una inestabilidad diferente relacionada con el Borde de la Estabilidad debido a valores propios del Hessiano no nulos.
• Contexto LLM: En el entrenamiento de Modelos de Lenguaje Grandes, aunque los picos de Resorte por lote completo son raros debido a la estocasticidad de mini-lotes, el mecanismo NFI aún impulsa un crecimiento anormal de logits. Los autores notan que en el modelado de lenguaje, los desequilibrios en la frecuencia de tokens (ley de Zipf) pueden crear una gran media global $W_G$ incluso sin colapso numérico, conduciendo a una dinámica diferente donde la baja precisión puede suprimir parcialmente la divergencia inducida por frecuencia.

5. Significado y Afirmaciones

El trabajo reinterpreta el Mecanismo de Resorte no como un comportamiento de optimización misterioso, sino como una dinámica numérica del entrenamiento de precisión finita.

• Reinterpretación del Flujo de Gradiente: Los resultados destacan una brecha entre los análisis teóricos de flujo de gradiente (que asumen precisión infinita y a menudo predicen crecimiento logarítmico de normas) y el entrenamiento práctico, donde los errores numéricos pueden impulsar un crecimiento exponencial.
• Estabilidad en Entrenamiento de Baja Precisión: A medida que el campo se desplaza hacia precisiones más bajas (BF16, FP8, FP4) para el entrenamiento a gran escala, comprender la NFI es crítico. Los autores argumentan que los errores numéricos en el cálculo de log-probabilidades representan un modo de fallo distinto que puede desestabilizar el entrenamiento, separado de los errores de cuantización en la multiplicación de matrices o los sumideros de atención.
• Alcance Moderado: Los autores reconocen limitaciones, señalando que su análisis se basa en la suposición del Modelo de Características Sin Restricciones e investiga principalmente la interacción entre la última capa y las características penúltimas. No afirman explicar todas las formas de inestabilidad en el entrenamiento, sino que se dirigen específicamente a los picos de pérdida periódicos observados en el entrenamiento a largo plazo sin regularización.

En conclusión, el trabajo afirma que el Resorte es un fallo, no una característica, impulsado por los errores de absorción de la aritmética de punto flotante, y ofrece explicaciones concretas y comprobables e intervenciones para el crecimiento anormal de parámetros en las etapas tardías del entrenamiento.