Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

1. Declaração do Problema

O artigo aborda a análise teórica de convergência do otimizador Shampoo Estilo AdamW, um algoritmo de última geração que venceu a competição de treinamento de redes neurais AlgoPerf. Embora o Shampoo (e suas variantes como SOAP e Muon) tenha demonstrado desempenho empírico superior aos métodos de pré-condicionamento diagonal (como Adam) ao utilizar pré-condicionadores de matriz completa, suas garantias teóricas de convergência para otimização não convexa têm sido ausentes ou limitadas.

Especificamente, a literatura existente enfrentou três lacunas principais:

1. Não Convexidade: A maioria das análises anteriores restringia-se a cenários convexos.
2. Escopo do Pré-condicionamento: Análises não convexas anteriores (por exemplo, Xie et al., 2026) limitavam-se ao pré-condicionamento unilateral, ignorando o pré-condicionamento bilateral mais complexo e eficaz utilizado em implementações práticas.
3. Componentes Práticos: Teorias existentes frequentemente omitiam recursos práticos críticos, como momento e decaimento de peso desacoplado (estilo AdamW), que são essenciais para o treinamento moderno de modelos de linguagem grandes (LLM).

O artigo visa estabelecer uma taxa de convergência rigorosa para o algoritmo geral Shampoo Estilo AdamW, unificando tanto o pré-condicionamento unilateral quanto o bilateral sob um framework não convexo.

2. Metodologia

Os autores analisam o Algoritmo 1, que representa o otimizador Shampoo Estilo AdamW. A regra de atualização envolve:

• Médias Móveis Exponenciais (EMA): Rastreamento dos primeiros momentos ($M_k$) e segundos momentos ($L_k, R_k$) dos gradientes.
• Pré-condicionamento Bilateral: A etapa de atualização é $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$, onde $L_k$ e $R_k$ são produtos externos acumulados de gradientes (aproximando $E[GG^T]$ e $E[G^TG]$).
• Expoentes Unificados: Os parâmetros $p$ e $q$ (onde $1/p + 1/q = 1$) permitem que o algoritmo alterne entre pré-condicionamento bilateral ($p, q < \infty$) e unilateral ($p=1, q=\infty$ ou vice-versa).

Inovações Técnicas Chave:
Para lidar com a complexidade do pré-condicionamento matricial bilateral em cenários não convexos, os autores desenvolveram técnicas de prova inovadoras:

1. Desigualdade de Hölder de Schatten-p: Em vez das desigualdades de Hölder baseadas em vetores padrão usadas na análise do Adam, os autores utilizam a desigualdade de Hölder de Schatten-p para limitar a norma nuclear ($\|\cdot\|_*$) do gradiente. Isso é crucial porque o pré-condicionador atua sobre os valores singulares dos parâmetros matriciais.
2. Desigualdade de Cauchy-Schwarz Matricial: Para limitar a norma espectral da etapa de atualização ($L^{-1/2p} M R^{-1/2q}$), eles empregam uma desigualdade de Cauchy-Schwarz específica para matrizes para lidar com a interação entre os pré-condicionadores esquerdo e direito.
3. Análise de Piso Espectral: Eles introduzem uma condição (Condição 4) referente ao limite inferior dos pré-condicionadores ($\hat{\epsilon}$). Eles argumentam que, na prática, o piso espectral efetivo dos pré-condicionadores é significativamente maior que a regularização numérica $\epsilon$ (por exemplo, $10^{-12}$), alinhando os limites teóricos com observações empíricas.

3. Contribuições Principais

• Taxa de Convergência Unificada: O artigo estabelece a primeira taxa de convergência para o Shampoo Estilo AdamW em otimização não convexa que abrange tanto o pré-condicionamento unilateral quanto o bilateral.
• Derivação de Taxa Ótima: Eles provam que a norma nuclear esperada média do gradiente converge a uma taxa de:
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  onde $m, n$ são as dimensões da matriz, $K$ é a contagem de iterações e $C$ corresponde à constante na taxa ótima do SGD.
• Justificativa Teórica para Cenários Práticos: A análise explica por que definir o regularizador numérico $\epsilon$ para valores extremamente pequenos (por exemplo, $10^{-12}$) não prejudica a convergência. A teoria mostra que o piso do pré-condicionador efetivo é determinado pela escala de ruído do gradiente, e não pelo $\epsilon$ numérico.
• Relação de Normas: Os autores demonstram que a taxa de norma nuclear derivada é análoga à taxa de norma de Frobenius ótima do SGD ($O(K^{-1/4})$) sob a condição ideal onde $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. Resultados

• Limites Teóricos:
  • O limite do pior caso depende de $\sqrt{m+n}$, que é significativamente mais apertado que a dependência de $\sqrt{mn}$ (dimensão total $d$) encontrada em alguns métodos diagonais ou limites inferiores de matriz completa.
  • A taxa de convergência corresponde ao limite inferior conhecido para otimização estocástica não convexa ($O(K^{-1/4})$) até um fator de $\sqrt{m+n}$.
• Validação Empírica (Pré-treinamento GPT-2):
  • Razão da Norma do Gradiente: Experimentos no pré-treinamento do GPT-2 mostraram que a razão $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ permanece consistentemente próxima do limite superior teórico $\sqrt{\min(m,n)}$. Isso valida a suposição de que a taxa de norma nuclear é efetivamente equivalente à taxa de norma de Frobenius na prática.
  • Piso Espectral: Os autovalores mínimos dos pré-condicionadores ($L_k, R_k$) foram observados sendo ordens de magnitude maiores que o $\epsilon$ numérico ($10^{-12}$) e comparáveis à escala dependente de ruído $\hat{\sigma}^2/(m+n)$. Isso confirma que a suposição "não prática" de um piso espectral maior se mantém no treinamento real de LLMs.

5. Significado

• Ponte entre Teoria e Prática: Este trabalho fornece a primeira base teórica rigorosa para a variante específica do Shampoo que venceu a competição AlgoPerf, explicando por que ela funciona tão bem na prática.
• Avanço sobre Métodos Diagonais: Valida teoricamente a superioridade do pré-condicionamento não diagonal (matriz completa) sobre métodos diagonais (como Adam) para parâmetros com estrutura matricial, mostrando que as correlações cruzadas entre parâmetros capturadas pelo Shampoo levam a direções de convergência melhores.
• Orientação para Hiperparâmetros: A análise oferece justificativa teórica para a escolha de hiperparâmetros (por exemplo, a relação entre o momento $\theta$ e o decaimento do segundo momento $\beta$, e o papel de $\epsilon$), guiando praticantes no ajuste desses otimizadores para modelos em grande escala.
• Fundação para Futuros Otimizadores: Ao unificar o pré-condicionamento unilateral e bilateral e incorporar momento/decaimento de peso, este framework estabelece um novo padrão para analisar otimizadores avançados semelhantes a segunda ordem em aprendizado profundo.

Em resumo, o artigo unifica com sucesso a análise teórica do Shampoo, provando que ele alcança uma taxa de convergência ótima para problemas não convexos enquanto captura as complexas estruturas matriciais inerentes aos modelos modernos de aprendizado profundo.