Resumen Técnico: Principio Compatible con Simetría para el Diseño de Optimizadores

1. Declaración del Problema

Existe una disparidad geométrica fundamental en la práctica actual del aprendizaje profundo. Mientras que las arquitecturas de redes neuronales (incrustaciones, cabezas de LM, MLPs SwiGLU, enrutadores MoE) poseen ricas propiedades de simetría intrínseca y equivarianza, los métodos de optimización dominantes como Adam, AdamW y RMSprop operan coordenada a coordenada. Estos métodos tratan los parámetros matriciales como vectores aplanados de coordenadas independientes, ignorando así la estructura matricial y rompiendo la equivarianza natural del espacio de parámetros.

Este enfoque "ciego a la geometría" falla al distinguir entre las geometrías distintas de diferentes tipos de capas. En consecuencia, los optimizadores populares pueden descartar la estructura espectral, inducir geometrías desajustadas y no respetar las estructuras de invariancia y equivarianza inherentes a bloques de parámetros específicos (por ejemplo, invariancia de permutación en vocabularios o intercambiabilidad de expertos en MoEs). La comprensión teórica actual y las comparaciones empíricas entre métodos adaptativos coordenada a coordenada y optimizadores espectrales siguen siendo limitadas, a menudo dependiendo de benchmarks a gran escala sin un análisis profundo del comportamiento geométrico.

2. Metodología: El Principio Compatible con Simetría

El artículo propone un principio compatible con simetría a nivel de capa para el diseño de optimizadores. La tesis central es que la regla de actualización $\mathcal{U}$ del optimizador para un bloque de parámetros $W$ debe ser equivariante bajo el grupo de simetría $G$ que actúa sobre ese bloque específico.

2.1 Marco Teórico

Sea $W \in \mathbb{R}^{m \times n}$ una matriz de parámetros. Si las coordenadas de los espacios de entrada y salida se transforman mediante matrices invertibles $P$ y $Q$, la aplicación lineal se transforma como $\tilde{W} = P W Q^{-1}$. El gradiente se transforma de manera contravariante/covariante en consecuencia. El mapa de actualización del optimizador $\mathcal{U}$ debe satisfacer:
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
donde $D$ es una dirección de actualización (por ejemplo, gradiente o momento) y $g$ es un elemento del grupo de simetría.

El artículo deriva clases específicas de optimizadores basadas en los grupos de simetría de diferentes tipos de capas:

1. Capas Matriciales Ordinarias (Lineales/Atención):

   • Simetría: Bi-ortogonal ($O_m \times O_n$).
   • Clase de Optimizador: Optimizadores Espectrales. La actualización debe ser un operador espectral que actúe sobre los valores singulares preservando los vectores singulares. Esto unifica el Descenso Espectral Estocástico (SSD), Muon, Scion y los métodos de Gradiente Polar. La actualización es típicamente el factor polar ortogonal de la dirección del momento.

2. Incrustaciones y Cabezas de LM:

   • Simetría: Permutación izquierda ($P_v$) y Ortogonal derecha ($O_d$). Las filas indexan elementos del vocabulario (permutación discreta), las columnas indexan características ocultas (rotación continua).
   • Clase de Optimizador: Actualizaciones Equivariantes Permutación-Izquierda Ortogonal-Derecha (LPRO).
     • Espectral Derecha: Actúa globalmente a través de la matriz Gram derecha ($D^\top D$).
     • Norma de Fila: Actúa localmente sobre las normas de las filas (elementos del vocabulario).
     • Híbrido: Composiciones de actualizaciones de norma de fila y espectrales derechas.

3. Proyecciones MLP SwiGLU:

   • Simetría: Permutación de neuronas intermedias ($P_{d_{ff}}$). Las proyecciones de puerta (gate) y subida (up) tienen filas correspondientes a neuronas intermedias; la proyección de bajada (down) tiene columnas correspondientes a ellas.
   • Clase de Optimizador: Variantes conscientes de fila y columna. Las proyecciones de puerta/subida utilizan actualizaciones compatibles con LPRO (norma de fila/espectral derecha); las proyecciones de bajada utilizan la geometría transpuesta.

4. Enrutadores MoE:

   • Simetría: Permutación de expertos ($P_e$) e invariancia de desplazamiento de logit compartido (añadir una constante a todos los logits no cambia la distribución softmax).
   • Clase de Optimizador: Actualizaciones compatibles con enrutadores. Estas operan sobre la dirección centrada $D_c = \Pi^\perp D$ (eliminando el desplazamiento compartido).
     • Espectral Izquierda Centrada: Actúa sobre la matriz Gram de expertos centrada.
     • Norma de Fila Centrada: Actúa sobre las normas de las filas de expertos centradas.

2.2 Implementación Práctica

Los autores materializan estos principios en optimizadores prácticos basados en momento:

• RightPolarGradM / LeftPolarGradM: Actualizaciones espectrales unilaterales utilizando iteraciones de Newton-Schulz para raíces cuadradas inversas de matrices.
• RowNormM: Actualizaciones de normalización fila a fila.
• HybridPolarGradM: Composiciones de pasos de norma de fila y espectrales unilaterales.
• Centrado: Para enrutadores MoE, las actualizaciones se aplican al momento centrado $M_{k,c} = \Pi^\perp M_k$.

3. Contribuciones Clave

1. Principio Compatible con Simetría: Un argumento formal de que los optimizadores coordenada a coordenada son geométricamente desajustados para parámetros matriciales. El artículo establece que las reglas de actualización deben coincidir con el grupo de simetría del bloque de parámetros.
2. Visión Unificadora de Optimizadores Espectrales: Proporciona una interpretación basada en simetría de los optimizadores espectrales (Muon, SSD, etc.) como la clase única de mapas de actualización bi-ortogonalmente equivariantes.
3. Nuevas Clases de Optimizadores Equivariantes: Deriva reglas de actualización específicas para capas con simetrías no estándar:
   • Actualizaciones equivariantes LPRO para incrustaciones y cabezas de LM.
   • Actualizaciones conscientes de fila/columna para MLPs SwiGLU.
   • Actualizaciones de permutación de expertos centradas para enrutadores MoE.
4. Evidencia de Pre-entrenamiento de Extremo a Extremo: Construye una pila completa de optimizadores a nivel de capa donde cada clase principal de matrices recibe una actualización coincidente con su simetría. Esto se evalúa en modelos de lenguaje densos y MoE dispersos (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, gpt-oss reducido).

4. Resultados Experimentales

Los autores realizaron experimentos de pre-entrenamiento en subconjuntos de 10 mil millones de tokens de FineWeb-Edu a través de cuatro arquitecturas.

• Modelos Densos (Qwen3-0.6B, Gemma 3 1B):

  • Reemplazar AdamW en matrices indexadas por vocabulario (incrustaciones/cabezas de LM) con RowNormM o HybridPolarGradM mejoró consistentemente la pérdida de validación final.
  • Las ganancias fueron modestas para el modelo Qwen3 más pequeño pero más pronunciadas para el modelo Gemma 3 más grande, lo que sugiere que las actualizaciones conscientes de la geometría se vuelven más críticas a medida que las matrices de vocabulario crecen y los gradientes se vuelven más anisotrópicos.
  • Utilizar HybridPolarGradM para las proyecciones MLP SwiGLU mejoró aún más el rendimiento, probablemente debido a la corrección del desequilibrio de escala de fila en matrices altas y delgadas.

• Modelos MoE Dispersos (OLMoE-1B-7B, gpt-oss reducido):

  • Las configuraciones que utilizan actualizaciones compatibles con simetría para incrustaciones, cabezas de LM y enrutadores (por ejemplo, RowNormM + LeftPolarGradM) superaron a las configuraciones que utilizan AdamW para estos componentes.
  • La línea base totalmente coordenada a coordenada (AdamW en todas partes) tuvo el peor rendimiento.
  • Las actualizaciones de enrutadores compatibles con simetría mejoraron sobre las actualizaciones de enrutadores coordenada a coordenada y, en algunos casos, redujeron los picos de pérdida de entrenamiento, sugiriendo una mayor estabilidad.

• Tendencia General: En todos los modelos, reemplazar AdamW en matrices grandes indexadas por vocabulario con optimizadores compatibles con simetría redujo consistentemente la pérdida de validación final. Los beneficios persistieron a través de diferentes arquitecturas y niveles de dispersión.

5. Significado y Afirmaciones

El artículo afirma ofrecer un cambio principista en el diseño de optimizadores, pasando de una visión coordenada a coordenada "ciega a la geometría" a una visión matricial "consciente de simetría".

• Co-diseño Arquitectura-Optimizador: El trabajo argumenta que el diseño de optimizadores debe ser un problema consciente de la arquitectura. Diferentes clases de parámetros (incrustaciones, enrutadores, capas lineales) tienen simetrías naturales distintas y requieren geometrías de optimizador distintas.
• Fundamentación Teórica: Proporciona una caracterización matemática rigurosa (mediante equivarianza) de por qué los optimizadores espectrales funcionan para capas lineales y extiende esta lógica para derivar nuevos optimizadores para otras capas.
• Utilidad Práctica: Los resultados demuestran que las asignaciones compatibles con simetría no solo son teóricamente sólidas, sino también beneficiosas en la práctica, mejorando la pérdida final y potencialmente la estabilidad del entrenamiento en pre-entrenamiento a gran escala.
• Modestia: Los autores declaran explícitamente que no afirman que los optimizadores equivariantes dominen a los métodos coordenada a coordenada en todos los regímenes. En cambio, presentan evidencia de que el principio compatible con simetría a nivel de capa conduce a asignaciones competitivas y a menudo beneficiosas en configuraciones de pre-entrenamiento representativas. Consideran los resultados como evidencia de la utilidad del principio en lugar de un benchmark exhaustivo.

El artículo concluye que introducir simetría y equivarianza en el diseño de optimizadores abre un camino hacia un pre-entrenamiento de modelos a gran escala más estable, consistente con la parametrización y fundamentado teóricamente.