Resumo Técnico: Correção do Viés de Atualização Estocástica em Otimizadores de Modelos de Linguagem Pré-condicionados

1. Declaração do Problema

Otimizadores pré-condicionados (por exemplo, AdamW, Sophia, Shampoo) são centrais para o treinamento de grandes modelos de linguagem (LLMs). Esses métodos estimam uma direção de descida e um termo de geometria (pré-condicionador) a partir de mini-lotes para aproximar a atualização ideal da população:
$$u^\star(\theta) = P(\theta)^{-1}\nabla f(\theta)$$
Na prática, os otimizadores aplicam uma atualização estocástica usando estimativas de mini-lote $\hat{g}$ e $\hat{P}$:
$$\hat{u} = \hat{P}^{-1}\hat{g}$$

O artigo identifica que, mesmo quando $\hat{g}$ e $\hat{P}$ são estimadores não viesados do gradiente da população e do pré-condicionador, o produto $\hat{P}^{-1}\hat{g}$ é geralmente um estimador viesado da atualização ideal. Esse viés de amostra finita surge de duas fontes distintas:

1. Viés de Acoplamento Gradiente–Pré-condicionador: Quando $\hat{g}$ e $\hat{P}$ são estimados a partir do mesmo mini-lote, eles são estatisticamente dependentes. Esse acoplamento introduz um termo de momento cruzado que se desvia do campo vetorial populacional alvo.
2. Viés de Inversão de Pré-condicionador: A inversão de matriz (ou operações de raiz inversa) é uma transformação não linear. Pelo princípio da desigualdade de Jensen, mesmo que $\hat{P}$ seja não viesado para $P$, a expectativa da inversa $\mathbb{E}[\hat{P}^{-1}]$ não é igual à inversa da expectativa $(\mathbb{E}[\hat{P}])^{-1}$. Esse viés é particularmente pronunciado quando as estimativas do pré-condicionador são ruidosas, agressivas ou de valor matricial.

Esses vieses não são erros de implementação, mas consequências inerentes da aplicação de regras adaptativas não lineares a estimativas estocásticas finitas.

2. Metodologia

Os autores propõem uma estrutura de correção de viés de único lote que aborda ambas as fontes de viés sem exigir o retorno ao mesmo valor de parâmetro (o que é impraticável no treinamento de redes neurais). O método opera em um passo de parâmetro fixo $\theta_t$ dividindo o lote de treinamento atual em grupos de micro-lotes.

2.1 Pré-condicionamento Cruzado (Endereçando o Acoplamento)

Para eliminar o acoplamento gradiente–pré-condicionador, o método estima o numerador (gradiente) e o denominador (pré-condicionador) a partir de grupos de micro-lotes independentes.

• Seja $A$ e $B$ dois grupos independentes amostrados da mesma distribuição de treinamento.
• O gradiente é estimado a partir de $A$: $\hat{g}_A$.
• O pré-condicionador é estimado a partir de $B$: $\hat{P}_B$.
• A atualização torna-se $\hat{u}_{cf} = \hat{P}_B^{-1}\hat{g}_A$.
• Devido à independência, $\mathbb{E}[\hat{P}_B^{-1}\hat{g}_A] = \mathbb{E}[\hat{P}_B^{-1}]\mathbb{E}[\hat{g}_A]$, removendo efetivamente o termo de viés de acoplamento.

2.2 Inversão Corrigida por Variância (Endereçando a Não Linearidade)

Para reduzir o viés introduzido pela não linearidade da inversão, o método utiliza a variabilidade dentro do grupo de estimativa do pré-condicionador $B$.

• O grupo $B$ é dividido em $m$ micro-lotes $B_1, \dots, B_m$.
• Caso Diagonal (AdamW, Sophia): Os autores aplicam uma correção pelo método delta. Para uma entrada escalar do pré-condicionador $p$, o funcional inverso $T(p) = (p+\lambda)^{-1}$ é expandido até a segunda ordem. O termo de viés principal, proporcional à variância da média do micro-lote, é estimado e subtraído:
  $$\tilde{T}(\bar{p}) = \frac{1}{\bar{p} + \lambda} - \frac{\widehat{\text{Var}}(\bar{p})}{(\bar{p} + \lambda)^3}$$
  Essa correção é aplicada elemento a elemento a pré-condicionadores diagonais.
• Caso Matricial (Shampoo): Para pré-condicionadores de valor matricial, o método aplica a correção na base de autovetores do pré-condicionador médio $\bar{P}$. Os pré-condicionadores de micro-lote são projetados nessa base, e a correção escalar pelo método delta é aplicada aos autovalores (ou autovalores de raiz inversa) antes de reconstruir a matriz.
• Alternativa Jackknife: Uma correção jackknife mais geral de "deixar um de fora" (LOO) também é descrita, que cancela o termo de viés principal $O(1/m)$, embora possa ser computacionalmente cara para pré-condicionadores matriciais.

3. Contribuições Principais

1. Decomposição Estatística: O artigo decompõe explicitamente o erro de atualização estocástica pré-condicionada em dois vieses de amostra finita: acoplamento gradiente–pré-condicionador e a não linearidade da inversão.
2. Estrutura de Correção: Propõe uma estrutura prática de único passo combinando pré-condicionamento cruzado (para remover o acoplamento) e inversão corrigida por variância (para remover o viés de inversão) usando a variabilidade de micro-lotes.
3. Instanciação Empírica: A estrutura é instanciada para três famílias principais de otimizadores:
   • AdamW: Baseado em momentos diagonais (usando ajuste cruzado LOO e correção Jensen/delta).
   • Sophia: Baseado em curvatura diagonal.
   • Shampoo: Pré-condicionado por matriz (usando correção na base de autovetores).

4. Resultados Experimentais

A estrutura foi avaliada no Qwen2.5-0.5B sob dois regimes:

1. Pré-treinamento: A partir de inicialização aleatória em sequências FineWeb-Edu (limpas e de qualidade mista/ruidosas).
2. Ajuste de Instrução: Ajuste fino supervisionado (SFT) em dados no estilo Alpaca.

Principais Descobertas:

• Desempenho de Pré-treinamento: A correção de viés reduziu consistentemente a perda de pré-treinamento retida (entropia cruzada) em todas as três famílias de otimizadores:
  • AdamW: Redução de perda de 0,149 nats (de 4,836 para 4,687).
  • Sophia: Redução de perda de 0,070 nats (de 6,665 para 6,595).
  • Shampoo: Redução de perda de 0,110 nats (de 5,792 para 5,681).
  • No pré-treinamento de qualidade mista (ruidoso), a variante corrigida do AdamW mostrou um ganho menor, mas positivo (0,019 nats) em dados de avaliação limpos, enquanto ajustava menos agressivamente segmentos corrompidos.
• Ajuste de Instrução: Os resultados foram "neutros a positivos".
  • AdamW: Mostrou pequenos ganhos (por exemplo, -0,001 nats) após o reajuste das taxas de aprendizado.
  • Sophia: Quase empatou com a otimização padrão em conjuntos pequenos, ligeiramente melhor em conjuntos maiores de 5K.
  • Shampoo: Eficazmente neutro no cenário SFT apenas com atenção.
• Abações: O artigo observa que a divisão ingênua de dois do lote (usando metade do lote para o denominador) pode ser excessivamente conservadora devido ao aumento do ruído. A construção proposta de "deixar um de fora" (LOO) para AdamW preserva a escala do denominador enquanto remove o acoplamento, provando-se crítica para o desempenho.

5. Significado e Alegações

O artigo posiciona este trabalho como um mecanismo prático para melhorar otimizadores pré-condicionados ao abordar vieses estatísticos inerentes a atualizações de amostra finita, em vez de propor uma nova arquitetura de otimizador ou aproximação de curvatura.

• Insight Teórico: Estabelece que atualizações adaptativas padrão são estimadores viesados da direção pré-condicionada da população devido ao acoplamento e à não linearidade.
• Utilidade Prática: A correção é leve para otimizadores diagonais e extensível a otimizadores matriciais. Melhora a eficiência de amostragem e a estabilidade, particularmente em regimes de pré-treinamento onde as estimativas são ruidosas e as execuções de treinamento são longas.
• Escopo Moderado: Os autores são modestos quanto à generalizabilidade do tamanho do efeito. Reconhecem que os resultados são baseados em uma única escala de modelo (0,5B) e execuções de semente única. Observam que, em execuções curtas de ajuste de instrução iniciadas a partir de checkpoints pré-treinados fortes, o custo de variância da divisão de amostras pode compensar os benefícios de redução de viés, levando a resultados neutros.
• Conclusão: A contribuição principal é demonstrar que corrigir o viés de amostra finita em atualizações estocásticas pode tornar a otimização de modelos de linguagem pré-condicionados mais eficaz, particularmente no pré-treinamento, sem exigir mudanças fundamentais na pegada de memória do otimizador ou na estrutura de atualização.