Resumen Técnico: Dos Velocidades de Aprendizaje: Una Descomposición Representación-Lectura del Grokking y la Doble Descensión

Enunciado del Problema

Las señales estándar para monitorear el entrenamiento de redes neuronales profundas, como la pérdida de entrenamiento y la precisión, a menudo no logran capturar la dinámica compleja de la generalización. Dos fenómenos prominentes complican este panorama:

1. Grokking: Un modelo alcanza un rendimiento de entrenamiento casi perfecto mientras la precisión de prueba permanece en niveles de azar durante un período prolongado, seguida de una aparición abrupta de la generalización.
2. Doble Descensión por Época: La pérdida de entrenamiento disminuye monótonamente, mientras que la pérdida o el error de prueba exhiben un aumento y una caída no monótonos.

Las explicaciones existentes son frecuentemente específicas de la tarea (por ejemplo, dependiendo de estructuras algebraicas en aritmética modular) o se basan en transiciones de régimen binarias (por ejemplo, de aprendizaje "perezoso" a "rico"). Existe una carencia de un marco unificado, agnóstico a la tarea, capaz de diagnosticar estos fenómenos en diversas arquitecturas y tareas, distinguiendo la generalización genuina de artefactos de aprendizaje espurio inducidos por recetas de entrenamiento no estándar.

Metodología: Descomposición Representación-Lectura

Los autores proponen un marco que descompone una red neuronal $f(x)$ en dos componentes distintos:

1. Codificador ($\phi_\theta$): Un mapa no lineal desde el espacio de entrada a un espacio de representación de $N$ dimensiones.
2. Lectura ($W$): Un mapa lineal desde el espacio de representación a los logits de salida.

Esta descomposición permite el seguimiento separado del aprendizaje de representaciones (cambios en $\phi_\theta$) y la calibración de la lectura (cambios en $W$). Para analizar la interacción entre estos componentes, los autores emplean cuatro medidas diagnósticas agnósticas a la tarea derivadas de la geometría representacional y la teoría del Núcleo Tangente Neural (NTK):

1. Dimensión Crítica ($N_{crit}$): Cuantifica la robustez de la separabilidad lineal en el espacio de representación. Se define como la dimensión de proyección más pequeña esperada requerida para preservar la separabilidad con una probabilidad $\geq 0.5$. Una $N_{crit}$ más baja indica mejores variedades "desenredadas" y una mayor generalizabilidad.
2. Medidas Geométricas Relacionadas con la Tarea (GLUE): Un conjunto de cuatro métricas ($D, R, \rho_c, \rho_a$) derivadas de la teoría de la Geometría Vinculada a la Eficiencia de Desenredo (GLUE). Estas caracterizan la dimensión efectiva de la variedad, el radio y la alineación de centros y ejes entre variedades condicionadas por etiquetas.
3. Precisión de Sonda Lineal: Un clasificador lineal regularizado entrenado sobre representaciones congeladas. La brecha entre la precisión de la sonda y la precisión del modelo diagnostica la desalineación de la lectura (sobreajuste o subajuste en relación con las capacidades del codificador).
4. Alineación NTK-Etiqueta de la Última Capa: Mide la alineación entre la matriz Gram de características (NTK de la última capa) y el núcleo de etiquetas. Esto cuantifica qué tan bien la geometría de características actual soporta el flujo de gradiente hacia soluciones discriminativas de etiquetas, evaluado por separado en los conjuntos de entrenamiento y prueba.

Contribuciones y Resultados Clave

1. Dos Velocidades de Aprendizaje en el Grokking

Al aplicar la descomposición a tareas canónicas de grokking (suma modular, composición de permutaciones, paridad dispersa) en arquitecturas MLP y Transformer, los autores refutan el relato estricto de "perezoso a rico", que postula que el aprendizaje de representaciones está latente hasta el inicio del grokking.

• Aprendizaje Continuo de Representaciones: $N_{crit}$ disminuye gradual y continuamente durante todo el entrenamiento, incluida la fase previa al grokking. El aprendizaje de representaciones es lento pero no ausente; procede continuamente en lugar de como una transición de fase repentina.
• Lectura Sesgada hacia el Entrenamiento: Antes del inicio del grokking, la dirección de aprendizaje de la lectura está sistemáticamente sesgada hacia los datos de entrenamiento. La brecha de alineación NTK (entrenamiento menos prueba) es grande y positiva, lo que indica que la lectura sobreajusta el conjunto de entrenamiento mientras las representaciones aún no son lo suficientemente robustas para apoyar la generalización.
• Mecanismo: El grokking surge no de un cambio repentino en los regímenes de aprendizaje, sino de la igualación gradual de la dirección de la lectura a medida que la calidad de la representación ($N_{crit}$) eventualmente alcanza la velocidad de la calibración de la lectura.

2. Diferencias Arquitectónicas en Mecanismos Geométricos

El marco revela que diferentes arquitecturas logran la generalización a través de vías geométricas distintas.

• MLP vs. Transformers: En tareas de suma modular y composición de permutaciones, los MLP y los Transformers exhiben tendencias opuestas en la alineación de centros ($\rho_c$) y la alineación de ejes ($\rho_a$). Mientras que ambas arquitecturas reducen la dimensión y el radio de la variedad, manipulan la geometría interclase de manera diferente para lograr el mismo comportamiento superficial de grokking.

3. Diagnóstico de Aprendizaje Espurio

El marco distingue con éxito la generalización genuina de los artefactos causados por recetas de entrenamiento no estándar.

• Artefacto de Grokking en MNIST: Un ejemplo de grokking previamente reportado en MNIST (inducido por normas de peso iniciales altas y conjuntos de entrenamiento pequeños) se diagnostica como espurio. El marco revela degradación de representaciones (aumento de $N_{crit}$ y disminución de la precisión de la sonda lineal) y sobreajuste de la lectura durante la fase "previa al grokking". La aparente generalización retrasada es un artefacto de que la lectura finalmente corrige una representación degradada, en lugar de una mejora genuina en el aprendizaje de características.
• Doble Descensión por Época: En escenarios de doble descensión inducidos por ruido en las etiquetas, el marco identifica un proceso de dos etapas: una fase inicial de deterioro de representaciones seguida de una fase donde las actualizaciones de lectura están dominadas por señales espurias (etiquetas ruidosas). Eliminar los ingredientes no estándar (por ejemplo, ruido en las etiquetas) elimina tanto el patrón de doble descensión como las firmas diagnósticas.

4. Marco de Análisis de Arriba hacia Abajo

Los autores posicionan esta descomposición como un complemento de "arriba hacia abajo" a la interpretabilidad mecánica de "abajo hacia arriba" (por ejemplo, descubrimiento de circuitos). Al caracterizar la estructura global de las representaciones (geometría de variedades, separabilidad, alineación), el marco proporciona observables consistentes a través de divisiones de entrenamiento/prueba y diversas tareas, ofreciendo una señal confiable para monitorear la dinámica de aprendizaje donde las curvas de pérdida estándar fallan.

Significancia

El artículo establece la descomposición representación-lectura como un marco unificado y agnóstico a la tarea para comprender la dinámica de aprendizaje. Su principal significancia radica en:

• Refinar Cuentas Teóricas: Avanza más allá de la dicotomía binaria perezoso-rico, demostrando que el aprendizaje de representaciones es continuo y que el grokking es impulsado por las velocidades relativas de la mejora de representaciones versus la calibración de la lectura.
• Utilidad Diagnóstica: Proporciona firmas específicas para identificar dinámicas de aprendizaje espurias (por ejemplo, degradación de representaciones o desalineación de lectura) que podrían confundirse con fenómenos de aprendizaje genuinos en configuraciones de entrenamiento no estándar.
• Generalizabilidad: El marco se aplica consistentemente en tareas algorítmicas, conjuntos de datos del mundo real (MNIST) y arquitecturas (MLP, Transformer, ResNet), ofreciendo una herramienta principista para la investigación de interpretabilidad y el monitoreo de modelos críticos para la seguridad.

Los autores concluyen que, aunque el marco se aplica actualmente en la capa lineal final, ofrece una base robusta para futuras investigaciones sobre capas intermedias, estructuras específicas de la tarea y escenarios de aprendizaje continuo.