Resumo Técnico: OrScale – Otimização Ortogonalizada com Escalonamento de Razão de Confiança por Camada

1. Declaração do Problema

O otimizador Muon emergiu como uma ferramenta prática para treinamento em grande escala, utilizando atualizações ortogonalizadas para corrigir a direção de passos de valor matricial via iterações de Newton–Schulz (aproximando o fator polar). No entanto, o Muon sofre de uma limitação crítica: enquanto otimiza a direção das atualizações, deixa a magnitude controlada por uma única taxa de aprendizado essencialmente global.

Essa escala global é subótima porque diferentes camadas (por exemplo, projeções de atenção, projeções de subida/descida MLP, embeddings) exibem distribuições de gradiente e curvaturas efetivas vastamente diferentes. Embora a variante Moonlight tenha introduzido um fator estático por forma para equalizar o RMS por elemento, ela não consegue se adaptar à deriva dinâmica das escalas das camadas durante o treinamento.

Tentativas anteriores de portar o Escalonamento Adaptativo de Taxa por Camada (LARS) e o Escalonamento Adaptativo de Momento por Camada (LAMB) — que desacoplaram com sucesso os tamanhos dos passos das taxas de aprendizado globais no SGD e no Adam — para o Muon falharam. O artigo identifica que portos ingênuos colapsam devido a três razões estruturais específicas:

1. Denominadores Degenerados: O uso da norma de Frobenius da direção ortogonalizada ($\|Q_\ell\|_F$) como denominador da razão de confiança falha porque é quase uma constante determinada pela forma da matriz (rank), não pela magnitude da atualização.
2. Incompatibilidade de Unidades e Saturação de Corte: O uso da norma bruta do momento ($\|\tilde{M}_\ell\|_F$) introduz uma incompatibilidade de unidades (unidades de gradiente vs. unidades de peso), fazendo com que a razão de confiança atinja o limite superior de corte ($r_{max}$) quase universalmente, tornando a adaptação ineficaz.
3. Fuga de Decaimento de Peso Desacoplado: A combinação de denominadores calibrados com decaimento de peso desacoplado cria um ciclo de realimentação positiva onde as normas de peso crescem descontroladamente (10–50$\times$), desestabilizando o treinamento.

2. Metodologia

Os autores propõem o OrScale, um framework construído sobre um único princípio de design: O denominador de uma razão de confiança por camada deve medir a norma de Frobenius da direção de atualização real que o otimizador subtrairá dos parâmetros.

Princípio Central

Seja $D_{\ell,t}$ a direção real no espaço de parâmetros subtraída dos pesos $W_{\ell,t}$ antes que a taxa de aprendizado global seja aplicada. A razão de confiança é definida como:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Crucialmente, $D_{\ell,t}$ deve incluir tanto o termo de decaimento de peso quanto a direção ortogonalizada.

Duas Variantes

O framework produz duas implementações específicas:

1. OrScale (Padrão Geral/Visão):

   • Direção de Atualização: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (onde $\lambda$ é o decaimento de peso e $Q$ é a direção ortogonalizada).
   • Fator de Forma: $s_\ell = 1$ (sem escalonamento estático).
   • Calibração: Nenhuma.
   • Decaimento de Peso: Acoplado diretamente ao passo de atualização escalado pela razão de confiança.

2. OrScale-LM (Padrão para LLM):

   • Direção de Atualização: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Fator de Forma: Adota o fator estático do Moonlight $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ para corresponder ao RMS por elemento do AdamW.
   • Calibração: Introduz uma constante de calibração preguiçosa por camada ($c_{\text{denom}, \ell}$) calculada uma vez na inicialização. Isso ancora a razão de confiança inicial exatamente em 1 ($r_{\ell,0}=1$).
   • Significado da Calibração: Isso permite que o OrScale-LM herde os hiperparâmetros de taxa de aprendizado e decaimento de peso ajustados para o Moonlight (e, por extensão, o AdamW) sem exigir uma nova varredura de hiperparâmetros. Garante invariância de largura na inicialização.

Detalhes Algorítmicos

• Decaimento de Peso Acoplado: O decaimento de peso é integrado ao cálculo da razão de confiança ($\lambda W$ faz parte do denominador) em vez de ser aplicado como um puxão separado. Isso previne o ciclo de realimentação de fuga observado em variantes desacopladas.
• Eficiência: O método adiciona sobrecarga negligenciável (<1% do tempo de relógio), exigindo apenas um escalar adicional por camada de matriz e dois cálculos de norma de Frobenius por passo, que são dominados pelas iterações de Newton–Schulz existentes.

3. Contribuições Principais

• Algoritmo: Uma receita unificada e leve em parâmetros (Algoritmo 1) que integra com sucesso o escalonamento de razão de confiança por camada no Muon. Os autores identificam o único "cantinho viável" no espaço de design (Denominador de direção de atualização real + Decaimento de Peso Acoplado).
• Análise de Falhas: Um mapeamento rigoroso do espaço de design da razão de confiança do Muon (Tabela 1) que caracteriza analítica e empiricamente três modos de falha (denominador degenerado, saturação de corte e crescimento descontrolado de peso), demonstrando por que falham e por que a solução proposta tem sucesso.
• Teoria:
  • Convergência: Prova uma garantia de convergência não convexa de norma nuclear $O(1/\sqrt{T})$.
  • Ganho Adaptativo: Estabelece uma constante de descida adaptativa por camada $\kappa_{\text{eff}} > 1$ sob heterogeneidade de camada mensurável, melhorando estritamente sobre a constante homogênea do Muon.
  • Separação: Prova uma separação estrita do MuTrust/MuScale sob condições de saturação de corte.
  • Calibração: Demonstra que o OrScale-LM preserva a transferência de taxa de aprendizado estilo $\mu$P na inicialização.
• Empíricos: Avaliação abrangente em tarefas de modelos de visão e linguagem.

4. Resultados Experimentais

Visão (CIFAR-10 / DavidNet)

• Desempenho: O OrScale alcançou 94,05% $\pm$ 0,08 de precisão top-1 de validação, classificando-se em primeiro lugar entre cinco baselines.
• Melhoria: Isso representa uma melhoria de +0,35 pontos sobre o Muon padrão (93,70%) e +0,30 sobre o Muon+Moonlight (93,75%).
• Estabilidade: O OrScale demonstrou uma janela de taxa de aprendizado estável mais ampla em comparação com o Muon, consistente com as previsões teóricas de ganhos adaptativos por camada.
• Ablação: Um estudo de ablação de modo de falha confirmou que variantes correspondentes aos modos de falha identificados (por exemplo, usando $\|Q_\ell\|_F$ ou momento bruto) exibiram os sintomas previstos (por exemplo, saturação de corte de 99,5–100% ou fuga de norma de peso) e falharam em igualar o desempenho do OrScale.

Grandes Modelos de Linguagem (Pré-treinamento FineWeb-Edu)

• Escalonamento: Avaliado em quatro escalas de modelo (125M a 1,1B de parâmetros).
• Comparação: O OrScale-LM superou o Muon+Moonlight em três das quatro escalas (125M, 545M, 1,1B) e venceu o AdamW em todas as escalas.
• Redução de Perda: Na escala de 1,1B, o OrScale-LM alcançou uma entropia cruzada final de validação de 2,6251, comparado a 2,6360 para o Muon+Moonlight e 2,7304 para o AdamW.
• Eficiência: A vantagem foi preservada em toda a faixa de computação, com o OrScale-LM mostrando uma vitória limpa no regime onde a computação mais importa para LLMs.

5. Significado e Alegações

O artigo afirma que o OrScale conecta com sucesso a lacuna entre a geometria espectral do Muon e a adaptabilidade por camada do LARS/LAMB. Seu significado principal reside em:

1. Resolver o Problema da Magnitude: Fornece o primeiro controle de magnitude variável no tempo e por camada para o Muon que é robusto a modos de falha estruturais.
2. Transferibilidade Prática: Através da variante OrScale-LM e da calibração preguiçosa, permite que praticantes transfiram taxas de aprendizado ajustadas para AdamW ou Moonlight diretamente para o OrScale-LM sem reajuste, reduzindo a barreira de adoção.
3. Solvência Teórica: Oferece uma justificativa teórica rigorosa para por que tentativas híbridas anteriores falharam e prova que o método proposto produz constantes de convergência estritamente melhores sob heterogeneidade.
4. Validação Empírica: Demonstra que essas melhorias teóricas se traduzem em ganhos tangíveis tanto em tarefas de visão em pequena escala quanto em pré-treinamento de LLMs em grande escala, estabelecendo o OrScale como um otimizador de última geração para camadas de valor matricial.

Os autores mantêm modéstia em relação aos resultados de LLM, observando que a vantagem na escala de 1,1B é comparável ao ruído de uma única semente em escalas menores, e, portanto, os resultados devem ser vistos como evidência direcional da eficácia do método, em vez de uma superioridade estatística definitiva em todos os regimes. No entanto, a melhoria consistente através das escalas e a ablação bem-sucedida dos modos de falha apoiam fortemente o princípio de design proposto.