Rethinking Bregman Divergences in Kronecker-Factored Optimizers

Resumen Técnico: Replanteando las Divergencias de Bregman en Optimizadores Factorizados por Kronecker

Planteamiento del Problema
Los optimizadores de estilo Shampoo aproximan la matriz de covarianza del gradiente $C = E[\text{vec}(G)\text{vec}(G)^\top]$ utilizando una estructura factorizada por Kronecker $L \otimes R$ para reducir la complejidad computacional. Un trabajo reciente (Lin et al., 2026) unificó varias variantes de Shampoo enmarcándolas como problemas de minimización de divergencias de matrices de Bregman entre $C$ y $L \otimes R$. Sin embargo, persiste una brecha crítica: para matrices de covarianza densas genéricas, la factorización de Kronecker exacta es imposible, lo que significa que el error de aproximación es inevitable. No está claro cómo las diferentes divergencias de Bregman distribuyen este error a través del espectro de la covarianza y si los precondicionadores resultantes confían por igual en todas las partes del espectro. Además, dado que la covarianza del gradiente suele diferir significativamente de la verdadera Hessiana, no todos los componentes espectrales son adecuados para el precondicionamiento. La pregunta central abordada es: Cuando la coincidencia exacta de Kronecker es imposible, ¿qué partes del espectro de la covarianza debería aproximar y confiar un optimizador estructurado?

Metodología y Análisis
Los autores abordan este problema a través de dos lentes analíticas principales: la geometría espectral de las aproximaciones de Bregman y la fiabilidad empírica de los eigenspacios de la covarianza.

1. Geometría Espectral de la Aproximación de Bregman:
   El artículo caracteriza cómo las divergencias de Frobenius, von Neumann y LogDet penalizan los errores de aproximación de manera distinta. Al descomponer la divergencia en términos de espectro marginal y un término de acoplamiento espectral cruzado $g_F(\lambda_i, \omega_j)$, los autores demuestran que:

   • Divergencia de Frobenius: El término de acoplamiento $-\lambda_i \omega_j$ impone una penalización de magnitud bilineal, enfatizando fuertemente el espectro superior (eigenvalores grandes).
   • Divergencia de von Neumann: El término $-\lambda_i \log \omega_j$ retiene una preferencia por el espectro superior pero la atenúa logarítmicamente.
   • Divergencia de LogDet: El término $\lambda_i / \omega_j$ crea una distorsión de razón relativa, haciendo que la aproximación sea altamente sensible a los eigenvalores pequeños (el espectro de la cola).

   El Teorema 2 establece que los puntos estacionarios de estos problemas de minimización no están determinados por el residuo bruto $(S^* - C)$, sino por un residuo ponderado por la divergencia $\nabla^2 F(S^*)[S^* - C]$. En consecuencia, los factores de Kronecker resultantes se alinean con las preferencias espectrales específicas de la divergencia elegida (por ejemplo, los factores de LogDet son más sensibles a los eigenvalores pequeños, mientras que los factores de Frobenius se centran en los subespacios dominantes).

2. Fiabilidad Espectral del Precondicionamiento de la Covarianza:
   Los autores investigan empíricamente la alineación entre la covarianza del gradiente y los eigenspacios de la Hessiana en un modelo Transformer entrenado en SST-2. Los resultados (Tabla 3) indican que el eigenspacio superior de la covarianza está sustancialmente mejor alineado con la Hessiana que los componentes de eigenvalores inferiores. Se observa que el espectro de la cola es más ruidoso y menos fiable para el precondicionamiento, lo cual es consistente con hallazgos previos (Kunstner et al., 2019; Thomas et al., 2020).

3. Solución Propuesta: Optimizador de Kronecker Consciente del Subespacio (BregTop):
   Motivados por el desajuste entre la preferencia espectral inducida por la divergencia y la fiabilidad real del espectro, los autores proponen BregTop. Este optimizador combina las fortalezas de ambas observaciones:

   • Selecciona los factores de Kronecker utilizando una divergencia sensible al espectro superior (por ejemplo, von Neumann o Frobenius).
   • Aplica precondicionamiento basado en eigenvalores solo dentro del subespacio superior (los $\lceil \rho mn \rceil$ entradas más grandes del espectro conjunto de Kronecker).
   • Trata el subespacio inferior complementario con una constante de aceleración isotrópica adaptativa ($\chi_{\rho, q}$), en lugar de intentar precondicionar la cola ruidosa con información de curvatura potencialmente engañosa.

   Teóricamente, esto descompone el precondicionador en un componente superior escalado por la curvatura y un componente de aceleración uniforme para el espacio inferior (Teorema 4).

Contribuciones Clave

• Caracterización de la Geometría Espectral: El artículo proporciona una caracterización formal de cómo las divergencias de Frobenius, von Neumann y LogDet inducen diferentes penalizaciones dependientes de los eigenvalores, mostrando que los factores de Kronecker estacionarios están gobernados por residuos ponderados por la divergencia en lugar de errores brutos.
• Fiabilidad Espectral Empírica: El trabajo demuestra que, si bien el espectro completo de la covarianza es poco fiable para el precondicionamiento, el espectro superior mantiene una fuerte alineación con la Hessiana, justificando un enfoque consciente del subespacio.
• Diseño de Optimizador Consciente del Subespacio: Los autores introducen BregTop, un algoritmo que cambia adaptativamente entre el escalado basado en eigenvalores en el subespacio superior confiable y la aceleración isotrópica en el subespacio inferior ruidoso.

Resultados Experimentales
Los autores evaluaron BregTop frente a Shampoo estándar, KL-Shampoo, VN-Shampoo y F-Shampoo en un Transformer de 4 capas entrenado en el conjunto de datos completo de SST-2. La métrica utilizada fue el número de pasos requeridos para que la pérdida de entrenamiento suavizada por EMA cayera por debajo de 0.05.

• BregTop-VN-v2 alcanzó la pérdida objetivo en 1003 pasos, superando al Shampoo estándar (1119 pasos) y al baseline VN-Shampoo (1127 pasos).
• Las variantes de BregTop mejoraron consistentemente sus bases de divergencia correspondientes (por ejemplo, BregTop-VN mejoró sobre VN-Shampoo; BregTop-F mejoró sobre F-Shampoo), validando el beneficio de aislar el subespacio superior para el precondicionamiento.

Significancia y Reivindicaciones
El artículo sostiene que el diseño de optimizadores estructurados debe tener en cuenta el error de aproximación inevitable en la factorización de Kronecker y la fiabilidad variable del espectro de la covarianza. Al reconocer que diferentes divergencias de Bregman priorizan inherentemente diferentes regiones espectrales, y que solo el espectro superior está alineado de manera fiable con la Hessiana, el enfoque propuesto consciente del subespacio ofrece una estrategia de precondicionamiento más robusta. Los autores señalan que su diseño es consistente con la dinámica de entrenamiento observada en Song et al. (2025), donde el aprendizaje ocurre en el subespacio inferior; por lo tanto, aplicar aceleración isotrópica allí preserva el flujo de gradiente necesario mientras se amortizan las actualizaciones ruidosas en las direcciones de alta curvatura mediante el precondicionamiento del subespacio superior. Los resultados sugieren que este enfoque híbrido mejora la eficiencia de los pasos sin requerir información de segundo orden completa.