Princípio Compatível com Simetria para Projeto de Otimizadores: Embeddings, Cabeças de LM, MLPs SwiGLU e Roteadores MoE

Resumo

Uma notável disparidade geométrica persiste há muito tempo na prática do aprendizado profundo. Enquanto arquiteturas modernas de redes neurais naturalmente exibem ricas propriedades de simetria e equivariância, otimizadores populares como Adam e suas variantes operam inerentemente de forma coordenada a coordenada, tornando-os incapazes de respeitar as estruturas de equivariância do espaço de parâmetros. Abordamos essa disparidade introduzindo um princípio compatível com simetria para o projeto de otimizadores: a regra de atualização do gradiente deve ser equivariante sob o grupo de simetria que atua no bloco de pesos correspondente. Seguindo esse princípio, fornecemos primeiro uma perspectiva unificada sobre atualizações bi-ortogonalmente equivariantes para camadas de matrizes gerais, conforme empregado pelo descida espectral estocástica, Muon, Scion e métodos de gradiente polar. Mais importante, ao passar de grupos ortogonais para simetrias de permutação e deslocamento compartilhado, derivamos otimizadores compatíveis com simetria para blocos de parâmetros cujas simetrias diferem daquelas de camadas de matrizes gerais: matrizes de embedding e cabeça LM, projeções MLP SwiGLU e matrizes de roteador MoE. Essas construções incluem atualizações espectrais unilaterais, de norma de linha, híbridas norma de linha/espectral, conscientes de linha, conscientes de coluna, de norma de linha centralizada e espectrais à esquerda. Elas resultam em uma pilha de otimizadores por camada de ponta a ponta na qual a cada classe principal de parâmetro matricial é atribuída uma atualização cuja equivariância corresponde ao seu grupo de simetria. Corroboramos esse princípio por meio de experimentos de pré-treinamento em modelos de linguagem MoE densos e esparsos, incluindo arquiteturas estilo Qwen3-0.6B, estilo Gemma 3 1B, estilo OLMoE-1B-7B e gpt-oss reduzidas. Nesses experimentos, as atualizações compatíveis com simetria melhoram consistentemente a perda de validação final e, em vários casos, a estabilidade do treinamento, em comparação com as atualizações AdamW correspondentes.

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.