Characterizing Optimizer-Dependent Training Dynamics Through Hessian Eigenvector Displacement and Localization

Resumen Técnico: Caracterización de la Dinámica de Entrenamiento Dependiente del Optimizador mediante el Desplazamiento y la Localización de los Autovectores de la Hessiana

Planteamiento del Problema
Comprender el entrenamiento de las redes neuronales sigue siendo un desafío debido a la naturaleza no convexa y de alta dimensión del paisaje de pérdida. Si bien la matriz Hessiana (la matriz de segundas derivadas de la pérdida con respecto a los parámetros) es una herramienta estándar para analizar la dinámica de entrenamiento, la generalización y la nitidez (sharpness), la mayor parte del trabajo previo se centra en los autovalores de la Hessiana (magnitud de la curvatura) o en la estabilidad del subespacio de los $k$ principales como un todo. Existe una brección crítica en la comprensión de la evolución de los autovectores individuales líderes, los cuales codifican las direcciones específicas de curvatura en el espacio de parámetros. No está claro si estos autovectores se estabilizan, derivan o se reorganizan durante el entrenamiento por mini-lotes (mini-batch), y si diferentes optimizadores inducen trayectorias geométricas cualitativamente distintas a través del paisaje de pérdida.

Metodología
Los autores investigan la evolución dependiente del tiempo de los autovectores líderes de la Hessiana, $\{v^{(i)}(t)\}$, durante el entrenamiento de perceptrones multicapa (MLP) en los conjuntos de datos MNIST y FashionMNIST. El estudio compara tres optimizadores: Gradiente Descendente Estocástico (SGD), Adam y Minimización de Nitidez Consciente de la Forma (SAM).

Para cuantificar la dinámica de los autovectores, el artículo emplea dos métricas complementarias:

1. Desplazamiento del Autovector: Inspirado en los análisis de sistemas vítreos (glassy systems), esta métrica mide la distancia euclidiana al cuadrado entre un autovector en un tiempo inicial $t_w$ y su estado tras un tiempo de espera $t$. Esto actúa como una función de correlación de dos tiempos para distinguir entre dinámicas estacionarias y comportamientos de "envejecimiento" (aging), donde la tasa de cambio depende de la edad del entrenamiento. Un modelo nulo aleatorio, basado en la Hessiana esperada bajo una inicialización de parámetros aleatoria, sirve como línea de base.
2. Razón de Participación Inversa (IPR): Esta métrica cuantifica la localización de la masa del autovector. Un valor de IPR cercano a $1/N$ (donde $N$ es el número de parámetros) indica delocalización (dispersión uniforme), mientras que valores cercanos a $O(1)$ indican localización (concentración en un pequeño subconjunto de parámetros). En este contexto, el IPR mide cómo se distribuye la sensibilidad a la curvatura a través de los pesos de la red.

Los experimentos utilizan un MLP totalmente conectado de 3 capas entrenado con mini-lotes. La Hessiana se evalúa sobre el conjunto de datos completo en intervalos, y los autovectores se computan mediante la iteración de Lanczos utilizando productos Hessiano-vector para evitar la formación de la matriz completa.

Resultados Clave
El estudio revela comportamientos claros y dependientes del optimizador tanto en el desplazamiento como en la dinámica de localización:

• Desplazamiento y Estabilidad:

  • SGD: Exhibe un comportamiento de "envejecimiento" donde el autovector líder ($v^{(1)}$) cambia rápidamente al principio pero se estabiliza progresivamente. El desplazamiento se satura por debajo de la línea de base aleatoria, lo que indica que el SGD conduce a direcciones de curvatura líder progresivamente más estables.
  • Adam: Muestra una reorganización sustancialmente más fuerte. El desplazamiento del autovector se aplana cerca de la línea de base aleatoria, lo que sugiere una descorrelación casi completa y un cambio persistente de las direcciones de curvatura líderes durante todo el entrenamiento.
  • SAM: Muestra un comportamiento intermedio con una deriva más lenta pero persistente en comparación con SGD.
  • Estas dinámicas son robustas ante variaciones en el tamaño del lote y la tasa de aprendizaje, excepto en el régimen de lote completo (full-batch).

• Localización (IPR):

  • Adam: Induce un fenómeno de fuerte localización donde el autovector líder se concentra en un pequeño subconjeto de parámetros (IPR $\approx 0.1$). Esto sugiere que, bajo Adam, un pequeño subconjunto de pesos contribuye desproporcionadamente a las direcciones de curvatura líderes.
  • SGD: Muestra una tendencia a la delocalización, con valores de IPR que disminuyen hacia la línea de base uniforme ($1/N$), lo que indica que la sensibilidad de la curvatura se distribuye de manera más uniforme entre los parámetros.
  • SAM: Exhibe una localización ligeramente mayor que SGD, pero mantiene una baja varianza entre ejecuciones.

• Alineación: El artículo señala que, si bien los gradientes generalmente se alinean con los autovectores superiores, la alineación de la actualización de Adam permanece baja (por debajo de 0.1), probablemente debido a su parámetro de momento, mientras que las actualizaciones de SGD y SAM se alinean más estrechamente con el gradiente.

Significancia y Reivindicaciones
El artículo sostiene que el seguimiento de la dinámica de los autovectores de la Hessiana captura diferencias clave en el comportamiento de los optimizadores que no son visibles mediante el análisis espectral por sí solo. La principal significancia radica en demostrar que:

1. La elección del optimizador dicta la geometría: Diferentes optimizadores no solo atraviesan el mismo paisaje a diferentes velocidades; inducen formas cualitativamente distintas de exploración del paisaje. SGD estabiliza las direcciones de curvatura, mientras que Adam las reorganiza continuamente.
2. La localización como una estadística estable: A pesar de la mezcla de autovectores bajo Adam, la métrica de localización (IPR) permanece como una estadística estable durante el entrenamiento. Esto sugiere que, incluso cuando las direcciones específicas de curvatura rotan, la naturaleza de la curvatura (si es impulsada por unos pocos parámetros o distribuida) permanece consistente para un optimizador dado.
3. Implicaciones para el análisis: Estos hallazgos sugieren que las inestabilidades de los autovectores pueden facilitar la optimización al permitir la exploración de nuevas regiones del paisaje de pérdida. Además, la localización observada bajo Adam puede tener implicaciones prácticas para la poda (pruning) dirigida y la cuantificación de la incertidumbre, ya que identifica subconjuntos específicos de pesos que dominan las direcciones más convexas.

Los autores concluyen que estas dinámicas son robustas a las variaciones de hiperparámetros (dentro de los regímenes de mini-lote) y que el trabajo futuro podría aprovechar estas observaciones para mejorar la dinámica de los optimizadores y explotar la redundancia de los parámetros.