Resumo Técnico: AMUSE (Muon a Qualquer Momento com Avaliação de Gradiente Estável)

Declaração do Problema

A otimização moderna de aprendizado profundo geralmente depende do otimizador AdamW combinado com cronogramas de taxa de aprendizado prescritos (por exemplo, decaimento cosseno). Avanços recentes desafiaram esses componentes individualmente: a otimização Schedule-Free (SF) remove cronogramas explícitos avaliando gradientes em pontos interpolados entre as iterações atuais e as médias, enquanto o Muon melhora a geometria de atualização para parâmetros matriciais ao ortogonalizar o momento.

Apesar do sucesso empírico do Muon no pré-treinamento de modelos de linguagem grandes (LLM), seu mecanismo subjacente não é totalmente compreendido. O artigo identifica uma instabilidade específica no Muon: enquanto sua ortogonalização acelera o progresso ao longo de direções "planas" (o "rio" da paisagem de perda), ela amplifica inadvertidamente o ruído em direções "íngremes" de alta curvatura (as "paredes do vale"), levando a trajetórias oscilatórias. Por outro lado, embora os métodos Schedule-Free estabilizem as trajetórias mantendo-as próximas ao rio, eles podem carecer do progresso rápido em massa oferecido pela ortogonalização do Muon. O problema central é projetar um otimizador que combine a convergência rápida do Muon com a estabilidade da média Schedule-Free, sem exigir cronogramas manuais de taxa de aprendizado.

Metodologia: AMUSE

Os autores propõem o AMUSE (Muon a Qualquer Momento com Avaliação de Gradiente Estável), que integra as atualizações matriciais ortogonalizadas do Muon com o framework Schedule-Free usando um coeficiente de interpolação variante no tempo.

Fundamento Teórico: Paisagem de Perda Rio-Vale

A análise fundamenta-se na perspectiva da paisagem de perda rio-vale:

• O Rio: Um subespaço volumétrico de alta dimensão e baixa curvatura onde ocorre o progresso útil do treinamento.
• As Paredes do Vale: Um subespaço dominante de baixa dimensão e alta curvatura que induz oscilações.
• Mecanismo do Muon: A ortogonalização da matriz de momento aumenta o componente volumétrico relativo das atualizações, acelerando o progresso ao longo do rio. No entanto, ela também impede que os componentes dominantes decaiam, fazendo com que pequenos ruídos em direções dominantes se amplifiquem em oscilações através das paredes do vale.
• Mecanismo Schedule-Free: Avaliar gradientes em uma trajetória média (mais próxima do rio) filtra flutuações de alta curvatura antes da etapa de atualização.

Design do Algoritmo

O AMUSE emprega um coeficiente de interpolação variante no tempo $\beta_t$ para determinar o ponto de avaliação do gradiente $Y_t$ entre a sequência base rápida $Z_t$ (atualizações do Muon) e a sequência média estável $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

O cronograma para $\beta_t$ é projetado para equilibrar a adaptação inicial e a estabilidade em estágios tardios:

1. Fase de Aquecimento ($t \le T_0$): $\beta_t$ é fixado em um valor pequeno $\beta_1$. Isso permite que o otimizador avalie gradientes próximos à sequência rápida do Muon ($Z_t$) para adaptação inicial rápida.
2. Pós-Aquecimento ($t > T_0$): $\beta_t$ aumenta gradualmente em direção a 1 de acordo com:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   Aqui, $\rho \in [0, 1]$ controla a taxa de transição. À medida que $\beta_t$ aumenta, o ponto de avaliação do gradiente se desloca em direção à sequência média $X_t$. Isso suprime os componentes de alta curvatura (dominantes) do gradiente antes que eles sofram a ortogonalização do Muon, estabilizando assim a trajetória.

Crucialmente, este cronograma é independente do horizonte total de treinamento, permitindo treinamento a qualquer momento (o treinamento pode ser interrompido a qualquer ponto sem degradação de desempenho). O AMUSE não requer cronograma de decaimento de taxa de aprendizado.

Contribuições Principais

1. Análise Geométrica do Muon: O artigo fornece uma explicação geométrica para a eficiência e instabilidade do Muon. Ele demonstra que a ortogonalização promove atualizações orientadas ao volume (acelerando o progresso no rio), mas amplifica o ruído em direções dominantes (causando oscilações nas paredes do vale).
2. Estabilização via Interpolação: Os autores mostram que avaliar gradientes mais próximos da trajetória média (via Schedule-Free) reduz componentes dominantes antes da ortogonalização, mitigando oscilações.
3. Algoritmo AMUSE: Um otimizador inovador que desloca dinamicamente o ponto de avaliação do gradiente da trajetória rápida do Muon para a trajetória média estável. Isso combina a convergência rápida do Muon com a estabilidade dos métodos Schedule-Free.
4. Suporte a Treinamento a Qualquer Momento: Ao usar um cronograma variante no tempo independente do número total de passos e evitar o decaimento de taxa de aprendizado, o AMUSE suporta horizontes de treinamento flexíveis.

Resultados Experimentais

O AMUSE foi avaliado em diversos benchmarks, incluindo classificação de imagens (CIFAR-10/100, SVHN, ImageNet-1k), segmentação de imagens (ISIC 2018) e pré-treinamento de LLM (modelos estilo Llama com 124M, 720M e 1B de parâmetros no FineWeb).

• Fronteira de Pareto Desempenho-Iteração: O AMUSE supera consistentemente o AdamW (Schedule-Free) e o Muon em todas as tarefas.
• Ganhos de Eficiência:
  • No pré-treinamento de LLM de 720M, o AMUSE atinge o desempenho final do Muon com 1,51$\times$ menos passos.
  • No ImageNet com ResNet-50, ele alcança o mesmo desempenho com 1,12$\times$ menos passos.
  • No ImageNet com ajuste fino de ViT, ele requer 3,08$\times$ menos passos.
• Estabilidade: Diferentemente do Muon puro, que se beneficia significativamente da média exponencial de pesos (EWA) pós-hoc ou do decaimento de taxa de aprendizado, o AMUSE mantém uma trajetória estável durante todo o treinamento, mostrando melhoria mínima quando essas técnicas pós-hoc são aplicadas.
• Escalabilidade: O AMUSE escala efetivamente para grandes tamanhos de lote (até 2048) e grandes escalas de modelo (1 bilhão de parâmetros), onde outras variantes Schedule-Free às vezes lutam.

Significado e Alegações

O artigo afirma que o AMUSE representa uma abordagem principista para o design de otimizadores derivada de uma compreensão sistemática da dinâmica de treinamento e da geometria da paisagem de perda. Ao abordar explicitamente o trade-off entre o progresso rápido em massa e a estabilidade das paredes do vale, o AMUSE alcança desempenho superior sem a necessidade de agendamento manual de taxa de aprendizado.

Os autores notam modestamente que, embora o AMUSE supere otimizadores existentes, ele incorre em uma pequena sobrecarga de memória (uma cópia extra de estado para a sequência base $Z_t$) em comparação com o Muon puro, embora permaneça comparável ao AdamW e ao AdamW Schedule-Free. Eles sugerem que trabalhos futuros poderiam focar em alcançar estabilidade e desempenho semelhantes sem esse custo adicional de memória. O artigo enfatiza que a perspectiva "rio-vale" oferece um framework robusto para projetar otimizadores que sejam ao mesmo tempo eficientes e estáveis.