Principio de Compatibilidad con la Simetría para el Diseño de Optimizadores: Embeddings, Cabezas de LM, MLPs SwiGLU y Enrutadores MoE

Resumen

Una disparidad geométrica llamativa ha persistido durante mucho tiempo en la práctica del aprendizaje profundo. Mientras que las arquitecturas modernas de redes neuronales exhiben naturalmente ricas propiedades de simetría y equivarianza, los optimizadores populares como Adam y sus variantes operan inherentemente de manera coordenada, lo que les impide respetar las estructuras de equivarianza del espacio de parámetros. Abordamos esta disparidad introduciendo un principio compatible con la simetría para el diseño de optimizadores: la regla de actualización del gradiente debe ser equivariante bajo el grupo de simetría que actúa sobre el bloque de pesos correspondiente. Siguiendo este principio, primero proporcionamos una perspectiva unificada sobre las actualizaciones biequivariantes ortogonales para capas de matrices generales, tal como se emplean en el descenso espectral estocástico, Muon, Scion y los métodos de gradiente polar. De manera más importante, al pasar de grupos ortogonales a simetrías de permutación y desplazamiento compartido, derivamos optimizadores compatibles con la simetría para bloques de parámetros cuyas simetrías difieren de las de las capas de matrices generales: matrices de embeddings y cabezas de modelos de lenguaje, proyecciones MLP SwiGLU y matrices de enrutadores MoE. Estas construcciones incluyen actualizaciones espectrales unilaterales, de norma de fila, híbridas norma de fila/espectrales, conscientes de fila, conscientes de columna, de norma de fila centrada y espectrales izquierdas. Proporcionan una pila de optimizadores capa por capa de extremo a extremo en la que a cada clase principal de parámetros matriciales se le asigna una actualización cuya equivarianza coincide con su grupo de simetría. Corroboramos este principio mediante experimentos de preentrenamiento en modelos de lenguaje densos y MoE dispersos, incluidas arquitecturas estilo Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B y gpt-oss reducidas. En todos estos experimentos, las actualizaciones compatibles con la simetría mejoran consistentemente la pérdida final de validación y, en varios casos, la estabilidad del entrenamiento, en comparación con las actualizaciones AdamW correspondientes.

English

A striking geometric disparity has long persisted in the practice of deep learning. While modern neural network architectures naturally exhibit rich symmetry and equivariance properties, popular optimizers such as Adam and its variants operate inherently coordinate-wise, rendering them unable to respect the equivariance structures of the parameter space. We address this disparity by introducing a symmetry-compatible principle for optimizer design: the gradient update rule should be equivariant under the symmetry group acting on the corresponding weight block. Following this principle, we first provide a unified perspective on bi-orthogonally equivariant updates for general matrix layers, as employed by stochastic spectral descent, Muon, Scion, and polar gradient methods. More importantly, by moving from orthogonal groups to permutation and shared-shift symmetries, we derive symmetry-compatible optimizers for parameter blocks whose symmetries differ from those of general matrix layers: embedding and LM head matrices, SwiGLU MLP projections, and MoE router matrices. These constructions include one-sided spectral, row-norm, hybrid row-norm/spectral, row-aware, column-aware, centered row-norm, and left-spectral updates. They yield an end-to-end layerwise optimizer stack in which each major matrix-valued parameter class is assigned an update whose equivariance matches its symmetry group. We corroborate this principle through pre-training experiments on dense and sparse MoE language models, including Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style, and downsized gpt-oss architectures. Across these experiments, symmetry-compatible updates consistently improve final validation loss, and in several cases training stability, over corresponding AdamW updates.