Resumo Técnico: Muon Aprende Características Mais Robustas e Transferíveis que o Adam

Declaração do Problema

Embora o otimizador Muon (Momentum Orthogonalized by Newton-Schulz) tenha demonstrado recentemente uma eficiência de estado da arte no pré-treinamento de Grandes Modelos de Linguagem (LLMs) e classificadores de visão — treinando frequentemente quase duas vezes mais rápido que o Adam e o SGD — a natureza específica de suas vantagens na aprendizagem de características permanece obscura. A literatura existente tem explicado amplamente o sucesso do Muon através de lentes de otimização teórica, como o gradiente descendente mais íngreme sob normas espectrais, dinâmica de memória associativa ou tratamento de dados de cauda pesada. No entanto, faltava uma investigação sistemática sobre se o Muon aprende características melhores do que o Adam e o SGD, e como essas características se manifestam nas representações do modelo. Este artigo aborda a questão: O Muon aprende características mais robustas e transferíveis do que o Adam e o SGD, e quais propriedades caracterizam essa vantagem?

Metodologia

Os autores empregam uma abordagem dupla combinando avaliação empírica extensiva em domínios de visão e linguagem com análise teórica em um problema de classificação estilizado.

Avaliação Empírica

O estudo avalia a qualidade das características através de duas métricas primárias: robustez e transferibilidade.

• Robustez: Modelos pré-treinados (ResNet-18, ViT-S, GPT-2, GPT-2 Medium) são testados em entradas corrompidas. Modelos de visão são avaliados no ImageNet-C (15 tipos de corrupção), enquanto modelos de linguagem são avaliados em um benchmark recém-construído, FineWeb10B-C, que aplica ataques de erro de caractere por nível (inserções, deleções, substituições, trocas).
• Transferibilidade: Características pré-treinadas são adaptadas para tarefas de jusante (downstream). Para visão, classificadores lineares são treinados em backbones congelados em quatro conjuntos de dados (EuroSAT, Food-101, Flowers-102, Stanford Cars). Para linguagem, modelos completos são ajustados (fine-tuned) em conjuntos de dados de seguimento de instruções (Alpaca, Dolly, WizardLM).
• Análise de Estados Ocultos: Para entender o mecanismo por trás desses ganhos de desempenho, os autores analisam os estados ocultos camada a camada usando:
  • Margens de Logit: Sondas lineares treinadas (framework tuned-lens) decodificam estados ocultos para computar a margem entre a classe/token correto e o distrator mais forte.
  • Diversidade Espectral: O posto efetivo (effective rank) e a fração de energia Top-k das matrizes de estados ocultos são computados para medir a diversidade e a distribuição de informação através de direções ortogonais.
• Protocolo Experimental: Todas as comparações seguem um protocolo de orçamento correspondente (matched-budget protocol) onde arquitetura, dados e duração de treinamento são fixos, enquanto hiperparâmetros (taxa de aprendizado, scheduler) são ajustados independentemente para cada otimizador para isolar as diferenças de representação do esforço de engenharia.

Análise Teórica

Os autores fornecem suporte teórico usando um modelo de classificação linear de uma camada com características de múltiplos componentes. O modelo de dados assume que as características consistem em um componente de nível de bloco compartilhado (classe grosseira) e um componente específico de classe (classe refinada). A análise compara a dinâmica de treinamento do Gradiente Descendente (GD), Adam e Muon sob tamanhos de passo infinitesimais e momentum zero, focando no estado em que cada otimizador atinge pela primeira vez um limite de perda específico.

Principais Contribuições e Resultados

1. Evidência Empírica de Qualidade Superior de Características

O artigo estabelece que o Muon consistentemente aprende características que são mais robustas e transferíveis do que aquelas aprendidas pelo Adam e SGD:

• Robustez: Modelos pré-treinados com Muon alcançam a maior acurácia média no ImageNet-C e a menor perplexidade no FineWeb10B-C em todos os tipos de corrupção e arquiteturas. Notavelmente, o Muon supera o Adam em transformadores de visão e modelos de linguagem, e supera o SGD em modelos de linguagem.
• Transferibilidade: Características pré-treinadas com Muon transferem-se mais efetivamente para tarefas de jusante. Em visão, o Muon iguala ou supera o Adam e o SGD em todos os benchmarks de jusante. Em linguagem, modelos pré-treinados com Muon alcançam menor perplexidade após o ajuste de instrução em comparação com modelos pré-treinados com Adam.

2. Mecanismos de Estados Ocultos

As vantagens empíricas estão diretamente ligadas a propriedades específicas das representações aprendidas:

• Maiores Margens de Logit: A sondagem camada a camada revela que o Muon produz margens de logit significativamente maiores em comparação ao Adam e SGD. Isso indica uma separação mais ampla entre previsões corretas e distratores, correlacionando-se diretamente com a robustez observada.
• Maior Posto Efetivo: A análise espectral mostra que o Muon atinge um posto efetivo mais alto e uma fração de energia Top-10 mais baixa através das camadas. Isso sugere que as representações do Muon espalham a informação de forma mais uniforme através de direções ortogonais, criando um subespaço de características mais rico que facilita uma melhor adaptação a tarefas de jusante.

3. Garantias Teóricas

Em um cenário de características de múltiplos componentes estilizado, os autores provam que a normalização espectral do Muon remove desequilíbrios de valores singulares na matriz de gradiente.

• Vantagem de Margem: O Muon atinge uma margem de classificação estritamente maior do que o Adam e o GD em qualquer perda de treinamento correspondente (fora de uma janela de perda terminal exponencialmente pequena). Isso ocorre porque a normalização espectral do Muon equilibra a contribuição das direções de características ortogonais, enquanto a normalização por coordenada do Adam e o fluxo de gradiente bruto do GD levam à dependência excessiva de componentes específicos.
• Vantagem de Posto Efetivo: O Muon atinge provadamente um posto efetivo de características estritamente superior ao do Adam e GD em perda correspondente, apoiando o achado empírico de representações mais diversas.

Significância e Alegações

O artigo afirma "desmistificar" as vantagens do Muon ao mudar o foco da eficiência de otimização para a qualidade das características. Os autores argumentam que o sucesso do Muon não é meramente um resultado de convergência mais rápida, mas decorre de sua capacidade de aprender representações que são inerentemente mais robustas a perturbações de entrada e mais adaptáveis a novas tarefas.

A significância deste trabalho reside em:

1. Conectar Otimização à Representação: Estabelece uma ligação concreta entre a regra de atualização consciente de matrizes do Muon e as propriedades geométricas das características aprendidas (margens de logit e posto efetivo).
2. Validação Teórica: Fornece a primeira prova teórica de que o Muon induz qualidades de características sistematicamente diferentes e superiores (maiores margens, maior posto efetivo) em comparação ao Adam e GD sob condições de perda correspondente.
3. Implicações Práticas: As descobertas sugerem que, para aplicações que exigem robustez a mudanças de distribuição ou forte transferibilidade para tarefas de jusante, o Muon é uma escolha superior aos otimizadores adaptativos padrão como o Adam, mesmo quando os orçamentos de treinamento são iguais.

Os autores observam modestamente que sua análise atual foca em LLMs e classificadores de visão, deixando o desempenho do Muon em outras famílias de modelos (ex: modelos de difusão) como uma direção para trabalhos futuros.