Resumo Técnico: LoRA-Muon: Descida do Gradiente Estepante Espectral na Variedade de Baixo Rank

1. Definição do Problema

A Adaptação de Baixo Rank (LoRA) é amplamente adotada para o ajuste fino de grandes modelos de aprendizado profundo devido à sua eficiência na redução de FLOPs de treinamento e requisitos de memória. No entanto, otimizadores padrão como o AdamW, quando aplicados independentemente aos fatores de LoRA ($A$ e $B$), introduzem dificuldades significativas de ajuste em comparação ao treinamento denso. Esses problemas incluem:

• Sensibilidade à Inicialização: O desempenho é altamente dependente das escolhas de inicialização.
• Baixa Transferibilidade de Hiperparâmetros: As taxas de aprendizado ideais não se transferem bem entre diferentes ranks, larguras, profundidades ou convenções de escala de fatores.
• Dinâmicas Divergentes: Mesmo partindo de pesos compostos idênticos ($W = AB^T$), as dinâmicas de treinamento podem divergir se as escalas dos fatores se afastarem.
• Desempenho Subótimo: O pré-treinamento nativo de baixo rank muitas vezes tem dificuldade em igualar os baselines densos sem estabilização adicional.

O desafio central é projetar um otimizador para a variedade de baixo rank que mantenha as propriedades geométricas dos otimizadores de rank total (especificamente Muon e a família Shampoo) e garanta que as taxas de aprendizado ideais sejam transferíveis através de rank, largura, profundidade e reescalonamento de fatores.

2. Metodologia

2.1 Fundamentação Geométrica

O artigo reformula o design de otimizadores como um problema geométrico: identificar a norma sob a qual um otimizador realiza o gradiente estepante.

• Otimizador Muon: Realiza o gradiente estepante sob a norma espectral ($\|\cdot\|_{2\to2}$) no espaço de matriz ambiente $\mathbb{R}^{m \times n}$. A direção de atualização é dada pela função de sinal de matriz, $-\text{msign}(G)$.
• Variedade de Baixo Rank (Low-Rank Manifold): O LoRA restringe os pesos a uma variedade $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$. Os autores propõem resolver o problema do gradiente estepante diretamente no espaço tangente desta variedade, $T_W \mathcal{M}_r$, em vez de otimizar os fatores $A$ e $B$ como coordenadas independentes.

2.2 Derivação do LoRA-Muon

Os autores derivam o LoRA-Muon resolvendo o problema de minimização linear restrita no espaço tangente de baixo rank sob a norma espectral.

1. Decomposição do Espaço Tangente: As atualizações em $W$ são decompostas em $\Delta W = \Delta A B^T + A \Delta B^T$.
2. Divisão da Região de Confiança (Trust Region Splitting): Para tornar o problema tratável, o orçamento da região de confiança $\eta$ é dividido igualmente entre os dois componentes tangentes. Isso é justificado pela observação empírica de que os dois componentes se alinham rapidamente durante o treinamento, tornando a desigualdade triangular quase estrita.
3. Formulação do Projetor: Usando invariância unitária, o problema é reduzido à busca de Oráculos de Minimização Linear (LMOs) em gradientes projetados. A atualização ideal é:
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   onde $P_A$ e $P_B$ são projetores nos espaços coluna de $A$ e $B$.
4. Especialização Espectral: Para a norma espectral, o LMO é $-\text{msign}(\cdot)$. As atualizações de fator resultantes (Equações 14–15) envolvem o branqueamento (whitening) dos gradientes dos fatores pelas matrizes de Gram atuais ($S_A = A^TA, S_B = B^TB$) e a aplicação da função de sinal espectral:
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 Decaimento de Peso Dividido (Split Weight Decay)

O decaimento de peso padrão por fator leva a um descompasso entre as dinâmicas de atualização de rank total e baixo rank. Os autores derivam uma regra de decaimento de peso dividida onde o fator de decaimento $s = \sqrt{1 - \lambda \eta}$ é aplicado assimetricamente:
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
Isso garante que a atualização do peso composto corresponda à dinâmica desejada de rank total, preservando a transferibilidade de hiperparâmetros.

2.4 Invariância de Gauge

Uma propriedade crítica do LoRA-Muon é a invariância de gauge. A atualização no peso composto $W$ permanece invariante sob transformações invertíveis arbitrárias $(A, B) \to (AR, BR^{-T})$. Isso é alcançado porque a atualização depende apenas dos projetores $P_A$ e $P_B$, que são, por si só, invariantes de gauge.

3. Principais Contribuições

1. Derivação do LoRA-Muon: Uma nova família de otimizadores derivada pela resolução do problema de gradiente estepante na variedade de baixo rank sob normas unitariamente invariantes, especializando-se para a norma espectral.
2. Regra de Decaimento de Peso Dividido: Uma regra de atualização específica que alinha as dinâmicas de treinamento de baixo rank com as de rank total, permitindo a transferência de hiperparâmetros.
3. Prova de Invariância de Gauge: Prova teórica de que a atualização LoRA-Muon é invariante sob a ação do grupo de gauge completo, garantindo que o otimizador dependa da atualização de peso induzida, e não de coordenadas de fatores arbitrárias.
4. Comparação com o Spectron: Demonstração de que o Spectron não é invariante de gauge e é sensível a reescalonamentos arbitrários de fatores, ao contrário do LoRA-Muon.
5. Relação com o LoRA-RITE: Prova de que o núcleo de coordenadas QR simplificado do LoRA-RITE é algebricamente idêntico à atualização espectral do LoRA-Muon, embora o LoRA-Muon evite o overhead computacional de decomposições QR e armazenamento de segundo momento.
6. Validação Empírica: Evidência de que o LoRA-Muon serve como um proxy eficaz para a busca de taxa de aprendizado do Muon denso.

4. Resultados Experimentais

Experimentos foram conduzidos em modelagem de linguagem ao nível de caractere usando o dataset TinyShakespeare com baselines densos e de baixo rank com computação equivalente.

• Transferência de Taxa de Aprendizado: A taxa de aprendizado ideal para o Muon denso ($lr = 0.1$) transfere-se perfeitamente para o LoRA-Muon através de varreduras de rank (2, 4, 8, 16, 32), largura, profundidade e reescalonamento de fator escalar.
  • Em Rank 2, o LoRA-Muon recupera a melhor taxa de aprendizado testada para o modelo denso.
  • Em Rank 32, o LoRA-Muon atinge uma perda de validação média menor ($1.776 \pm 0.002$) do que o baseline denso ($1.789 \pm 0.002$) na varredura de média de sementes.
• Rebalanceamento de Gauge: Quando os fatores são inicializados com um grande desequilíbrio escalar (escala 9), aplicar um passo de rebalanceamento tem um efeito negligenciável no LoRA-Moun (mudança de perda $\approx -5.1 \times 10^{-5}$), confirmando a invariância de gauge. Em contraste, o desempenho do Spectron muda significativamente ($2.0 \times 10^{-2}$), confirmando sua sensibilidade à escala de fatores.
• Eficiência: Um proxy de rank-2 usa apenas ~4,3% dos parâmetros treináveis do modelo denso, enquanto suporta um orçamento de passo de otimizador ~9,2x maior sob as mesmas restrições de computação.

5. Significância e Alegações

O artigo alega que o LoRA-Muon fornece uma solução teoricamente fundamentada para as dificuldades de ajuste do LoRA. Ao tratar o LoRA como um problema de otimização na variedade de baixo rank, em vez de um sistema de coordenadas fatorizado, os autores alcançam:

• Robustez: O otimizador é invariante à escolha arbitrária da fatoração (simetria de gauge).
• Transferibilidade: As taxas de aprendizado ideais (e, por extensão, outros hiperparâmetros via leis de escala) transferem-se perfeitamente através de rank, largura, profundidade e escalas de fator.
• Eficiência: Ele computa a atualização espectral sem decomposições QR caras ou armazenamento de segundo momento, tornando-o amigável para aceleradores.

Os autores posicionam o LoRA-Muon como um "bom proxy de baixo rank" para o Muon de rank total e para a família de otimizadores Shampoo. Eles enfatizam que a parametrização de baixo rank não deve ditar a geometria do otimizador; em vez disso, a atualização deve ser definida pela própria geometria da variedade de matriz de baixo rank.

Limitações: Os autores observam que as alegações empíricas são atualmente estabelecidas em modelos de escala TinyShakespeare e ainda não foram validadas em pré-treinamento de larga escala ou benchmarks de ajuste fino de downstream. O estudo foca na transferência de taxa de aprendizado, não em uma varredura completa de hiperparâmetros.