Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

Resumo Técnico: Muon-OGD

Enunciado do Problema
O desafio central abordado neste trabalho é o esquecimento catastrófico em Modelos de Linguagem de Grande Escala (LLMs) durante a aprendizagem contínua. Quando LLMs são ajustados sequencialmente em novas tarefas, seu desempenho em tarefas previamente aprendidas frequentemente degrada-se rapidamente. Métodos existentes de aprendizagem contínua baseados em projeção (por exemplo, OGD, GEM, O-LoRA) mitigam isso restringindo as atualizações de parâmetros a subespaços ortogonais às direções associadas a tarefas passadas. No entanto, esses métodos são tipicamente formulados sob geometria euclidiana de parâmetros, onde as magnitudes das atualizações e a ortogonalidade são governadas pela norma de Frobenius.

Os autores identificam uma incompatibilidade fundamental: otimizadores modernos para parâmetros matriciais, como o recentemente proposto otimizador Muon, operam sob geometria de norma espectral (norma de operador). O Muon aplica atualizações matriciais ortogonalizadas (aproximadas via iterações de Newton–Schulz) e pode ser interpretado como realizando o descenso de maior inclinação sob uma restrição de norma espectral. O artigo argumenta que confiar na geometria de norma de Frobenius para restrições de aprendizagem contínua é subótimo para parâmetros matriciais de LLMs, pois não se alinha com a geometria da dinâmica de otimização subjacente.

Metodologia: Muon-OGD
O artigo propõe o Muon-OGD (Projeção de Gradiente Ortogonal Espectral baseada em Muon), um framework que integra otimização consciente da norma espectral com restrições de aprendizagem contínua baseadas em projeção.

1. Formulação: Em vez do problema padrão restrito por norma de Frobenius, o Muon-OGD formula a atualização $\Delta$ como um problema de otimização restrito por norma espectral:
   $$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{s.t.} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$$
   Aqui, $G$ é a matriz de gradiente (ou momento), $\|\cdot\|_2$ é a norma espectral, $\eta$ é o tamanho do passo e $\{C_i\}$ são direções de restrição (por exemplo, gradientes de tarefas passadas) que devem permanecer ortogonais à atualização.

2. Otimização Dual: Como o problema carece de uma solução de forma fechada, os autores derivam uma formulação dual. O objetivo dual envolve minimizar a norma nuclear de uma matriz deslocada $H = G + \sum \lambda_i C_i$. A atualização ótima é recuperada via a função sinal da matriz: $\Delta^* = -\eta \cdot \text{msgn}(H)$.

3. Algoritmo: O método resolve o problema dual iterativamente usando descenso de subgradiente nas variáveis duais $\lambda$.

   • Loop Interno: Atualiza iterativamente $\lambda$ para impor restrições de ortogonalidade.
   • Aproximação do Sinal da Matriz: A função sinal da matriz $\text{msgn}(H)$ é aproximada eficientemente usando um pequeno número de iterações de Newton–Schulz (especificamente 5 iterações, denotadas como $NS_5$), evitando o custo computacional de SVD explícito.
   • Instanciação de Baixo Rango: Para eficiência, os autores propõem uma instanciação específica onde as restrições são derivadas dos $k$ vetores singulares superiores de tarefas passadas, permitindo que as variáveis duais sejam organizadas em uma matriz de baixo rango $\Lambda$, reduzindo a dimensionalidade do problema dual.

Principais Contribuições

1. Geometria de Norma Espectral para AC: O artigo formula a aprendizagem contínua baseada em projeção sob geometria de norma espectral, alinhando as restrições de estabilidade da aprendizagem contínua com a geometria de norma de operador de otimizadores modernos como o Muon.
2. Regra de Atualização Prática: Deriva um algoritmo prático (Muon-OGD) que combina restrições de ortogonalidade de tarefas passadas com atualizações de sinal de matriz no estilo Muon, resolvendo o problema restrito eficientemente via formulação dual e aproximações de Newton–Schulz.
3. Validação Empírica: O framework é instanciado para ajuste fino contínuo de LLMs e avaliado em diversos benchmarks, demonstrando melhorias consistentes em relação a contrapartes baseadas em norma de Frobenius.

Resultados Experimentais
Os autores avaliam o Muon-OGD em:

• Benchmarks Padrão: Benchmarks de aprendizagem contínua de 5 e 15 tarefas usando T5-Large.
• Benchmark TRACE: Um benchmark para LLMs ajustados por instrução (usando LLaMA-2-7B-Chat) cobrindo habilidades diversas como codificação, matemática e compreensão multilíngue.
• Currículos Específicos de Domínio: Aprendizagem sequencial em Codificação (BigCodeBench) $\to$ Matemática (GSM8K) $\to$ Médico (HuatuoGPT-o1) usando modelos Qwen2.5 e LLaMA3.2.

Descobertas:

• Desempenho: O Muon-OGD supera consistentemente o ajuste fino sequencial (SeqFT) e baselines competitivas de gradiente ortogonal (por exemplo, O-LoRA, Sculpting Subspaces, OSFT) em termos de Precisão Média (AA) e Transferência Reversa (BT).
  • No benchmark de 15 tarefas do T5-Large, o Muon-OGD alcançou 72,0% de precisão média, superando o O-LoRA (69,6%) e o OSFT (71,3%).
  • No benchmark TRACE, o Muon-OGD alcançou a maior AA (49,4%) e a menor taxa de esquecimento (-4,7%) entre os métodos comparados.
• Escalabilidade: O método permanece computacionalmente escalável, aproveitando a eficiência das iterações de Newton–Schulz e representações de restrição de baixo rango.
• Agnosticismo de Arquitetura: A abordagem é eficaz tanto em arquiteturas codificador-decodificador (T5) quanto apenas decodificador (LLaMA, Qwen).

Significado e Alegações
O artigo alega que a geometria de atualização consciente da norma espectral fornece uma alternativa prática e eficaz à projeção de norma de Frobenius para aprendizagem contínua em LLMs. Ao preencher a lacuna entre a geometria de otimizadores modernos (como o Muon) e os requisitos de estabilidade da aprendizagem contínua, o Muon-OGD melhora o compromisso estabilidade-plasticidade. Os autores sugerem que métodos existentes baseados em projeção podem ser sistematicamente atualizados para operar sob geometria de norma espectral para preservar melhor os benefícios de condicionamento e estabilidade observados no Muon, enquanto previnem o esquecimento catastrófico. O trabalho não alega resolver todos os desafios da aprendizagem contínua, mas oferece uma alternativa geométrica principista que produz ganhos empíricos em retenção e adaptação em diversos domínios e escalas de modelos.