To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters

Este artigo investiga as desvantagens do otimizador Muon, argumentando que, embora ofereça maior velocidade de treinamento, ele elimina o viés de simplicidade inerente a métodos como o SGD, o que pode levar a modelos mais propensos a capturar características espúrias e ter dificuldade em identificar estruturas subjacentes comuns entre tarefas.

O essencial

🚀 O Dilema do "Super Carro" vs. O "Caminhão de Mudança"

Imagine que você precisa aprender a dirigir em uma cidade nova. Você tem duas opções de veículos:

1. O Caminhão de Mudança (SGD - Gradiente Descendente): É um veículo antigo, lento e pesado. Ele segue as ruas principais, faz curvas largas e para em cada esquina para verificar se está no caminho certo.
2. O Carro Esportivo Turbo (Muon): É um veículo novo, extremamente rápido e ágil. Ele ignora as curvas, corta por atalhos e chega ao destino em metade do tempo.

Por anos, todos usaram o Caminhão (ou o Adam, um carro um pouco mais moderno) porque era confiável. Recentemente, o Carro Esportivo (Muon) explodiu em popularidade porque é muito mais rápido para treinar redes neurais (os "cérebros" da inteligência artificial).

Mas a pergunta que este artigo faz é: Chegar rápido é sempre bom? O que acontece se você chegar rápido demais e não entender a cidade?

🧠 O Segredo: O "Viés de Simplicidade"

A descoberta principal do artigo é que o Caminhão (SGD) tem um superpoder secreto: ele tem um viés de simplicidade.

• Como o Caminhão (SGD) aprende: Ele aprende passo a passo. Primeiro, ele aprende a regra mais importante e simples da cidade. Só depois que domina essa, ele aprende a próxima regra um pouco mais complexa. É como aprender a andar de bicicleta: primeiro você aprende a equilibrar, depois a pedalar, depois a virar. Ele constrói o conhecimento de baixo para cima, criando uma estrutura sólida.
• Como o Carro Esportivo (Muon) aprende: Ele tenta aprender tudo ao mesmo tempo. Ele não espera dominar a regra simples para ir para a complexa. Ele "engole" todas as informações de uma vez só.

⚠️ O Perigo de Aprender Tudo de Uma Vez

O artigo mostra que, embora o Muon seja mais rápido, essa velocidade tem um custo: ele perde a capacidade de encontrar a "verdadeira estrutura" das coisas.

O autor usa duas analogias para mostrar onde o Muon falha:

1. O Exemplo do "Mapa Mágico" (Representações Compartilhadas)

Imagine que você está ensinando um robô a reconhecer números (1, 2, 3, 4) em diferentes línguas e escritas (inglês, árabe, grego).

• O Caminhão (SGD): Percebe que, no fundo, o "1" é sempre o mesmo conceito, não importa a escrita. Ele aprende o conceito abstrato de "1" e aplica a todas as línguas. Se você mostrar uma escrita nova que ele nunca viu, ele ainda consegue adivinhar o número.
• O Carro Esportivo (Muon): Como ele aprende tudo de uma vez e muito rápido, ele apenas decora cada exemplo que vê. Ele memoriza que "na escrita A, o 1 é assim" e "na escrita B, o 1 é assado". Se você mostrar uma nova escrita, ele entra em pânico e falha, porque ele não aprendeu a regra geral, apenas decorou os casos específicos.

Conclusão: O Muon é ótimo para decorar, mas ruim para entender o "porquê".

2. O Exemplo da "Mancha de Batata" (Correlações Espúrias)

Imagine que você está ensinando um robô a reconhecer se uma foto é de um gato ou de um cachorro. Mas, por acaso, todas as fotos de gatos no seu treinamento têm uma pequena mancha de batata no canto da imagem (um erro de dados).

• O Caminhão (SGD): Ele olha para o gato com cuidado. Ele vê o gato, depois vê a mancha. Ele decide: "Ok, vou focar no gato primeiro". Só depois de entender o gato, ele nota a mancha. Se você tirar a mancha na hora do teste, ele ainda reconhece o gato.
• O Carro Esportivo (Muon): Ele vê o gato e a mancha ao mesmo tempo, com a mesma velocidade. Como a mancha é um sinal fácil e rápido de identificar, ele decide: "Ah, é mais fácil olhar para a mancha de batata do que analisar o gato!". Ele aprende a confiar na mancha. Se você tirar a mancha no teste, o robô fica confuso e erra tudo.

Conclusão: O Muon é tão ágil que pula direto para o "atalho" óbvio (mesmo que seja falso), enquanto o Caminhão toma o tempo necessário para entender a realidade.

🎯 Qual é a Lição para o Futuro?

O artigo não diz que o Muon é "ruim". Ele é incrível para velocidade e para certos tipos de dados desbalanceados. Mas ele nos dá um aviso importante:

Não basta olhar apenas para a velocidade de treinamento.

Quando engenheiros criam novos otimizadores (novos "carros"), eles devem perguntar: "Que tipo de viés (tendência) este carro introduz?"

• Às vezes, queremos a velocidade do Muon.
• Outras vezes, queremos a sabedoria e a estrutura do Caminhão (SGD).

Se usarmos apenas o Muon em problemas complexos (como matemática, medicina ou ciências), podemos acabar com modelos que são rápidos, mas que apenas "decoraram" os dados em vez de realmente "entender" o mundo.

Resumo em uma frase: O Muon é um Ferrari que chega rápido ao destino, mas às vezes esquece de olhar pela janela e entender a paisagem; o SGD é um caminhão lento que, passo a passo, constrói um mapa mental sólido e confiável. O segredo é saber quando usar qual veículo.

Resumo técnico

Título: Usar ou Não Usar Muon: Como o Viés de Simplicidade em Otimizadores Importa

Autores: Sara Dragutinović e Rajesh Ranganath (NYU)

1. Problema e Motivação

O otimizador Adam tem sido historicamente a escolha padrão para treinar redes neurais profundas. Recentemente, o otimizador Muon (MomentUm Orthogonalized by Newton-Schulz) ganhou popularidade devido à sua velocidade de treinamento superior, sendo adotado como novo padrão em projetos como o NanoChat.

No entanto, a comunidade carece de uma compreensão teórica sobre os vieses indutivos (inductive biases) introduzidos pelo Muon. Enquanto o Muon converge mais rápido em termos de tempo de relógio (wallclock time), não se sabe quais trajetórias ele percorre no espaço de perda (loss landscape) e quais são as consequências disso para a qualidade da solução final. O problema central abordado é que o Muon pode estar removendo um viés de simplicidade naturalmente preservado por métodos mais antigos, como o Gradiente Descendente Estocástico (SGD), o que pode levar a modelos que memorizam em vez de generalizar estruturas subjacentes.

2. Metodologia

Os autores utilizam uma abordagem combinada de análise teórica e experimentação empírica:

• Análise Teórica (Redes Lineares Profundas):

  • Para tornar a análise tratável, os autores estudam redes lineares profundas (2 camadas) otimizadas por uma versão simplificada do Muon, chamada Spectral Gradient Descent (Spectral GD).
  • A Spectral GD assume uma decomposição em valores singulares (SVD) exata e remove o momento, focando na propriedade de ortonormalização do passo de atualização.
  • Comparam-se as dinâmicas de aprendizado do Gradiente Descendente (GD) clássico com as da Spectral GD.
  • O foco está em como cada método lida com os valores singulares da matriz de pesos e como isso afeta a complexidade da solução ao longo do tempo.

• Experimentos Empíricos:

  • Tarefa de "Roteamento" (Routing Task): Um cenário onde múltiplas fontes de entrada mapeiam para um espaço oculto comum. O objetivo é aprender uma representação compartilhada (estrutura subjacente) que generalize para pares entrada-saída não vistos durante o treinamento.
  • Correlações Espúrias (MNIST): Um cenário onde os dados contêm uma característica espúria (intensidade de um pixel específico por classe) além da característica real (o dígito). O objetivo é avaliar qual otimizador aprende a depender da característica real versus a espúria.
  • Comparações diretas entre Muon, SGD e Adam.

3. Contribuições Principais

A. Análise Teórica: Perda do Viés de Simplicidade

• Dinâmica do SGD/GD: O Gradiente Descendente exibe um comportamento de "saddle-to-saddle" (de sela para sela). Ele aprende os componentes principais (valores singulares) de forma sequencial. Primeiro, aprende os componentes de maior magnitude, aumentando gradualmente o posto (rank) da solução. Isso atua como uma regularização implícita, favorecendo soluções mais simples inicialmente.
• Dinâmica da Spectral GD (Muon): O Muon/Spectral GD aprende todos os componentes principais simultaneamente e na mesma velocidade. Ele não passa pelas "saddles" que o SGD atravessa.
• Consequência: Embora a Spectral GD seja mais rápida (evitando os platôs de perda causados pelas saddles), ela perde o viés de simplicidade. Ela tende a aprender todas as modalidades (frequências) dos dados ao mesmo tempo, em vez de priorizar a estrutura dominante primeiro.

B. Experimentos: Generalização e Estrutura Compartilhada

• Falha na Generalização de Estrutura: Na tarefa de roteamento, o SGD conseguiu aprender a representação compartilhada subjacente, generalizando perfeitamente para pares de entrada-saída nunca vistos. O Spectral GD (e o Muon), por outro lado, memorizou os pares vistos durante o treinamento, falhando em generalizar para os não vistos.
• Análise Espectral: A solução do SGD convergiu para uma estrutura de baixo posto (rank = 4, igual à dimensão da tarefa real), enquanto a solução do Spectral GD exibiu um posto efetivo muito mais alto e um espectro de cauda pesada, indicando memorização em vez de aprendizado estrutural.

C. Experimentos: Correlações Espúrias

• Em um cenário com o dataset MNIST e pixels espúrios, o SGD demonstrou maior robustez.
• O SGD aprendeu a depender dos dígitos reais (característica dominante) por mais tempo antes de começar a aprender o pixel espúrio.
• O Muon aprendeu ambas as características (dígito e pixel espúrio) simultaneamente. Isso resultou em uma dependência mais rápida da característica espúria, levando a uma queda de desempenho mais cedo em conjuntos de validação sem o pixel espúrio (especialmente se o uso de early stopping for aplicado).
• O artigo demonstra que a "velocidade" do Muon pode ser uma faca de dois gumes: se a característica espúria for forte, o Muon a aprende tão rápido quanto a característica real, prejudicando a generalização.

4. Resultados Chave

1. Velocidade vs. Qualidade: O Muon é mais rápido porque evita as saddles e aprende tudo simultaneamente, mas isso custa a perda do viés de simplicidade que ajuda na generalização.
2. Memorização vs. Descoberta: Em tarefas que exigem descobrir uma estrutura comum subjacente (como em aprendizado multi-modal ou generalização fora da distribuição), o Muon tende a memorizar os dados de treinamento em vez de inferir a regra geral.
3. Sensibilidade a Correlações Espúrias: O Muon é mais propenso a ajustar-se a características espúrias rapidamente, pois não prioriza a estrutura dominante do dado da mesma forma sequencial que o SGD.
4. Dependência do Cenário: A preferência entre Muon e SGD depende da distribuição dos dados e da dominância das características. Em dados desbalanceados, o Muon pode ser superior, mas em tarefas que exigem generalização estrutural, o SGD pode ser melhor.

5. Significado e Implicações

• Reavaliação de Otimizadores: O trabalho alerta que a escolha do otimizador não deve ser baseada apenas na velocidade de convergência ou no desempenho in-distribution. Os vieses indutivos introduzidos pelo algoritmo são fundamentais para o comportamento final do modelo.
• Riscos para Modelos de Linguagem e IA: Se o Muon tende a memorizar em vez de generalizar estruturas, seu uso em larga escala para treinar LLMs (Large Language Models) pode levar a modelos que "decoram" dados em vez de aprender regras lógicas ou estruturas causais, o que é crítico para tarefas complexas como matemática ou raciocínio.
• Futuro de Pesquisa: O artigo sugere que novos otimizadores devem ser projetados considerando como eles navegam no espaço de perda. Uma meta futura seria criar otimizadores que mantenham a velocidade do Muon (evitando saddles) mas preservem o viés de simplicidade e a trajetória sequencial do SGD.

Em resumo, o paper argumenta que a simplicidade (aprender o que é mais importante primeiro) é um mecanismo crucial para a generalização, e que otimizar puramente pela velocidade (como o Muon faz) pode comprometer a capacidade do modelo de descobrir a verdadeira estrutura dos dados.