🤖 machine learning

Approximate Muon with low-rank adapters

Este artigo apresenta o sMuon, um método eficiente que aproxima o otimizador Muon para o ajuste fino eficiente de parâmetros de baixo posto ao linearizar o objetivo e resolvê-lo via mínimos quadrados, permitindo assim os benefícios de desempenho do Muon em configurações como SFT e pré-treinamento ReLoRA, onde ele era anteriormente incompatível com LoRA.

Autores originais: Ben Anson, Conor Houghton, Edward Milsom

Publicado 2026-08-17

📖 4 min de leitura☕ Leitura rápida

CC0 1.0

Autores originais: Ben Anson, Conor Houghton, Edward Milsom

Artigo original dedicado ao domínio público sob CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar uma nova habilidade a um robô gigante e superinteligente, como escrever poesia ou resolver problemas matemáticos. Você não quer retreinar todo o cérebro do robô do zero porque isso leva uma eternidade e custa uma fortuna. Em vez disso, você quer anexar um pequeno e leve "módulo de treinamento" ao cérebro existente dele. Isso é chamado de Ajuste de Parâmetros Eficiente (PEFT - Parameter-Efficient Fine-Tuning). Pense nisso como colocar um headset especializado em um robô de uso geral; o headset ensina a ele o novo truque sem alterar o hardware principal do robô.

Uma maneira popular de construir esses headsets é chamada de LoRA (Low-Rank Adaptation). É como construir o headset a partir de duas folhas finas e planas de papel que, quando empilhadas, criam a forma da nova habilidade. Para ensinar o robô, precisamos de um otimizador — um treinador matemático que diz ao headset como ajustar sua forma para melhorar. O treinador mais comum é chamado de AdamW, mas existe um treinador mais novo e poderoso chamado Muon. O Muon é famoso por ser incrivelmente eficiente durante o pré-treinamento (ensinando as primeiríssimas lições ao robô do zero). Ele funciona garantindo que as atualizações do cérebro do robô ocorram de uma forma perfeitamente equilibrada e "ortogonal", como um dançarino girando perfeitamente sobre um pé sem oscilar.

No entanto, há um problema. Quando você tenta usar o Muon para treinar o pequeno headset LoRA, a matemática falha. O Muon quer equilibrar a camada inteira do cérebro do robô, mas o headset LoRA é apenas uma pequena fatia de baixo posto dessa camada. É como tentar usar a batuta de um maestro para dirigir um único violinista de uma forma que faça sentido para toda a orquestra; a geometria simplesmente não se encaixa. A solução paliativa padrão tem sido ignorar a orquestra e apenas dizer ao violinista para girar por conta própria, mas isso perde a magia do Muon. Este artigo pergunta: Podemos encontrar uma maneira de fazer o Muon funcionar para esses pequenos headsets sem quebrar a matemática ou atrasar tudo?

Os autores, Ben Anson, Conor Houghton e Edward Milsom, dizem que sim. Eles introduzem um novo método chamado sMuon (small Muon). Em vez de tentar forçar o impossível, eles usam um truque matemático inteligente: eles "linearizam" o problema. Imagine que você está tentando caminhar em uma corda bamba enquanto segura um guarda-chuva gigante e instável. Em vez de tentar equilibrar o guarda-chuva inteiro de uma vez, você finge que o guarda-chuva é apenas uma vara reta por um breve segundo, calcula o passo perfeito e depois se ajusta. Ao fazer isso, eles encontam a melhor maneira possível de atualizar o headset LoRA que aproxima a dança perfeita do Muon.

O artigo constata que este novo método sMuon funciona surpreendentemente bem. Quando testaram em vários modelos de linguagem, o sMuon frequentemente teve um desempenho melhor do que o treinador padrão AdamW e outras tentativas de Muon. Foi particularmente impressionante quando usado em um robô que já havia sido pré-treinado com Muon (chamado Moonlight), onde o s_Muon alcançou as pontuações mais altas em seis de onze tarefas diferentes. Em um experimento de pré-treinamento, o sMuon empatou com o melhor resultado na redução da taxa de erro, mas o fez cerca de 30% mais rápido do que seu rival mais próximo, um método chamado Riemannion.

Crucialmente, os autores mostram que seu método não é apenas uma ideia teórica; é também prático. Enquanto outros métodos avançados exigem cálculos pesados e lentos (como desmontar e remontar matrizes complexas), o sMuon baseia-se quase inteiramente em multiplicações de matrizes simples — o tipo de matemática para o qual os chips de computador modernos são construídos para operar em velocidade de raio. Isso significa que o sMuon é rápido o suficiente para ser usado no treinamento do mundo real sem atrasar o processo. O artigo sugere que, ao respeitar a geometria de toda a camada em vez de tratar as partes LoRA como peças separadas e não relacionadas, podemos obter um aumento significativo de desempenho com muito pouco custo adicional. É um lembrete de que, às vezes, a melhor maneira de ensinar um novo truque a um robô gigante é dar a ele um empurrãozinho pequeno, inteligente e matematicamente elegante.

=== RESUMO ===

Resumo Técnico: Aproximação de Muon com Adaptadores de Baixo Rank (sMuon)

Declaração do Problema

O otimizador Muon demonstrou um desempenho superior no pré-treinamento de grandes modelos de linguagem (LLMs) ao ortogonalizar as atualizações de peso em cada camada, efetivamente equalizando os valores singulares. Este comportamento corresponde ao declive mais íngreme sob a norma espectral, controlando a mudança de pior caso para uma camada como um operador linear em vez de um vetor achatado. No entanto, o Muon é raramente utilizado para Ajuste Fino de Parâmetros Eficientes (PEFT), particularmente com Adaptação de Baixo Rank (LoRA).

A obstrução matemática central é que o Muon requer a ortogonalização do passo de declive para a matriz de peso inteira da camada. No cenário LoRA, onde a atualização de peso é restringida a uma parametrização de baixo rank $\Delta W = BA$ (com $B \in \mathbb{R}^{d_1 \times r}$ e $A \in \mathbb{R}^{r \times d_2}$ ), é matematicamente impossível encontrar, de forma geral, atualizações $\delta B$ e $\delta A$ de modo que o $\delta \Delta W$ resultante corresponda perfeitamente ao gradiente ortogonalizado de rank total exigido pelo Muon.

As soluções alternativas existentes incluem:

Muon por fator: Ortogonaliza $B$ e $A$ separadamente, o que ignora a geometria da camada completa.
LoRA-Muon: Relaxa o objetivo utilizando a desigualdade triangular.
Riemannion: Formula o problema como otimização Riemanniana em uma variedade de baixo rank, o que requer Decomposições de Valor Singular (SVD) truncadas caras durante o passo de retração, levando a uma sobrecarga computacional significativa.

Metodologia

Os autores propõem o sMuon (small Muon), um método que aproxima a solução de um objetivo Muon relaxado no cenário de baixo rank. A abordagem baseia-se em linearização e otimização de mínimos quadrados para encontrar o passo de declive ideal.

1. Linearização e Formulação de Mínimos Quadrados

O objetivo é encontrar atualizações $\delta B, \delta A$ de tal forma que o $\delta \Delta W$ implícito $\approx \delta B A + B \delta A$ aproxime o passo ideal de Muon $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ , onde $\text{msign}(X) = UV^\top$ é o fator polar da SVD.

Como o passo ideal é geralmente de rank total e não pode ser realizado exatamente por fatores de baixo rank, os autores linearizam o problema descartando termos de segunda ordem ( $\delta B \delta A$ ) e resolvendo um problema de mínimos quadrados:
$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$
Isso produz uma solução de forma fechada envolvendo inversas de Moore-Penrose ( $B^\dagger, A^\dagger$ ). Crucialmente, os autores mostram que a atualização ótima não é única em termos dos fatores $(\delta B, \delta A)$ , mas a mudança de primeira ordem resultante na camada completa $\delta \Delta W$ é única.

2. Computação Eficiente Baseada Apenas em Matmul

Uma implementação direta da solução de forma fechada exige a ortogonalização de uma grande matriz $d_1 \times d_2$ , o que é computacionalmente proibitivo. Para resolver isso, os autores derivam uma formulação eficiente (Algoritmo 2) que evita a ortogonalização explícita de matrizes grandes:

Gradiente de Baixo Rank: Em vez de usar o gradiente de rank total $G_{\Delta W}$ , eles utilizam uma projeção de baixo rank $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ , que é equivalente de primeira ordem para o objetivo de otimização.
Redução de Dimensionalidade: Eles exploram a identidade $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ , onde $U$ e $V$ são bases ortonormais que abrangem os espaços coluna e linha de $H$ . Como $H$ tem rank de no máximo $2r$ , isso reduz a ortogonalização dispendiosa de $d_1 \times d_2$ para uma ortogonalização de $2r \times 2r$ .
Implementação: Todo o algoritmo baseia-se exclusivamente em multiplicações de matrizes, raízes quadradas inversas de matrizes e decomposições polares de matrizes pequenas de $r \times r$ ou $2r \times 2r$ . Isso evita o overhead de decomposições QR e SVD exigidos pelo Riemannion.

3. Detalhes Algorítmicos

O otimizador sMuon incorpora:

Decaimento de Peso Dividido (Split Weight Decay): Para corresponder à dinâmica do treinamento de parâmetros totais, o decaimento de peso é aplicado via um fator de escala $s = \sqrt{1-\lambda\eta}$ aos fatores $A$ e $B$ antes da atualização do gradiente.
Transporte de Momento: Para garantir a invariância de parametrização (ou seja, o otimizador se comporta de forma consistente independentemente da fatoração específica $B, A$ ), os buffers de momento são projetados nos espaços atuais dos fatores a cada passo.

Principais Contribuições

Derivação do sMuon: Um método de forma fechada para treinar adaptadores LoRA com Muon, utilizando uma aproximação de primeira ordem do objetivo de declive mais íngreme com norma espectral regularizada.
Implementação Eficiente: Um algoritmo "apenas matmul" que é computacionalmente eficiente para hardware moderno, evitando decomposições caras. Oferece velocidade comparável ao Per-factor Muon e LoRA-Muon, sendo significativamente mais rápido que o Riemannion, especialmente para ranks LoRA maiores.
Avaliação Empírica: Experimentos extensos em quatro modelos base (incluindo um pré-treinado com Muon e três pré-treinados com AdamW) e onze benchmarks (senso comum e codificação).

Resultados

O artigo reporta as seguintes descobertas:

Desempenho: O sMuon apresenta um desempenho favorável em benchmarks de SFT. Embora os resultados dependam do modelo, o sMuon alcançou a melhor precisão em seis de onze tarefas ao realizar o ajuste fino do modelo Moonlight-16B-A3B (que foi pré-treinado com Muon). Ele geralmente supera o AdamW e o Per-factor Muon.
Pré-treinamento ReLoRA: Em um cenário de pré-treinamento usando ReLoRA (fusão e atualização periódica de adaptadores), o sMuon alcançou a menor perda de validação final, empatando com o Riemannion, mas sendo aproximadamente 30% mais rápido.
Propriedades Espectrais: A análise empírica dos espectros de atualização de peso mostra que o sMuon mantém um espectro mais uniforme (mais próximo do comportamento ideal do Muon) em comparação ao Per-factor Muon e LoRA-Muon.
Eficiência: A Tabela 4 demonstra que, embora o sMuon seja ligeiramente mais lento que o AdamW, ele é significativamente mais rápido que o Riemannion. Para um rank de $r=64$ , o Riemannion adiciona cerca de 33% de overhead ao tempo de passo, enquanto o sMuon adiciona apenas ~0,8%.

Significância e Alegações

O artigo afirma que o sMuon fornece uma maneira principiada, eficiente em termos de memória e computacionalmente tratável de aplicar a otimização Muon a adaptadores de baixo rank. Ao respeitar a geometria da camada completa em vez de tratar os fatores LoRA como matrizes independentes, o sMuon alcança melhorias moderadas de desempenho, particularmente quando o modelo base também foi pré-treinado com Muon. Os autores enfatizam que seu método evita as "restrições de variedade caras" da otimização Riemanniana sem sacrificar o desempenho, tornando o ajuste fino baseado em Muon mais acessível para aplicações práticas.

Os autores mantêm a modéstia, observando que os resultados são "dependentes do modelo e da avaliação" e que a principal contribuição é uma aproximação limpa e eficiente que preenche a lacuna entre os benefícios teóricos do Muon e as restrições práticas do PEFT.