Tilde Research Apresenta Aurora: Um Otimizador Sensível à Alavancagem Que Corrige Um Problema Oculto de Morte de Neurônios no Muon

Pesquisadores da Tilde Research lançaram Aurora, um novo otimizador para o treinamento de redes neurais que aborda uma falha estrutural no amplamente utilizado otimizador Muon. A falha discretamente mata uma fração significativa de neurônios MLP durante o treinamento e os mantém permanentemente inativos. Aurora vem com um experimento de pré-treinamento de 1.1B parâmetros, um novo resultado de última geração no benchmark de speedrun modded-nanoGPT e códigos abertos. O que é Muon? Para entender Aurora, ajuda primeiro entender Muon. O otimizador Muon atraiu a atenção da comunidade de ML depois de superar o AdamW em tempo de clock para convergência na competição de speedrun nanoGPT – um benchmark da comunidade que mede a rapidez com que se pode treinar um modelo estilo GPT para uma perda de validação alvo. Desde então, Muon foi adotado no treinamento de modelos em escala de fronteira por vários grupos de pesquisa. A etapa algorítmica chave do Muon é o cálculo do fator polar da matriz gradiente. Para uma matriz gradiente G com Decomposição de Valor Singular (SVD) fina G = UΣVᵀ, Muon calcula polar(G) = UVᵀ, que é a matriz semi-ortogonal mais próxima de G na norma de Frobenius. Este gradiente ortogonalizado é então usado para atualizar os pesos: W ← W − η UVᵀ para uma taxa de aprendizado η. O uso de algoritmos iterativos apenas de multiplicação de matrizes para calcular o fator polar é o que torna Muon prático em escala. O Quebra-Cabeça NorMuon: Normalização de Linhas Ajuda, Mas Por Que? Antes de Aurora, NorMuon liderava o speedrun modded-nanoGPT. Ele introduziu um passo de normalização de linha – semelhante à escala por parâmetro de Adam – que ajustava o fator polar pela sua norma RMS inversa. Embora isso frequentemente afaste a atualização de um gradiente estritamente ortogonal, NorMuon ainda produz resultados impressionantes. A equipe da Tilde se propôs a entender exatamente qual lacuna na formulação de Muon NorMuon estava abordando. O Problema Central: Anisotropia de Norma de Linha e Morte de Neurônios em Matrizes Altas A equipe de pesquisa descobriu que o otimizador Muon, sem intenção, “mata” uma grande parte dos neurônios em matrizes de pesos altas, como as encontradas em camadas MLP baseadas em SwiGLU. Como é matematicamente impossível para estas formas de matriz específicas permanecerem perfeitamente ortogonais enquanto mantêm as atualizações de linha equilibradas, o otimizador acaba dando atualizações massivas a alguns neurônios enquanto virtualmente ignora outros. Isso resulta em uma “espiral da morte” onde os neurônios com desempenho inferior recebem menos sinal ao longo do tempo, eventualmente tornando-se permanentemente inativos. O estudo de pesquisa revelou que até o 500º treinamento