Resumo Técnico: Princípio Compatível com Simetria para Projeto de Otimizadores

1. Declaração do Problema

Existe uma disparidade geométrica fundamental na prática moderna de aprendizado profundo. Enquanto arquiteturas de redes neurais (embeddings, cabeças de LM, MLPs SwiGLU, roteadores MoE) possuem ricas propriedades de simetria e equivariância intrínsecas, os métodos de otimização dominantes, como Adam, AdamW e RMSprop, operam coordenada a coordenada. Esses métodos tratam parâmetros matriciais como vetores achatados de coordenadas independentes, ignorando assim a estrutura matricial e quebrando a equivariância natural do espaço de parâmetros.

Essa abordagem "cega à geometria" falha em distinguir entre as geometrias distintas de diferentes tipos de camadas. Consequentemente, otimizadores populares podem descartar estrutura espectral, induzir geometrias incompatíveis e falhar em respeitar as estruturas de invariância e equivariância inerentes a blocos de parâmetros específicos (por exemplo, invariância de permutação em vocabulários ou intercambiabilidade de especialistas em MoEs). A compreensão teórica atual e as comparações empíricas entre métodos adaptativos coordenada a coordenada e otimizadores espectrais permanecem limitadas, frequentemente dependendo de benchmarks em grande escala sem análise profunda do comportamento geométrico.

2. Metodologia: O Princípio Compatível com Simetria

O artigo propõe um princípio compatível com simetria por camada para o projeto de otimizadores. A tese central é que a regra de atualização $\mathcal{U}$ do otimizador para um bloco de parâmetros $W$ deve ser equivariante sob o grupo de simetria $G$ atuando naquele bloco específico.

2.1 Estrutura Teórica

Seja $W \in \mathbb{R}^{m \times n}$ uma matriz de parâmetros. Se as coordenadas dos espaços de entrada e saída forem transformadas por matrizes invertíveis $P$ e $Q$, o mapa linear se transforma como $\tilde{W} = P W Q^{-1}$. O gradiente se transforma contravariantemente/covariantemente de acordo. O mapa de atualização do otimizador $\mathcal{U}$ deve satisfazer:
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
onde $D$ é uma direção de atualização (por exemplo, gradiente ou momento) e $g$ é um elemento do grupo de simetria.

O artigo deriva classes específicas de otimizadores com base nos grupos de simetria de diferentes tipos de camadas:

1. Camadas Matriciais Comuns (Linear/Attention):

   • Simetria: Bi-ortogonal ($O_m \times O_n$).
   • Classe de Otimizador: Otimizadores Espectrais. A atualização deve ser um operador espectral atuando sobre os valores singulares, preservando os vetores singulares. Isso unifica a Descida Espectral Estocástica (SSD), Muon, Scion e métodos de Gradiente Polar. A atualização é tipicamente o fator polar ortogonal da direção do momento.

2. Embeddings e Cabeças de LM:

   • Simetria: Permutação à esquerda ($P_v$) e Ortogonal à direita ($O_d$). Linhas indexam itens do vocabulário (permutação discreta), colunas indexam características ocultas (rotação contínua).
   • Classe de Otimizador: Atualizações Equivariantes Permutação à Esquerda Ortogonal à Direita (LPRO).
     • Espectral à Direita: Atua globalmente via a matriz Gram à direita ($D^\top D$).
     • Norma de Linha: Atua localmente sobre as normas de linha (itens do vocabulário).
     • Híbrido: Composições de atualizações de norma de linha e espectral à direita.

3. Projeções MLP SwiGLU:

   • Simetria: Permutação de neurônios intermediários ($P_{d_{ff}}$). As projeções de porta e subida têm linhas correspondentes aos neurônios intermediários; a projeção de descida tem colunas correspondentes a eles.
   • Classe de Otimizador: Variantes conscientes de linha e coluna. Projeções de porta/subida usam atualizações compatíveis com LPRO (norma de linha/espectral à direita); projeções de descida usam a geometria transposta.

4. Roteadores MoE:

   • Simetria: Permutação de especialistas ($P_e$) e invariância de deslocamento de logit compartilhado (adicionar uma constante a todos os logits não altera a distribuição softmax).
   • Classe de Otimizador: Atualizações compatíveis com roteador. Estas operam sobre a direção centralizada $D_c = \Pi^\perp D$ (removendo o deslocamento compartilhado).
     • Espectral à Esquerda Centralizada: Atua sobre a matriz Gram de especialistas centralizada.
     • Norma de Linha Centralizada: Atua sobre as normas de linha de especialistas centralizadas.

2.2 Implementação Prática

Os autores instanciam esses princípios em otimizadores práticos baseados em momento:

• RightPolarGradM / LeftPolarGradM: Atualizações espectrais de um lado usando iterações de Newton-Schulz para raízes quadradas inversas de matrizes.
• RowNormM: Atualizações de normalização linha a linha.
• HybridPolarGradM: Composições de passos de norma de linha e espectrais de um lado.
• Centralização: Para roteadores MoE, as atualizações são aplicadas ao momento centralizado $M_{k,c} = \Pi^\perp M_k$.

3. Contribuições Principais

1. Princípio Compatível com Simetria: Um argumento formal de que otimizadores coordenada a coordenada são geometricamente incompatíveis para parâmetros matriciais. O artigo estabelece que as regras de atualização devem corresponder ao grupo de simetria do bloco de parâmetros.
2. Visão Unificadora de Otimizadores Espectrais: Fornece uma interpretação baseada em simetria de otimizadores espectrais (Muon, SSD, etc.) como a classe única de mapas de atualização bi-ortogonalmente equivariantes.
3. Novas Classes de Otimizadores Equivariantes: Deriva regras de atualização específicas para camadas com simetrias não padrão:
   • Atualizações equivariantes LPRO para embeddings e cabeças de LM.
   • Atualizações conscientes de linha/coluna para MLPs SwiGLU.
   • Atualizações de permutação de especialistas centralizadas para roteadores MoE.
4. Evidência de Pré-treinamento End-to-End: Constrói uma pilha completa de otimizadores por camada onde cada classe matricial principal recebe uma atualização compatível com simetria. Isso é avaliado em modelos de linguagem MoE densos e esparsos (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, gpt-oss reduzido).

4. Resultados Experimentais

Os autores conduziram experimentos de pré-treinamento em subconjuntos de 10B de tokens do FineWeb-Edu em quatro arquiteturas.

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

  • Substituir o AdamW em matrizes indexadas por vocabulário (embeddings/cabeças de LM) por RowNormM ou HybridPolarGradM melhorou consistentemente a perda de validação final.
  • Os ganhos foram modestos para o modelo Qwen3 menor, mas mais pronunciados para o modelo Gemma 3 maior, sugerindo que atualizações conscientes de geometria tornam-se mais críticas à medida que as matrizes de vocabulário crescem e os gradientes se tornam mais anisotrópicos.
  • O uso de HybridPolarGradM para projeções MLP SwiGLU melhorou ainda mais o desempenho, provavelmente devido à correção do desequilíbrio de escala de linha em matrizes altas e finas.

• Modelos MoE Esparsos (OLMoE-1B-7B, gpt-oss reduzido):

  • Configurações usando atualizações compatíveis com simetria para embeddings, cabeças de LM e roteadores (por exemplo, RowNormM + LeftPolarGradM) superaram configurações usando AdamW para esses componentes.
  • A linha de base totalmente coordenada a coordenada (AdamW em todos os lugares) teve o pior desempenho.
  • Atualizações de roteador compatíveis com simetria melhoraram em relação às atualizações de roteador coordenada a coordenada e, em alguns casos, reduziram picos de perda de treinamento, sugerindo maior estabilidade.

• Tendência Geral: Em todos os modelos, substituir o AdamW em grandes matrizes indexadas por vocabulário por otimizadores compatíveis com simetria reduziu consistentemente a perda de validação final. Os benefícios persistiram em diferentes arquiteturas e níveis de esparsidade.

5. Significado e Alegações

O artigo alega oferecer uma mudança principista no projeto de otimizadores, de uma visão coordenada a coordenada "cega à geometria" para uma visão matricial "consciente de simetria".

• Co-Design Arquitetura-Otimizador: O trabalho argumenta que o projeto de otimizadores deve ser um problema consciente de arquitetura. Diferentes classes de parâmetros (embeddings, roteadores, camadas lineares) possuem simetrias naturais distintas e requerem geometrias de otimizadores distintas.
• Fundamentação Teórica: Fornece uma caracterização matemática rigorosa (via equivariância) do porquê otimizadores espectrais funcionam para camadas lineares e estende essa lógica para derivar novos otimizadores para outras camadas.
• Utilidade Prática: Os resultados demonstram que atribuições compatíveis com simetria não são apenas teoricamente sólidas, mas praticamente benéficas, melhorando a perda final e potencialmente a estabilidade de treinamento em pré-treinamento em grande escala.
• Modéstia: Os autores afirmam explicitamente que não alegam que otimizadores equivariantes dominam métodos coordenada a coordenada em todos os regimes. Em vez disso, apresentam evidências de que o princípio compatível com simetria por camada leva a atribuições competitivas e frequentemente benéficas em configurações representativas de pré-treinamento. Eles veem os resultados como evidência da utilidade do princípio, e não como um benchmark exaustivo.

O artigo conclui que introduzir simetria e equivariância no projeto de otimizadores abre um caminho para pré-treinamento de modelos em grande escala mais estável, consistente com a parametrização e fundamentado teoricamente.