Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Resumen Técnico: Generalización Rápida tras la Interpolación mediante Optimización de Momento Críticamente Amortiguado

Planteamiento del Problema
En regímenes de alta dimensionalidad y baja muestra (comunes en genómica, medicina y finanzas), los modelos de aprendizaje automático suelen alcanzar un rendimiento de entrenamiento casi perfecto (interpolación) mientras fallan al generalizar a datos no vistos. Esta brecha surge porque el paisaje de optimización contiene muchas configuraciones de parámetros que producen una pérdida de entrenamiento idéntica pero difieren significativamente en complejidad y rendimiento de prueba. Mientras que los métodos de regularización (por ejemplo, el decaimiento de pesos) se utilizan tradicionalmente para sesgar la optimización hacia soluciones más simples, el fenómeno del "grokking" revela una desconexión temporal: los modelos suelen memorizar los datos primero y solo transitan hacia la generalización mucho después. Esta transición retardada es computacionalmente ineficiente, requiriendo órdenes de magnitud más épocas para la generalización que para el ajuste. El problema central es cómo seleccionar eficientemente una solución interpolante de baja complejidad inmediatamente después de que el modelo alcanza el umbral de interpolación, en lugar de depender de dinámicas implícitas lentas.

Metodología
Los autores proponen GROKtimizer, una estrategia de optimización bifásica diseñada para separar explícitamente la fase de ajuste de la fase de generalización.

1. Marco Teórico:

   • Dinámica Post-Interpolación: El artículo modela el régimen post-interpolación como un sistema dinámico cuadrático local. Bajo la suposición de que el paisaje de pérdida cerca de una solución interpolante puede aproximarse mediante una forma cuadrática que involucra la matriz Hessiana ($H$) y un término de regularización $L_2$ ($\lambda$), las dinámicas a lo largo de las direcciones planas (espacio nulo de $H$) se asemejan a un oscilador amortiguado.
   • Amortiguación Crítica: Los autores demuestran que el Momento Críticamente Amortiguado (CDM) proporciona la trayectoria óptima para minimizar la norma de los parámetros dentro de este pozo plano. Al ajustar el parámetro de momento $\beta$ a $\beta = 1 - 2\sqrt{\lambda\eta}$ (donde $\eta$ es la tasa de aprendizaje), el optimizador logra una amortiguación crítica.
   • Aceleración de la Convergencia: Teóricamente, el CDM ofrece una aceleración cuadrática sobre el Descenso de Gradiente (GD) clásico en este régimen. Mientras que el GD converge con una escala de tiempo proporcional a $\lambda^{-1}$, el CDM converge con una escala de tiempo proporcional a $\lambda^{-1/2}$. Esto resulta en una reducción cuadrática demostrable en el número de épocas requeridas para alcanzar la solución de mínima norma.
   • Estimación de la Tasa de Aprendizaje: Para implementar el CDM de forma práctica sin calcular la Hessiana completa, los autores proponen un método de Iteración de Potencia corregido por Aitken para estimar el autovalor más grande ($\lambda_{max}$) de la Hessiana. Esto permite el cálculo en línea de la tasa de aprendizaje óptima $\eta_{max}$ en el borde de la estabilidad.

2. Implementación Algorítmica (GROKtimizer):

   • Fase 1 (Interpolación): El optimizador se ejecuta con cero decaimiento de peso para priorizar la convergencia rápida al umbral de interpolación. Esta fase utiliza el momento estándar para ajustar los datos de entrenamiento rápidamente.
   • Fase 2 (Generalización): Una vez alcanzada la interpolación, el optimizador cambia a Momento Críticamente Amortiguado con decaimiento de peso. El parámetro de momento se ajusta a la condición de amortiguación crítica derivada en la teoría, y se aplica el decaimiento de peso para conducir los parámetros hacia la solución de mínima norma dentro de la variedad interpolante.

Contribuciones Clave

1. Caracterización del Régimen Post-Interpolación: Los autores caracterizan formalmente la fase de generalización retardada como un problema de oscilador amortiguado a lo largo de las direcciones planas del paisaje de pérdida, vinculando el fenómeno del grokking con la geometría del espacio nulo de la Hessiana.
2. Aceleración Teórica: Demuestran que el momento críticamente amortiguado acelera la convergencia hacia soluciones interpolantes de baja norma de forma cuadráticamente más rápida que el descenso de gradiente estándar, estableciendo la optimalidad entre los métodos de primer orden bajo el modelo cuadrático local.
3. GROKtimizer: Introducen un esquema de optimización bifásico práctico que cambia dinámicamente de un ajuste rápido no regularizado a una minimización de norma críticamente amortiguada.
4. Validación Empírica: El método se valida en diversos benchmarks sintéticos de grokking y en conjuntos de datos reales de alta dimensionalidad, demostrando mejoras consistentes sobre optimizadores estándar (Adam, AdamW, SGD, Muon).

Resultados Experimentales
El artículo evalúa GROKtimizer en varios entornos de prueba:

• Benchmarks Sintéticos: En tareas como Suma Modular, Suma Binaria y Paridad Dispersa, GROKtimizer supera significativamente a los modelos de referencia. Por ejemplo, en la tarea Gaussiana, redujo la pérdida de validación en varios órdenes de magnitud en comparación con el modelo de referencia más fuerte.
• Datos Reales de Alta Dimensionalidad: En tareas de clasificación biomédica (Leukemia, TCGA) y predicción de propiedades moleculares (QM9), GROKtimizer alcanzó la mayor precisión y la menor pérdida de validación. Las ganancias fueron más pronunciadas en entornos de escasez de datos donde el número de características supera ampliamente al número de muestras.
• Eficiencia de Muestra: En experimentos de MNIST con pocos datos (10–50 muestras por clase), GROKtimizer mostró la mayor ventaja, confirmando que la minimización explícita de la norma mejora la eficiencia de muestra cuando el espacio de soluciones está subdeterminado.
• Preentrenamiento de Modelos de Lenguaje: Los experimentos en modelos WikiText2 y BabyLM mostraron que, si bien la correlación directa entre la norma de los parámetros y la pérdida de validación es menos clara en los transformers (debido a las simetrías de escala y actualizaciones adaptativas), el patrón bifásico sigue siendo válido. Cambiar a la fase de minimización de norma después de la interpolación produjo reducciones adicionales en la pérdida de validación.

Significado y Reivindicaciones
El artículo afirma que GROKtimizer ofrece una solución fundamentada al problema del "grokking" al tratar el entrenamiento post-interpolación como un problema de optimización explícito en lugar de un periodo de espera pasivo. Al aprovechar la geometría del paisaje de pérdida, el método acelera la selección de soluciones de baja complejidad.

Los autores reconcilian sus hallazgos con la hipótesis de los mínimos planos, argumentando que la dimensionalidad del espacio nulo de la Hessiana sirve como una medida robusta de la planitud y la complejidad. Sugieren que minimizar la norma dentro de este espacio nulo es equivalente a navegar hacia mínimos más planos y generalizables.

Limitaciones
Los autores reconocen algunas limitaciones:

• El enfoque está diseñado principalmente para regímenes de datos restringidos y puede ser menos crítico en el preentrenamiento de lenguaje o visión a gran escala donde la abundancia de datos domina el rendimiento.
• El método asume que el modelo tiene la capacidad suficiente para alcanzar la interpolación; no es una solución para modelos subparametrizados.
• El análisis teórico depende de una aproximación cuadrática local del paisaje de pérdida, aunque los resultados empíricos sugieren que el método es robusto ante desviaciones de esta suposición.
• La norma de los parámetros se utiliza como un proxy de la complejidad, lo que podría no capturar todos los aspectos de la generalización del modelo en cada arquitectura.