Muon+: Towards Better Muon via One Additional Normalization Step

Este trabalho apresenta o Muon+, uma melhoria simples e eficaz para o otimizador Muon que introduz uma etapa adicional de normalização após a ortogonalização, demonstrando consistentemente melhorias no desempenho de pré-treinamento de modelos de linguagem de grande escala.

O essencial

Imagine que você está tentando ensinar um grupo de gigantes (os modelos de Inteligência Artificial) a escrever histórias incríveis. Para fazer isso, você precisa dar a eles "lições" (dados) e corrigir os erros deles. O processo de corrigir esses erros é chamado de otimização.

Até agora, a ferramenta mais famosa para fazer essa correção era o Adam (como um professor rigoroso que olha para cada erro individualmente). Mas, recentemente, surgiu um novo professor chamado Muon. O Muon é especial porque ele não olha apenas para os erros individuais; ele olha para o padrão de todos os erros juntos, organizando-os de forma que não se repitam (isso é chamado de "ortogonalização"). Isso funciona muito bem, mas os autores deste artigo perceberam que ainda havia um pequeno problema: às vezes, a "força" dessas correções ficava desequilibrada.

Aqui entra o MUON+, a nova e melhorada versão proposta neste artigo.

A Analogia do "Equilíbrio Perfeito"

Para entender o MUON+, vamos usar uma analogia de uma orquestra:

1. O Problema Antigo (Muon): Imagine que o Muon é um maestro que organiza a orquestra para que cada músico toque uma nota diferente, evitando que todos toquem a mesma coisa ao mesmo tempo (isso é a "ortogonalização"). Isso já é ótimo e evita o caos.
2. O Novo Passo (MUON+): No entanto, o maestro percebeu que, mesmo com notas diferentes, alguns músicos estavam tocando muito alto e outros muito baixo. O som final ficava desequilibrado.
3. A Solução: O MUON+ adiciona um passo extra simples: ele pede para todos os músicos ajustarem o volume para que fiquem no mesmo nível de intensidade, sem mudar a nota que estão tocando. Isso é a normalização.

Em termos técnicos, o MUON+ pega a correção que o Muon já fez e "padroniza" o tamanho dela antes de aplicá-la ao modelo. É como se, antes de dar o próximo passo na estrada, você verificasse se seus sapatos estão amarrados e do tamanho certo, garantindo que você não tropece.

O Que Eles Descobriram?

Os pesquisadores testaram essa ideia em vários tamanhos de modelos de IA (desde modelos pequenos de 130 milhões de "cérebros" até gigantes de 1 bilhão). Eles usaram dois tipos de arquitetura populares: GPT (como o ChatGPT antigo) e LLaMA (o modelo da Meta).

Os resultados foram impressionantes:

• Melhoria Consistente: Em todos os testes, o MUON+ aprendeu mais rápido e ficou mais inteligente (teve menos "perplexidade", que é uma medida de quão confuso o modelo fica ao prever a próxima palavra) do que o Muon original.
• Funciona em Longas Distâncias: Eles treinaram os modelos com uma quantidade gigantesca de dados (como se lessem toda a internet várias vezes). Mesmo com tanta informação, o MUON+ manteve a estabilidade, enquanto outros métodos começariam a oscilar ou perder eficiência.
• Simplicidade é Poder: A grande sacada é que eles não precisaram criar uma fórmula matemática complexa. Eles apenas adicionaram um passo de "ajuste de volume" (normalização) depois que o Muon já organizou a música.

Por Que Isso Importa?

Treinar modelos de IA hoje em dia custa milhões de dólares em energia e tempo de computador. Qualquer melhoria que faça o modelo aprender melhor com a mesma quantidade de dados é um grande avanço.

O MUON+ mostra que, às vezes, a solução para problemas complexos de IA não é criar algo mais complicado, mas sim garantir que o que já funciona esteja equilibrado. É como dizer: "Não adianta ter um carro de Fórmula 1 se os pneus estiverem descalibrados; o MUON+ apenas calibra os pneus."

Resumo em uma Frase

O MUON+ é uma melhoria simples e brilhante no otimizador Muon que, ao apenas "nivelar o campo" das correções matemáticas, faz os modelos de Inteligência Artificial aprenderem mais rápido, mais estáveis e com melhor qualidade, seja eles pequenos ou gigantes.

Resumo técnico

1. Problema

O treinamento de modelos de linguagem fundamentais (Foundation Models) em grande escala enfrenta desafios significativos de custo computacional e estabilidade de otimização. Otimizadores tradicionais como Adam e AdamW dominam o campo, mas recentemente o otimizador Muon ganhou destaque por sua capacidade de lidar com o "colapso de rank" do gradiente através da ortogonalização da matriz de momento.

No entanto, apesar do sucesso do Muon, os autores observam que ainda há espaço para melhorias na estabilidade da otimização e na qualidade final do modelo, especialmente em regimes de treinamento computacionalmente ótimos e em longas durações (overtraining). A questão central é: como aprimorar o Muon de forma simples e eficaz para extrair ainda mais desempenho sem aumentar drasticamente a complexidade computacional?

2. Metodologia: MUON+

O artigo propõe o MUON+, uma melhoria direta e simples ao otimizador Muon. A inovação principal é a introdução de um passo adicional de normalização aplicado imediatamente após a etapa de ortogonalização.

• Mecanismo de Atualização:

  1. Momento: Calcula-se o momento padrão ($M_t = \mu M_{t-1} + (1-\mu)G_t$).
  2. Ortogonalização: Aplica-se a função de ortogonalização (geralmente via iterações de Newton-Schulz para aproximar a decomposição SVD) para obter a matriz ortogonal $O_t$.
  3. Normalização (Novo Passo): Aplica-se um operador de normalização $\text{Norm}(d)$ na matriz ortogonalizada antes da atualização dos pesos.
  4. Atualização: Os pesos são atualizados usando a matriz normalizada.

• Direções de Normalização: Os autores investigam diferentes eixos de normalização $\ell_2$:

  • Coluna ($col$): Normaliza cada coluna da matriz.
  • Linha ($row$): Normaliza cada linha da matriz.
  • Combinadas ($col\_row$ e $row\_col$): Aplica-se normalização sequencialmente em ambas as direções.

• Abordagem Experimental: O estudo foi conduzido em arquiteturas GPT e LLaMA, variando de 60M a 1 bilhão de parâmetros. Os experimentos cobriram dois regimes:

  1. Treinamento Computacionalmente Ótimo: Relação Token-para-Parâmetro (T2P) de ~20.
  2. Overtraining (Longa Duração): Relação T2P estendida para ~200 (72 bilhões de tokens).

3. Principais Contribuições

• Simplicidade e Eficácia: Demonstra que uma única etapa de normalização, adicionada após a ortogonalização, é suficiente para gerar ganhos consistentes de desempenho, sem a necessidade de mecanismos complexos de adaptação de segunda ordem ou formulações de variedades (manifolds) complexas.
• Análise de Normalização: O trabalho isola o efeito da normalização, mostrando que ela é o principal motor de melhoria de desempenho, superando a contribuição de outros componentes adicionais encontrados em variantes anteriores do Muon (como o NorMuon).
• Escalabilidade: Valida que o MUON+ funciona eficazmente em uma ampla gama de escalas de modelos (de 60M a 1B) e arquiteturas (GPT e LLaMA).
• Robustez à Taxa de Aprendizado: O MUON+ demonstra ser menos sensível à escolha da taxa de aprendizado (learning rate), mantendo desempenho estável em uma faixa mais ampla de hiperparâmetros em comparação ao Muon padrão.

4. Resultados

Os resultados experimentais mostram melhorias consistentes na perplexidade de validação (PPL) em todos os cenários testados:

• Modelos GPT:
  • GPT-Small (124M): Redução de PPL de 29.66 (Muon) para 27.64 (MUON+), uma melhoria de -2.02.
  • GPT-Base (362M): Redução de 21.70 para 19.98 (-1.72).
  • GPT-Large (774M): Redução de 17.82 para 16.91 (-0.91).
• Modelos LLaMA:
  • Em todas as escalas (60M a 1B), o MUON+ superou o AdamW e o Muon padrão. Por exemplo, no LLaMA-1B, a PPL caiu de 10.68 (Muon) para 10.31 (MUON+).
• Overtraining (T2P ~200):
  • Mesmo com 72 bilhões de tokens de treinamento, o MUON+ manteve a vantagem. No GPT-Base, a PPL foi de 15.84 vs 16.97 do Muon.
  • As curvas de perda de treinamento mostraram que o MUON+ escala favoravelmente e não degrada em estágios tardios de otimização.
• Ablação de Direções: A normalização combinada (especialmente col_row e row_col) tendeu a oferecer o melhor desempenho, embora a normalização apenas por linha (row) tenha superado consistentemente a apenas por coluna (col).

5. Significado e Impacto

O MUON+ representa um avanço significativo na otimização de grandes modelos de linguagem por várias razões:

1. Custo-Benefício: A melhoria é obtida com um custo computacional marginal (apenas uma operação de normalização extra), tornando-a altamente atrativa para a indústria.
2. Estabilidade em Longa Duração: A capacidade de manter a estabilidade e a melhoria de desempenho em regimes de overtraining (T2P alto) é crucial para o treinamento de modelos de próxima geração, onde os dados são escassos e o treinamento é prolongado.
3. Insight Teórico: O trabalho sugere que a normalização estrutural das atualizações ortogonais é um fator chave, talvez mais importante do que adaptações complexas de variância (segunda ordem), redefinindo o foco para estratégias de normalização mais simples e eficazes.
4. Adoção Prática: Dado que o Muon já é utilizado em modelos de ponta (como Kimi e GLM), o MUON+ oferece um caminho imediato para melhorar o estado da arte (SOTA) em eficiência e qualidade de modelos fundamentais.

Em resumo, o MUON+ valida que "menos é mais" na otimização de LLMs, demonstrando que um passo de normalização bem posicionado pode desbloquear ganhos substanciais de desempenho e estabilidade.