Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

Resumo Técnico: Muon Aprende Soluções Equilibradas em Fatoração de Matrizes sem a Lenta Dinâmica de Sela-para-Sela

Enunciado do Problema
O artigo investiga a dinâmica de aprendizado do otimizador Muon no contexto de fatoração de matrizes, definida como a minimização de $L(P, Q) = \frac{1}{2}\|M^* - PQ^\top\|_F^2$. A fatoração de matrizes serve como um modelo minimalista para estudar a dinâmica de parâmetros não lineares e o aprendizado de representação, exibindo comportamentos relevantes para o aprendizado profundo, como o viés de baixo posto (low-rank bias) e o alinhamento de representação. Os autores comparam especificamente o Muon com o Gradiente Descendente (GD) padrão para entender como a regra de atualização única do Muon altera a trajetória da otimização, particularmente em relação à velocidade de convergência, estabilidade e à natureza das soluções encontradas a partir de uma inicialização aleatória pequena.

Metodologia
O estudo foca na regra de atualização do Muon, que ortogonaliza o gradiente antes de aplicar a taxa de aprendizado. Para uma matriz de pesos $W$, a atualização é $\Delta W = -\eta \cdot \widehat{\nabla_W L}$, onde $\widehat{G} = UV^\top$ é a decomposição polar (ortogonalização) do gradiente $G=USV^\top$. Isso garante que todas as atualizações sejam matrizes (semi-)ortogonais escalonadas por $\eta$.

Os autores analisam as dinâmicas através de três lentes:

1. Redução Escalar: Assumindo que os pesos estão alinhados, o problema se desacopla em problemas escalares independentes onde o resíduo $r = m^* - pq$ dita a atualização.
2. Dinâmica de Alinhamento: Os autores derivam taxas teóricas para como os vetores singulares de $P$ e $Q$ se alinham entre si (alinhamento interno) e com a matriz alvo $M^*$ (alinhamento com o alvo) sob diferentes regimes de parâmetros (subparamatrizado, exatamente parametrizado e sobreparametrizado).
3. Grandezas Conservadas: O artigo identifica quantidades matriciais específicas conservadas durante o fluxo do Muon versus GD.

Principais Contribuições e Resultados

1. Evitação da Lenta Dinâmica de Sela-para-Sela
O Gradiente Descendente padrão, a partir de uma inicialização pequena, exibe dinâmicas de "sela-para-sela". Os valores singulares crescem multiplicativamente, o que significa que os modos menores crescem mais devagar que os maiores. O otimizador permanece próximo de sela de baixo posto, aprendendo os modos sequencialmente (aprendizado por etapas).
Em contraste, o Muon exibe dinâmicas aditivas. Uma vez ocorrido o alinhamento, todos os valores singulares das matrizes de peso crescem a uma taxa uniforme ($\eta$ por passo), independentemente de sua magnitude atual. Consequentemente, os modos menores convergem antes dos maiores, e o otimizador aprende todos os modos superiores de $M^*$ simultaneamente, sem pontos de sela intermediários.

2. Estabilidade Além das Taxas de Aprendizado Críticas
O Gradiente Descendente é limitado por um limiar de taxa de aprendizado crítica determinado pela nitidez local da perda (especificamente $2/\lambda_{\max}(\nabla^2 L)$). Exceder esse limite causa divergência, e a taxa de convergência é limitada pelo número de condição do problema.
O Muon permanece estável mesmo quando a taxa de aprendizado excede esse limiar crítico. Como a magnitude da atualização é fixa pela etapa de ortogonalização, os iterados não podem divergir. Em vez disso, com taxas de aprendizado elevadas, a otimização procede rapidamente para a variedade da solução (solution manifold) e oscila ali. Isso desacopla a taxa de aprendizado do número de condição, permitindo uma convergência rápida via amortecimento (annealing) geométrico da taxa de aprendizado.

3. Diferentes Grandezas Conservadas e Soluções Equilibradas
Ambos os otimizadores encontram "soluções equilibradas" (onde $P^\top P = Q^\top Q$) a partir de uma inicialização nula, mas conservam quantidades diferentes:

• Fluxo de Gradiente: Conserva a quantidade matricial $\Delta_2 = P^\top P - Q^\top Q$. As linhas de fluxo no espaço de parâmetros são hiperbólicas.
• Fluxo Muon: Conserva a quantidade matricial $\Delta_1 = \sqrt{P^\top P} - \sqrt{Q^\top Q}$. As linhas de fluxo são afins (raios paralelos).
  Os autores derivam que o fluxo Muon é gerado por um Hamiltoniano $H(P, Q) = \text{Tr}(\Delta_1) = \|P\|_* - \|Q\|_*$ (a diferença nas normas nucleares). Essa lei de conservação garante que, se os pesos começarem equilibrados, eles permanecerão equilibrados, levando à solução equilibrada.

4. Dinâmica de Alinhamento e Taxas
O artigo deriva taxas de alinhamento para a fase de "alinhamento silencioso" (onde os vetores singulares se alinham enquanto a perda está plana):

• Subparametrizado ($d < n$): O alinhamento interno é trivial; o alinhamento com o alvo decai polinomialmente como $(t/\alpha_{\text{eff}})^{-2(1-s^{-1})}$, onde $s$ é o gap espectral relativo.
• Sobreparametrizado ($d > n$): O alinhamento com o alvo é trivial; o desalinhamento interno decai rapidamente como $(t/\alpha_{\text{eff}})^{-4}$.
• Geral: O alinhamento ocorre mais cedo com escalas de inicialização menores. Os autores mostram que essas taxas teóricas predizem o alinhamento empírico em configurações gerais.

5. Cronograma de Taxa de Aprendizado com Pico (Spiked Learning Rate Schedule)
Aproveitando a compreensão das dinâmicas de alinhamento, os autores constroem um cronograma específico de taxa de aprendizado que alcança um alinhamento quase perfeito em exatamente dois passos de otimização (para $d \ge n$):

1. Passo 1: Definir $\eta^{(0)} = \alpha$ (escala de inicialização) para alinhar $P$ com $Q$.
2. Passo 2: Definir um pico elevado $\eta^{(1)} \gg \eta^{(0)}$ para alinhar simultaneamente o modelo com o alvo $M^*$ e levar os pesos para a variedade da solução.
3. Amortecimento: Realizar o amortecimento geométrico da taxa de aprendizado para alcançar convergência linear.
   Este cronograma demonstra que os insights teóricos sobre o fluxo podem ser traduzidos em estratégias práticas de otimização, embora os autores notem que isso é uma prova de conceito para regimes específicos.

Significância e Alegações
O artigo afirma que o Muon altera fundamentalmente o cenário de otimização da fatoração de matrizes ao substituir a dinâmica multiplicativa e repleta de sela por um crescimento uniforme e aditivo. Isso permite que o Muon:

• Aprenda todos os modos da matriz alvo à mesma taxa, evitando o aprendizado sequencial e lento do GD.
• Opere de forma estável com taxas de aprendizado muito superiores aos limiares críticos que restringem o GD, liberando a taxa de aprendizado do número de condição do problema.
• Alcance soluções equilibradas via uma lei de conservação distinta ($\sqrt{P^\top P} - \sqrt{Q^\top Q}$) que mantém o equilíbrio dos pesos sem exigir a geometria hiperbólica do GD.

Os autores enfatizam que, embora sua análise seja restrita ao gradiente espectral não modificado (ignorando momentum, decaimento de peso e iterações de Newton-Schulz) e à simples fatoração de matrizes, a compreensão mecanística do alinhamento e do crescimento uniforme fornece uma base para entender o sucesso do Muon em modelos maiores e mais complexos. O trabalho serve como uma explicação teórica de por que o Muon acelera o treinamento em configurações espectrais heterogêneas ou mal condicionadas.