Resumen Técnico: Explotación de las Simetrías en el Espacio de Pesos para la Aproximación de la Curvatura

Planteamiento del Problema

La estimación precisa de la curvatura (la matriz Hessiana) de una función de pérdida es fundamental para muchos subcampos del aprendizaje automático, incluyendo la optimización de segundo orden, el aprendizaje profundo bayesiano, el aprendizaje continuo y la compresión de modelos. Sin embargo, construir, almacenar e invertir la Hessiana completa es computacional y de memoria prohibitivo para las redes profundas modernas. Si bien las técnicas de aproximación existentes como (E)KFAC, Shampoo y Soap han logrado escalabilidad asumiendo estructuras específicas de factorización de Kronecker o de bloques diagonales, el campo carece de un marco teórico basado en principios que explique por qué estas estructuras específicas funcionan o cómo derivarlas sistemáticamente basándose en la arquitectura del modelo. Además, trabajos anteriores han pasado por alto en gran medida las restricciones estructurales sobre la curvatura impuestas por las simetrías de espacio de pesos, bien conocidas en las redes neuronales (por ejemplo, la invarianza de permutación de las unidades ocultas).

Metodología

Los autores proponen un nuevo marco para aproximar la Hessiana mediante el promedio analítico sobre la "órbita del grupo" generada por las simetrías del espacio de pesos. La idea central reside en el hecho de que muchas funciones de pérdida son invariantes bajo transformaciones de parámetros específicas (acciones de grupo), lo que implica que el gradiente de la pérdida es equivariante bajo estas mismas transformaciones.

1. Promedio de Órbita y Equivarianza del Gradiente

Dado un conjunto actual de parámetros $w$ y un gradiente $g$, los autores consideran el conjunto de todos los pares $(w^{(i)}, g^{(i)})$ alcanzables aplicando transformaciones de simetría de un grupo $G$. Debido a que la pérdida es invariante, el gradiente en cualquier punto de la órbita está relacionado con el gradiente original mediante la acción del grupo. Al promediar estos gradientes sobre toda la órbita, el método agrega información de curvatura sin requerir múltiples cálculos de gradiente.

2. Aproximación de la Hessiana Estructurada

El método deriva una aproximación estructurada de la Hessiana, $\hat{H}$, resolviendo una condición de secante basada en los promedios de órbita de las diferencias de parámetros y las diferencias de gradientes.

• Álgebra de Conmutantes: La aproximación resultante de la Hessiana pertenece al álgebra del conmutante del grupo de simetría. Esto significa que la aproximación puede representarse como una superposición lineal de un pequeño número de elementos base (tensores dispersos y binarios formados por productos de Kronecker de matrices identidad y de unos).
• Factorización: En lugar de almacenar la matriz completa, el método almacena solo los "factores" escalares o matriciales que ponderan estos elementos base. Esto permite una estimación, almacenamiento e inversión tratables.
• Simplificación: Los autores proponen una simplificación práctica donde se asume que el promedio de órbita de las diferencias de parámetros es proporcional a la matriz identidad. Esto conduce a una estimación de la curvatura $\hat{H}_g = S_g^{1/2}$, donde $S_g$ es el promedio de la órbita del producto exterior del gradiente.

3. El Optimizador "Symo"

El marco se instancia como un nuevo optimizador llamado Symo (optimizador basado en simetría). La regla de actualización es $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$. Los autores proporcionan un compilador de tiempo de ejecución (just-in-time) que toma las especificaciones simbólicas de las simetrías de una red y genera rutinas de PyTorch eficientes para calcular los factores necesarios.

Contribuciones Clave

1. Nueva Aproximación de la Hessiana: Un método para derivar una estimación estructurada de la Hessiana a partir de un único cálculo de gradiente mediante el promedio sobre la órbita del grupo de simetría. La aproximación es demostrablemente altamente estructurada y reside en el álgebra del conmutante.
2. Unificación Teórica de Shampoo y Muon: El artículo demuestra que elecciones específicas de grupos de simetría dentro de este marco recuperan exactamente las reglas de actualización de los optimizadores Shampoo y Muon.
   • Al restringir el grupo de simetría a subgrupos de permutación con signo específicos (por ejemplo, aplicando simetría de identidad a las capas de entrada/salida y permutaciones con signo a las capas ocultas), la matriz de curvatura derivada se vuelve de bloques diagonales con factores de Kronecker, lo cual es matemáticamente equivalente a Shampoo/Muon.
   • Esto proporciona una explicación fundamentada para la efectividad de estos optimizadores: explotan implícitamente las simetrías del espacio de pesos, logrando un "punto óptimo" entre la precisión de la aproximación y el costo computacional.
3. Control de Complejidad: El marco demuestra que la complejidad de la aproximación de la Hessiana está gobernada directamente por el tamaño del grupo de simetría elegido. Grupos más grandes producen álgebras de conmutantes de menor dimensión (menos factores, computación más barata) pero órbitas más grandes (potencialmente menor precisión local). Grupos más pequeños producen más términos pero mejor fidelidad local.
4. Validación Empírica:
   • Calidad de la Hessiana: Experimentos en un MLP pequeño muestran que las aproximaciones propuestas (específicamente la $\hat{H}_g$ simplificada) logran una alta similitud de coseno con la verdadera Hessiana a lo largo de la dirección del gradiente negativo, superando a las estimaciones de parámetros promediadas por órbita y igualando el rendimiento de Shampoo en esta métrica.
   • Equivalencia de Optimizadores: En autoencoders de MNIST y nanoGPT (un modelo de lenguaje pequeño), el optimizador Symo produce curvas de entrenamiento casi idénticas a las de Muon, validando la equivalencia teórica bajo restricciones de grupo específicas. Ambos superan significativamente a Adam.

Resultados

• Compromiso entre Precisión y Costo: El artículo explora empíricamente diferentes elecciones de grupos. Un "grupo más grande" (que permite permutaciones de todas las capas, incluyendo entrada/salida) sigue superando a Adam pero es menos efectivo que un "grupo más pequeño" (que restringe las permutaciones a las capas internas y utiliza aproximaciones de bloques diagonales), lo que se alinea con el rendimiento de Shampoo. Esto sugiere que restringir el grupo para que coincida con las simetrías específicas de la arquitectura es crucial para un rendimiento óptimo.
• Deficiencia de Rango: Los autores señalan que la matriz de promedio de órbita completa $S_w$ puede ser deficiente en rango, lo que desestabiliza la inversión. La simplificación propuesta ($S_w \approx I$) evita este problema y, empíricamente, produce mejores estimaciones de curvatura que la solución completa semidefinida positiva.
• Escalabilidad: El método escala con éxito a arquitecturas Transformer (nanoGPT), demostrando que el enfoque consciente de la simetría es aplicable más allá de los simples MLP.

Significado y Reivindicaciones

El artículo afirma proporcionar una lente teórica unificadora para ver los optimizadores de segundo orden existentes. Al enmarcar a Shampoo y Muon como instancias específicas de la aproximación de la Hessiana basada en simetrías, este trabajo mueve estos métodos de la ingeniería heurística a la derivación matemática fundamentada.

Los autores sostienen que su marco:

• Ofrece una receta sistemática para construir aproximaciones de curvatura adaptadas a presupuestos computacionales específicos mediante la selección de grupos de simetría apropiados.
• Proporciona una ruta fundamentada desde las simetrías del modelo hacia la curvatura estructurada, lo que potencialmente inspira nuevos optimizadores para diferentes clases de modelos.
• Sugiere aplicaciones más allá de la optimización, incluyendo la estimación de la incertidumbre, el aprendizaje continuo y la compresión de modelos, aunque estas se presentan como direcciones futuras potenciales más que como resultados demostrados.

El artículo mantiene la modestia respecto a sus limitaciones, reconociendo que la precisión de la expansión de Taylor alrededor del promedio de la órbita no está totalmente acotada para pesos entrenados, y que la garantía formal de alineación de autovectores entre el precondicionador y la verdadera Hessiana sigue siendo un problema abierto. Además, aunque el marco automatiza la derivación de factores, la selección del grupo óptimo para nuevas arquitecturas arbitrarias sigue siendo una tarea dependiente del usuario.