Sumário Técnico: Limite de Convergência e Tamanho de Lote Crítico do Otimizador Muon

Declaração do Proble Problema

Embora o otimizador Muon tenha demonstrado um forte desempenho empírico como um potencial sucessor para otimizadores padrão como o AdamW, seus fundamentos teóricos permanecem subdesenvolvidos. A literatura existente carece de provas formais de convergência para o Muon, particularmente em configurações práticas envolvendo o momentum de Nesterov e o decaimento de peso (weight decay). Este artigo visa preencher essas lacunas, fornecendo uma análise de convergência rigorosa e derivando um limite inferior para o tamanho de lote crítico do Muon.

Metodologia

Os autores conduzem uma análise teórica do Muon através de quatro configurações práticas:

1. Sem momentum de Nesterov e sem decaimento de peso.
2. Com momentum de Nesterov e sem decaimento de peso.
3. Sem momentum de Nesterov e com decaimento de peso.
4. Com momentum de Nesterov e com decaimento de peso.

Estrutura Teórica:

• Análise de Convergência: Os autores derivam limites superiores para a norma nuclear média esperada do gradiente, $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$. Eles utilizam suposições padrão de suavidade $L$ ($L$-smoothness) e variância limitada do gradiente. Uma inovação fundamental é a análise da direção de busca do Muon, que é obtida pela ortogonalização de uma matriz de momentum $C_t$. Os autores consideram o posto (rank) da matriz de erro de momentum ($r_1$) e da direção de busca ($r_2$), em vez da dimensão total dos parâmetros.
• Aproximação Newton-Schulz: Como o Muon aproxima na prática a ortogonalização via iteração Newton-Schulz em vez de uma SVD exata, os autores analisam o impacto desse erro de aproximação ($\delta(k)$) na convergência, mostrando que ele introduz apenas um termo negligenciável.
• Derivação do Tamanho de Lote Crítico (CBS): Os autores modelam o número de passos de treinamento $T$ necessários para atingir uma precisão alvo $\epsilon$ como uma função do tamanho de lote $b$. Ao minimizar a complexidade SFO (Stochastic First-Order Oracle), eles derivam um limite inferior para o tamanho de lote crítico $b^*$.

Validação Experimental:
As previsões teóricas são validadas através de uma estratégia de dois níveis:

1. Experimentos Controlados de Full-Muon: Uma tarefa sintética Professor-Aluno onde o Muon atualiza todos os parâmetros, correspondendo estritamente às suposições teóricas.
2. Experimentos Híbridos de Escala Prática: Treinamento de ResNet-18/VGG-16 em conjuntos de dados CIFAR e Llama3.1 (320M) no corpus C4. Nestas configurações, o Muon é aplicado a parâmetros com formato de matriz enquanto o AdamW lida com o restante, refletindo o uso padrão.

Principais Contribuições

1. Provas de Convergência para Quatro Variantes

O artigo estabelece garantias de convergência para todas as quatro variantes do Muon. Os limites superiores derivados para a norma do gradiente média esperada são da ordem de:
$$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$$
(com pequenas variações para os termos de Nesterov e decaimento de peso).

• Significância: Diferente de otimizadores baseados em vetores onde a variância escala com o número total de parâmetros, o limite do Muon escala com o posto $r_1$ da matriz de erro de momentum. Evidências empíricas sugerem que $r_1 \ll n$ em regimes superparametrizados, implicando um limite teórico mais justo.
• Norma Nuclear: A análise utiliza a norma nuclear, que é um critério mais rigoroso do que a norma de Frobenius usada em muitas análises baseadas em vetores, sugerindo que o Muon respeita a estrutura intrínseca dos parâmetros da rede neural.

2. Papel do Decaimento de Peso e Estabilidade

Os autores provam que a incorporação do decaimento de peso garante a limitação quase certa (almost-sure boundedness) tanto das normas dos parâmetros ($\|W_t\|_F$) quanto das normas do gradiente ($\|\nabla f(W_t)\|_F$) sem depender da suposição restritiva de gradientes limitados.

• Condição de Estabilidade: Eles derivam uma condição necessária para estabilidade: a taxa de aprendizado $\eta$ e o coeficiente de decaimento de peso $\lambda$ devem satisfazer $\eta \leq 1/\lambda$. Resultados experimentais confirmam que o treinamento torna-se instável quando $\eta$ excede este limiar.
• Regularização Implícita: Embora o decaimento de peso não melhore a ordem da taxa de convergência assintótica, ele fornece regularização implícita ao limitar as normas, o que explica seu desempenho empírico superior.

3. Limite Inferior do Tamanho de Lote Crítico (CBS)

O artigo deriva um limite inferior para o tamanho de lote crítico $b^*$, mostrando sua dependência de hiperparâmetros:
$$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$$

• Dependência de Hiperparâmetros: A fórmula revela que:
  • Momentum ($\beta$): À medida que $\beta \to 1$, o termo $(1-\beta)$ desaparece, sugerindo que o tamanho de lote crítico diminui.
  • Decaimento de Peso ($\lambda$): O fator $1/(1-\lambda)^2$ indica que aumentar o decaimento de peso aumenta o tamanho de lote crítico.
  • Posto ($r_1$): O limite escala com o posto do erro de momentum, ligando o CBS à estrutura de baixo posto dos gradientes.
• Escalonamento Qualitativo: Os autores enfatizam que a fórmula fornece uma lei de escalonamento qualitativa em vez de uma previsão absoluta, pois depende de constantes específicas do problema ($\sigma^2, r_1, \epsilon$). No entanto, a razão entre os CBS de diferentes variantes do Muon depende apenas de $\beta$ e $\lambda$.

Resultados

Descobertas Teóricas:

• A taxa de convergência coincide com métodos padrão do tipo SGD não convexo ($O(T^{-1/4})$ sob esquemas específicos de hiperparâmetros), mas com um termo de variância mais justo devido à dependência de posto.
• O erro de aproximação Newton-Schulz tem um efeito limitado na convergência, decaindo exponencialmente com o número de iterações.

Descobertas Experimentais:

• Estabilidade: Experimentos em ResNet-18/CIFAR-10 validam a condição de estabilidade $\eta \leq 1/\lambda$. As normas do gradiente e a perda de treinamento são minimizadas próximo a este limite.
• Velocidade de Convergência: O Muon com momentum de Nesterov e decaimento de peso atinge consistentemente a diminuição mais rápida da perda e da norma do gradiente em comparação com o AdamW, Momentum SGD e outras variantes do Muon.
• Validação do Tamanho de Lote Crítico (CBS):
  • Efeitos de Hiperparâmetros: Experimentos confirmam que adicionar momentum de Nesterov desloca o CBS para valores maiores, e aumentar o decaimento de peso aumenta ainda mais o CBS. Por exemplo, aumentar $\lambda$ de 0 para 0,5 resultou em um deslocamento de CBS de $\approx 1024$ para $\geq 4096$, correspondendo ao escalonamento previsto de $1/(1-\lambda)^2$ (razão de 4,0).
  • Escalonamento de Posto: Variar a largura do modelo mostrou que o posto efetivo $r_1$ escala com a largura, e o produto $r_1 \cdot \sigma^2$ permanece relativamente constante, explicando por que os valores empíricos de CBS pareceram planos através das larguras, apesar da dependência teórica de posto.
  • Eficiência: O Muon geralmente requer menos passos e menor complexidade SFO do que o AdamW, particularmente em tamanhos de lote pequenos a moderados.

Significância e Alegações

O artigo afirma fornecer a primeira análise de convergência rigorosa do Muon que incorpora simultaneamente o momentum de Nesterov e o decaimento de peso, abordando uma lacuna na compreensão teórica deste otimizador emergente.

• Insight Teórico: Demonstra que o decaimento de peso não é meramente uma ferramenta de regularização, mas um mecanismo que garante a limitação e a estabilidade, removendo a necessidade de suposições de gradientes limitados.
• Orientação Prática: O limite inferior de CBS derivado oferece orientação acionável para praticantes. Ele esclarece como o ajuste de $\beta$ e $\lambda$ afeta qualitativamente o tamanho de lote ideal, sugerindo que configurações de alto momentum podem ser treinadas eficientemente com lotes menores, enquanto um maior decaimento de peso requer lotes maiores para manter a eficiência.
• Validação: O trabalho une a teoria e a prática ao validar essas previsões através de cargas de trabalho de visão e modelagem de linguagem, confirmando que as leis de escalonamento teóricas se sustentam mesmo em configurações de otimização híbridas onde o Muon é aplicado apenas a parâmetros com formato de matriz.

Os autores observam modestamente que sua análise de matriz única não modela explicitamente a heterogeneidade camada a camada, mas argumentam que a alta correlação das escalas de ruído do gradiente entre as camadas apoia a generalizabilidade de suas descobertas. Eles também reconhecem que, embora o Muon supere o Shampoo em seus experimentos, derivar limites análogos para o Shampoo permanece um problema em aberto.