🤖 machine learning

Aurora: A Leverage-Aware Spectral Optimizer

O artigo apresenta o Aurora, um otimizador espectral consciente de alavancagem que impõe a uniformidade de linhas nas atualizações de parâmetros de matrizes enquanto preserva a geometria do fator polar do Muon, prevenindo assim a estagnação de neurônios e alcançando o estado da arte em desempenho, particularmente no treinamento de camadas MLP muito largas.

Autores originais: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

Publicado 2026-06-29

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está treinando uma equipe massiva de trabalhadores (uma rede neural) para resolver quebra-cabeças complexos. Nesta equipe, existem departamentos específicos chamados camadas MLP, onde o trabalho é dividido em estações individuais. Cada estação é um "neurônio" responsável por uma parte minúscula da resposta final.

O artigo apresenta um novo gerente chamado Aurora, que corrige uma falha grave na forma como o gerente anterior, Muon, estava distribuindo o trabalho.

O Problema: O Ciclo do "Trabalhador Preguiçoso"

Sob o sistema antigo (Muon), o gerente olhava para o desempenho de toda a equipe e tentava equilibrar o trabalho. No entanto, para certos tipos de tarefas (representadas matematicamente como "matrizes altas"), o Muon acidentalmente criou um ciclo vicioso:

A Divisão Desigual: Alguns trabalhadores (neurônios) recebiam atualizações enormes e empolgantes para suas habilidades, enquanto outros recebiam atualizações minúsculas, quase invisíveis.
A Espiral de Morte: Os trabalhadores que recebiam atualizações minúsculas começavam a desaparecer. Como não estavam aprendendo, paravam de contribuir para o resultado da equipe. Eventualmente, tornavam-se "neurônios mortos" — eles ainda estavam na folha de pagamento, mas não faziam nada.
O Ciclo de Feedback: À medida que esses trabalhadores paravam de contribuir, o gerente (Muon) via menos necessidade de atualizá-los, então eles recebiam atualizações ainda menores. Era um ciclo de negligência que se auto-reforçava.

O artigo mostra que, em modelos grandes treinados com Muon, um número significativo desses trabalhadores essencialmente "pediu demissão" do trabalho, especialmente nas camadas iniciais da rede.

A Solução Antiga: A Abordagem de "Força Bruta"

Alguns pesquisadores tentaram corrigir isso simplesmente forçando cada trabalhador a ter a mesma quantidade de energia (normalização de linha). Eles olhavam para a equipe, viam alguém fraco e apenas o impulsionavam.

No entanto, o artigo argumenta que essa abordagem bruta era como tentar consertar um relógio delicado com um martelo de demolição. Embora impedisse a morte dos trabalhadores, ela distorcia a forma do trabalho. Forçava a equipe a uma geometria que não se ajustava ao fluxo natural do quebra-cabeça, tornando toda a equipe menos eficiente. Era um compromisso: você salvava os trabalhadores, mas quebrava o ritmo do trabalho.

A Nova Solução: Aurora

Aurora é um otimizador que resolve este problema sem quebrar o ritmo. Pense na Aurora como uma mestre coreógrafa que entende duas regras ao mesmo tempo:

A Regra da Geometria: A equipe deve se mover em um padrão ortogonal específico e elegante (como uma dança perfeitamente sincronizada) para resolver o quebra-cabeça de forma eficiente. Este é o "fator polar" que o Muon era bom em fazer.
A Regra da Equidade: Cada trabalhador deve receber a mesma quantidade de energia e atenção.

A Aurora gerencia ambos simultaneamente. Em vez de apenas impulsionar os trabalhadores fracos a posteriori, ela calcula o movimento perfeito que satisfaz tanto o padrão de dança quanto o requisito de equidade de uma só vez.

Como funciona na prática:

Ela observa as partes "altas" da rede (as projeções de up e gate nas camadas MLP).
Executa um cálculo iterativo rápido (como algumas rodadas de cálculo mental) para encontrar a atualização perfeita que mantém os passos de dança da equipe perfeitos, garantindo que ninguém seja deixado no escuro.
Ela evita o problema do "neurônio morto" inteiramente.

Os Resultados

Os autores testaram a Aurora em grandes modelos de linguagem (340 milhões e 1,1 bilhão de parâmetros) e descobriram:

Sem Neurônios Mortos: Ao contrário do Muon, a Aurora manteve cada neurônio ativo e contribuindo.
Melhor Desempenho: Os modelos treinados com Aurora aprenderam mais rápido e obtiveram pontuações mais altas em testes de raciocínio e conhecimento (como o MMLU) do que aqueles treinados com Muon ou com os "corretores de força bruta".
Quanto Mais Larga, Melhor: O artigo descobriu que, quanto mais larga a equipe (mais neurônios nas camadas MLP), maior era a vantagem da Aurora. Isso sugere que a Aurora é a chave para treinar redes extremamente largas e poderosas sem desperdiçar recursos com trabalhadores mortos.

A Conclusão

Aurora é uma maneira mais inteligente de treinar modelos de IA. Ela corrige uma falha oculta onde os "neurônios" da IA estavam morrendo silenciosamente, garantindo que cada parte da rede seja utilizada de forma eficaz, tudo isso mantendo a estrutura matemática do processo de aprendizado perfeita. É como atualizar de um gerente que acidentalmente ignora metade do pessoal para um que garante que todos estejam dançando em perfeita sincronia.

Resumo Técnico: Aurora – Um Otimizador Espectral Consciente de Alavancagem (Leverage-Aware)

Declaração do Problema

O artigo identifica uma patologia crítica no treinamento de grandes modelos de linguagem usando o otimizador Muon, um popular otimizador espectral que aproxima as atualizações de pesos usando o fator polar da matriz de momentum do gradiente. Especificamente, para matrizes altas (onde o número de linhas $m$ excede o de colunas $n$ , como nas matrizes de projeção de up e gate em camadas MLP), as atualizações do Muon podem exibir normas de linha altamente não uniformes.

Essa não uniformidade cria um ciclo de feedback autorreforçável:

Neurônios correspondentes a linhas com normas pequenas recebem atualizações persistentemente pequenas.
Esses neurônios falham em contribuir significativamente para as saídas da rede, tornando-se efetivamente "mortos".
O artigo fornece evidência mecanística ligando este fenômeno aos scores de alavancagem (leverage scores), mostrando que baixas normas de linha no buffer de momentum correlacionam-se com baixos scores de alavancagem, os quais o Muon preserva em vez de corrigir.

Os autores observam que, embora os métodos existentes de normalização de linha (ex: NorMuon) mitiguem esse problema ao redimensionar as linhas, eles o fazem de uma forma que distorce a geometria da atualização do Muon, afastando-a do verdadeiro fator polar. Esse compromisso entre impor normas de linha uniformes e preservar a geometria precisa do fator polar é o problema central abordado.

Metodologia

1. Análise de Normalização e Precisão Polar

Os autores primeiro analisam a tensão entre precisão do fator polar e normalização de linha.

Erro de Aproximação Polar (PAE): Eles demonstram que uma maior precisão no cálculo do fator polar (via iterações Newton-Schulz como CANS-12) correlaciona-se com um menor loss de validação.
Conflito Geométrico: Eles provam que, para matrizes altas, é matematicamente impossível satisfazer simultaneamente a ortogonalidade de coluna (a propriedade do fator polar) e normas de linha unitárias. A normalização post-hoc força a atualização para longe do verdadeiro fator polar, introduzindo um erro geométrico significativo.
U-NorMuon: Como um passo intermediário, eles propõem o U-NorMuon, que aplica a normalização de linha apenas a matrizes altas e remove buffers de momentum estatais para normalização. Isso isola o benefexo da normalização de linha sem o overhead de estatísticas obsoletas, mostrando que ele supera o Muon padrão e o NorMuon.

2. O Otimizador Aurora

Para resolver o conflito entre preservar a geometria polar e impor normas de linha uniformes, os autores propõem o Aurora.

Formulação: O Aurora formula a atualização como um problema de otimização restrita que maximiza o traço do produto escalar do gradiente sujeito a duas restrições:
1. Semi-ortogonalidade à esquerda: $U^\top U = \eta^2 I$ (preservando a estrutura do fator polar).
2. Normas de linha iguais: $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$ para todas as linhas $i$ .
Riemannian-Aurora: Eles derivam um solver de referência teórica, o Riemannian-Aurora, que trata essas restrições como uma variedade conjunta (a interseção da variedade de Stiefel e a variedade de normas de linha iguais). Embora seja teoricamente preciso, é computacionalmente proibitivo devido à necessidade de resolver grandes sistemas lineares envolvendo restrições de alavancagem.
Aurora Prático: Eles propõem uma instanciação prática e iterativa que alterna entre:
1. Normalização de Linha: Redimensionar as linhas para equalizar as normas.
2. Aproximação do Fator Polar: Aplicar uma iteração Newton-Schulz para re-ortogonalizar a matriz.
  Este processo é amortecido com um parâmetro $\beta$ para garantir a convergência. Os autores provam que, sob condições de regularidade local, essa projeção alternada converge linearmente para uma solução que satisfaz ambas as restrições.

Principais Contribuições

Análise de Otimizadores Espectrais Normalizados: O artigo demonstra que a normalização de linha unitária post-hoc degrada a precisão do fator polar para matrizes altas. Introduz o U-NorMuon, uma variante sem estado que restringe a normalização a matrizes altas, provando que a uniformidade de linha e a ortogonalização precisa são simultaneamente desejáveis.
Estudo da Morte de Neurônios: Os autores fornecem uma explicação mecanística para a morte de neurônios em modelos treinados com Muon, ligando-a a baixos scores de alavancagem em parâmetros de matrizes altas. Eles verificam isso empiricamente em suas próprias execuções de treinamento e em modelos de pesos abertos (ex: Rnj-1, Laguna-XS.2), encontrando coortes significativas de neurônios mortos em camadas densas.
Aurora e Riemannian-Aurora: Eles derivam o Aurora como um método para impor normas de linha uniformes enquanto aderem estritamente ao framework teórico do Muon. Eles também fornecem o Riemannian-Aurora como um solver de referência de alto custo para o problema de dupla restrição.
Verificação Extensiva de Pré-treinamento: Os autores validam o Aurora em escalas de 340M e 1.1B de parâmetros, demonstrando desempenho superior comparado ao Muon, NorMuon e U-NorMuon.

Resultados

Desempenho em Tarefas Posteriores (Downstream)

Scores de Benchmark: Na escala de 1.1B, o Aurora supera o baseline Muon significativamente, alcançando um score 9.1 pontos maior no MMLU. Ele também alcança o maior score médio em um conjunto de benchmarks de raciocínio e conhecimento (ARC, HellaSwag, etc.) tanto na escala de 340M quanto na de 1.1B.
Speedrun: No modded-nanoGPT speedrun (trilha de otimizador), o Aurora combinado com métodos existentes (Contra-Muon, flooring) estabeleceu um novo estado da arte, atingindo o loss de validação alvo 25 passos mais rápido que o baseline NorMuon.

Utilização de Neurônios

Scores de Alavancagem: Modelos treinados com Muon exibem uma distribuição bimodal de scores de alavancagem, com um modo significativo próximo a zero (neurônios mortos). Em contraste, o Aurora e o U-NorMuon produzem distribuições de scores de alavancagem isotrópicas, eliminando efetivamente os neurônios mortos.
Verificação de Pesos Abertos: A análise de modelos de pesos abertos treinados com Muon (Rnj-1 e Laguna-XS.2) confirma a presença de neurônios mortos em camadas densas iniciais (até ~28% na primeira camada do Laguna-XS.2), validando a generalizabilidade do problema.

Eficiência e Escalonamento

Throughput: Em um ambiente distribuído, o overhead computacional adicional do Aurora é mínimo porque as multiplicações de matrizes extras podem ser sobrepostas com a comunicação de gradientes. Em uma única GPU, o throughput diminui ligeiramente, mas o overhead de tempo de execução (wall-clock) no treinamento distribuído é modesto.
Escalonamento de Largura de MLP: Os ganhos de desempenho do Aurora sobre o Muon escalam monotonicamente com o fator de expansão do MLP. À medida que as camadas MLP se tornam mais largas, a vantagem do Aurora aumenta, sugerindo que ele é particularmente eficaz para treinar redes muito largas.

Significância e Alegações

O artigo afirma que o Aurora resolve uma patologia especificamente observada empiricamente no treinamento de modelos de linguagem — a morte de neurônios em parâmetros de matrizes altas — ao derivar um otimizador que respeita tanto a geometria teórica da atualização do Muon quanto a necessidade prática de normas de linha uniformes.

Os autores enfatizam uma filosofia de design de baixo para cima (bottom-up): derivar algoritmos baseados em dinâmicas de treinamento observadas (morte de neurônios, desequilíbrios de scores de alavancagem) em vez de puramente modelos teóricos. Eles argumentam que o co-design arquitetura-otimizador, onde otimizadores são adaptados para classes específicas de parâmetros (ex: matrizes altas vs. largas), é uma direção promissora que supera estratégias padrão de "tamanho único".

Finalmente, o artigo sugere que o Aurora permite o treinamento eficaz de camadas MLP muito largas, uma capacidade que pode ser limitada no treinamento atual baseado em Muon devido ao acúmulo de neurônios mortos. O trabalho posiciona o Aurora não apenas como um impulsionador de performance, mas como uma correção necessária das suposições geométricas subjacentes à otimização espectral atual em modelos de linguagem de larga escala.