How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

Resumen Técnico: Cómo el Optimizador Moldea las Soluciones Aprendidas en Redes Neuronales Equivariantes

Enunciado del Problema
Las redes neuronales equivariantes codifican simetrías geométricas directamente en su arquitectura, proporcionando un sesgo inductivo natural para datos con simetrías geométricas conocidas. Sin embargo, las restricciones de equivarianza estrictas a menudo dificultan la optimización, lo que conduce a paisajes de pérdida que contienen puntos críticos y mínimos espurios. En consecuencia, estos modelos pueden rendir menos que alternativas menos restringidas a gran escala. La respuesta dominante en la literatura reciente ha sido suavizar las restricciones arquitectónicas mediante métodos de equivarianza aproximada o relajada. En contraste, el papel del optimizador en la configuración de las soluciones de los modelos equivariantes permanece poco explorado. Este artículo investiga si la elección del optimizador, independientemente de las modificaciones arquitectónicas, puede alterar significativamente las soluciones alcanzadas durante el entrenamiento y mejorar el rendimiento.

Metodología
Los autores realizan un estudio comparativo entre Adam (Kingma & Ba, 2015) y Muon (Jordan et al., 2024), un optimizador propuesto recientemente que ortogonaliza su búfer de momento mediante iteraciones de Newton-Schulz. La evaluación abarca tres configuraciones distintas:

1. Clasificación de Nubes de Puntos 3D: Utilizando el conjunto de datos ModelNet40 (y su variante corrupta, ModelNet40-C) con tres arquitecturas que poseen diferentes sesgos inductivos geométricos:
   • EGNN: Red neuronal de grafos explícitamente equivariante bajo $E(3)$.
   • DGCNN: CNN de grafos equivariante bajo permutaciones con conectividad dinámica de $k$-NN.
   • PointNet: Arquitectura invariante bajo permutaciones mediante agrupación simétrica.
2. Aprendizaje Molecular: Utilizando el conjunto de datos QM9 con GotenNet (equivariante bajo $E(3)$) y GINE (equivariante bajo permutaciones) en los conjuntos de datos de péptidos y ZINC.
3. Análisis Geométrico: Para comprender por qué difiere el rendimiento, los autores analizan los puntos de control entrenados utilizando:
   • Geometría Local de la Pérdida: Visualización de rebanadas de pérdida 2D y superficies 3D alrededor de los puntos de control, y estimación de resúmenes del Hessiano (valor propio superior y traza) mediante iteración de potencia y estimadores de Hutchinson.
   • Estructura Espectral: Cálculo del rango estable y el rango efectivo (basado en la entropía de Shannon de los valores singulares) tanto para las matrices de pesos como para las representaciones de las capas intermedias.

Resultados Clave

• Mejoras de Rendimiento: En ModelNet40, Muon supera consistentemente a Adam en todas las arquitecturas (EGNN, PointNet, DGCNN) tanto en configuraciones limpias como corruptas. Por ejemplo, en ModelNet40 limpio, Muon mejoró la precisión de EGNN en +5.17%, la de PointNet en +2.67% y la de DGCNN en +1.96%. Se observaron tendencias similares en el conjunto de datos molecular QM9, donde Muon mejoró la mayoría de las propiedades objetivo.
• Características del Paisaje de Pérdida: Las visualizaciones de la superficie de pérdida alrededor de los puntos de control entrenados revelan que las soluciones de Muon exhiben una geometría local más suave y regular en comparación con los contornos irregulares que a menudo se observan con Adam. Sin embargo, las estimaciones del Hessiano contradicen la intuición de que una geometría más suave implica una menor curvatura; los puntos de control de Muon poseen en realidad mayores estimaciones de valor propio superior y traza (indicando una mayor curvatura) que los puntos de control de Adam.
• Propiedades Espectrales: Muon produce consistentemente rangos estables y efectivos más altos tanto para los pesos como para las representaciones intermedias en todas las capas y arquitecturas. Esto sugiere que Muon distribuye la capacidad a través de las direcciones singulares de manera más uniforme, evitando la concentración de los espectros en unas pocas direcciones dominantes, a menudo asociada con métodos basados en gradientes.

Significado y Afirmaciones
El artículo postula que el optimizador es una fuente significativa de sesgo inductivo que interactúa con las restricciones geométricas. Los autores afirman que:

1. La Elección del Optimizador Importa: Cambiar el optimizador de Adam a Muon puede mejorar el rendimiento en modelos equivariantes sin modificar la arquitectura ni relajar las restricciones.
2. Espacios de Solución Distintos: Muon y Adam convergen a soluciones cualitativamente diferentes, caracterizadas por diferentes geometrías locales de pérdida y estructuras espectrales.
3. Hipótesis del Mecanismo: Las actualizaciones de momento ortogonalizadas en Muon parecen proporcionar un sesgo a nivel de trayectoria hacia regiones más favorables del paisaje de pérdida, tratando las direcciones singulares de las actualizaciones de manera más uniforme.
4. Principio de Diseño: Los resultados sugieren un posible principio de diseño para optimizadores equivariantes: las reglas de actualización que evitan que los espectros singulares de los pesos y las representaciones se concentren en pocas direcciones pueden mejorar el entrenamiento.

Los autores mantienen una postura modesta, señalando que, si bien la correlación entre las propiedades espectrales y las ganancias de precisión es consistente, no se ha establecido un vínculo causal. Concluyen que la interacción entre el diseño del optimizador y el sesgo inductivo geométrico merece una mayor investigación, particularmente en lo que respecta a cómo los optimizadores pueden explotar la estructura de los grupos de simetría en lugar de tratar la equivarianza como una restricción de caja negra.