Resumen Técnico: Cómo domar el Grokking: La geometría de la representación como señal de control

Declaración del Problema

El grokking es un fenómeno en redes neuronales sobreparametrizadas donde los modelos inicialmente memorizan los datos de entrenamiento y solo exhiben una fuerte generalización tras una optimización prolongada. Aunque estudios recientes han vinculado el grocking con el decaimiento de pesos (weight decay), la compresión de características, la formación de circuitos y la dinámica del aprendizaje de representaciones, los factores específicos que gobiernan el cronometraje y la emergencia de esta generalización tardía aún no se comprenden completamente. Un tema común en la literatura existente es que las representaciones aprendidas se comprimen progresivamente en estructuras de menor dimensión. Sin embargo, el trabajo previo trata mayoritariamente la geometría de la representación como una propiedad emergente en lugar de un mecanismo de control directo. Este artículo investiga si la manipulación explícita de la geometría de las representaciones ocultas puede servir como una señal de control para influir en el inicio y la velocidad del grokking.

Metodología

Los autores introducen la Regularización de Dimensionalidad Geométrica (GeomDR), un regularizador espectral a nivel de representación diseñado para controlar directamente la dimensionalidad efectiva de las representaciones ocultas durante el entrenamiento.

Caracterización Geométrica

Para una matriz de representación de una capa oculta $Z \in \mathbb{R}^{N \times d}$, el método primero centra los datos para obtener $\tilde{Z}$. Luego, calcula la matriz de covarianza empírica $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$. La dimensionalidad efectiva ($D_{eff}$) se calcula utilizando la razón de participación de los autovalores ($\mu_i$) de $C$:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
Además, se define una distancia de vecindad local ($\rho$) como la distancia euclidiana promedio a los $k$ vecinos más cercanos de las representaciones normalizadas.

El Objetivo GeomDR

GeomDR impone una penalización sobre la varianza contenida en las direcciones más allá de un subespacio objetivo de dimensión $d^*$. Dados los autovalores ordenados de forma descendente ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), la pérdida por capa es:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
El objetivo de entrenamiento total combina la pérdida de la tarea ($L_{task}$) y el regularizador geométrico:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Crucialmente, el regularizador no está activo desde el principio. Se introduce después de una fase inicial de entrenamiento sin restricciones (paso de activación $t_s$) y su fuerza $\lambda(t)$ se aumenta gradualmente mediante una rampa de coseno. Este diseño imita la trayectoria natural del grokking, donde la memorización precede a la generalización.

Protocolo Experimental

El marco fue evaluado en tres tareas algorítmicas conocidas por exhibir grokking:

1. Suma Modular: Predecir $(a + b) \mod 97$.
2. División Modular: Predecir $(a \cdot b^{-1}) \mod 97$.
3. Composición de Permutaciones: Predecir la composición de dos permutaciones en $S_5$.

Los experimentos utilizaron dos arquitecturas: un Perceptrón Multicapa (MLP) y un Transformer. El estudio varió sistemáticamente los esquemas de intervención (fuerza $\lambda_{max}$ y duración de la rampa $T_{ramp}$), las dimensionalidades objetivo ($d^*$) y los tiempos de inicio de la intervención ($t_s$).

Contribuciones Clave

1. Observación del Colapso Precedente: Los autores demuestran que el colapso de la dimensionalidad precede consistentemente a la transición de la memorización a la generalización (grokking) en los entornos base. Esto sugiere que la compresión geométrica no es meramente un correlato, sino un motor mecanístico de la transición.
2. Introducción de GeomDR: Proponen GeomDR, un regularizador espectral que suprime directamente el espectro de covarianza fuera de un subespacio objetivo, controlando así la dimensionalidad efectiva de las representaciones ocultas.
3. Control Sistemático de la Dinámica: El artículo proporciona un estudio empírico exhaustivo que muestra cómo las intervenciones geométricas pueden alterar sustancialmente las dinámicas de la generalización tardía. Dependiendo del esquema de intervención y la dimensionalidad objetivo, GeomDR puede acelerar el grokking, retrasarlo o suprimirlo por completo.

Resultados

Aceleración del Grokking

En todos los entornos evaluados, GeomDR alteró consistentemente las dinámicas del grokking:

• Rendimiento del MLP: En la suma modular, GeomDR aceleró el grokking hasta 52 veces (reduciendo los pasos de 362K a ~7K). Se observaron mejoras similares para la división modular (39x) y la composición de permutaciones (~23x). En todos los casos de MLP, GeomDR logró una tasa de éxito del 100% donde la base era inestable o fallaba.
• Rendimiento del Transformer: Aunque la dinámica de grokking base para los Transformers ya era más rápida que la de los MLPs, GeomDR aun así proporcionó una aceleración constante (p. ej., ~2x en suma modular), manteniendo una tasa de éxito del 100%.

Influencia de los Hiperparámetros

• Dimensionalidad Objetivo ($d^*$): Las dimensiones objetivo moderadas a grandes ($16 \le d^* \le 64$) produjeron la mayor aceleración. La compresión agresiva ($d^* \le 8$) a menudo ralentizó o desestabilizó el aprendizaje.
• Tiempo de Intervención ($t_s$): Se identificó una "ventana temporal crítica". Las intervenciones aplicadas demasiado pronto (antes de que se formen representaciones útiles) o demasiado tarde (después de que la memorización es rígida) fueron menos efectivas. Los tiempos de activación intermedios (p. ej., $t_s = 2000$ para MLPs) produjeron el grokking más rápido y estable.
• Esquema (Schedule): Los niveles de regularización e intervalos de rampa intermedios generalmente funcionaron mejor; las intervenciones excesivamente fuertes podrían desestabilizar el entrenamiento.

Correlaciones Geométricas

El estudio confirmó que las intervenciones que aceleran el grokking también aceleran el inicio del colapso de la dimensionalidad. En los modelos GeomDR, la brecha temporal entre el colapso de la dimensionalidad y la generalización exitosa se redujo significativamente en comparación con las bases, lo que sugiere un acoplamiento más estrecho entre la reorganización geométrica y la emergencia de la generalización.

Escalabilidad y Eficiencia

Los autores demostraron que GeomDR sigue siendo efectivo incluso cuando el espectro de covarianza se estima a partir de un subconjunto aleatorio de representaciones (p. ej., 5%–25% de los datos), ofreciendo una aceleración de tiempo de ejecución de 2.4x–2.7x sin afectar materialmente la aceleración del grokking.

Significado y Reivindicaciones

El artículo afirma que la geometría de la representación sirve como una señal de control efectiva para el grokking. Al modificar directamente el espectro de covarianza de las activaciones ocultas, los investigadores pueden influir sistemáticamente en cuándo y qué tan rápido una red transiciona de la memorización a la generalización.

Los autores postulan que este trabajo proporciona evidencia de que las intervenciones geométricas ofrecen un enfoque práctico para estudiar e influir en la generalización tardía. En lugar de observar pasivamente los cambios representacionales, el método permite la manipulación directa de las propiedades geométricas para probar hipótesis sobre la dinámica del aprendizaje. Los resultados sugieren que la relación entre la geometría de la representación y el grokking es robusta a través de diferentes tareas y arquitecturas, aunque la magnitud de la mejora varía (siendo más pronunciada en MLPs que en Transformers, probablemente debido a los sesgos inductivos más fuertes de estos últimos).

El artículo concluye que, si bien no establece la dimensionalidad de la representación como el único mecanismo causal, demuestra que las intervenciones geométricas son suficientes para alterar el cronometraje del grokking, ofreciendo una nueva metodología experimental para la investigación de la optimización y la generalización.