Resumo Técnico: Treinamento para o Modelo que Você Retorna

Definição do Problema

Os pipelines modernos de Modelos de Linguagem (LM) frequentemente retornam um modelo com média — tipicamente uma Média Móvel Exponencial (EMA) das iterações de treinamento — em vez da própria iteração final. Embora a média de iterações seja conhecida por estabilizar a otimização e melhorar a generalização, uma questão fundamental permanece: Dado que o resultado final será uma média de iterações, como o processo de treinamento em si deve ser modificado para otimizar o desempenho deste estimador específico?

Trabalhos anteriores abordaram isso implicitamente ao observar que a média permite taxas de aprendizado mais agressivas. No entanto, este artigo adota uma abordagem direta: busca projetar um algoritmo de otimização que controle explicitamente a trajetória de treinamento para minimizar o erro do estimador de média de iterações retornado.

Metodologia

Formulação Teórica

Os autores enquadram o problema como um problema de controle ótimo dentro de um modelo quadrático estocástico de tempo contínuo.

• Dinâmica: A trajetória dos parâmetros $\theta_t$ segue uma equação diferencial estocástica (SDE) representando o gradiente descendente ruidoso próximo a um mínimo quadrático (processo Ornstein–Uhlenbeck), aumentada por uma entrada de controle $u_t$.
• Objetivo: O objetivo é encontrar uma estratégia de controle $u_t$ que minimize o erro quadrático médio entre o estimador de média temporal retornado $\hat{\theta}_{EMA, T}$ e o ótimo real $\mu^*$, sujeito a uma penalidade quadrática na magnitude da intervenção de controle (parâmetro de regularização $\lambda$).
• Derivação: Utilizando a teoria de controle Linear-Quadrático-Gaussiano (LQG) e o princípio da equivalência de certeza, os autores derivam um controlador ótimo. A solução exata depende da média posterior do ótimo dado a trajetória, o que é intratável para LMs de grande escala.

Aproximação Prática: PACE

Para tornar a solução implementável, os autores introduzem uma série de aproximações heurísticas:

1. Substituição do Estimador: Substituir a média posterior desconhecida do ótimo pela média atual da iteração ($\hat{\theta}_{EMA, t}$).
2. Aproximação de Tempo: Assumir que o tempo atual $t$ é muito menor que o tempo total de treinamento $T$ ($t \ll T$), simplificando os coeficientes dependentes do tempo.
3. Discretização e Pré-condicionamento: Discretizar o controle contínuo e substituir a matriz Hessiana desconhecida pelo pré-condicionador diagonal usado no AdamW (a estimativa do segundo momento $v_t$).

Isso resulta no PACE (Pullback Averaging Control for Efficient Optimization), um invólucro leve em torno do AdamW. A regra de atualização adiciona um termo de "pullback" (puxada) que move os pesos vivos em direção aos pesos da EMA:
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
onde $C$ é uma matriz diagonal derivada do pré-condicionador Adam, e a força do pullback é limitada para garantir que a atualização permaneça uma combinação convexa do passo padrão e do pullback.

Principais Contribuições

1. Formulação de Controle Ótimo: O artigo fornece a primeira derivação fundamentada de um algoritmo de otimização projetado especificamente para minimizar o erro de um estimador de média de iterações, em vez da iteração final.
2. Garantias Teóricas:
   • Configuração Convexa: Os autores provam que uma versão estilizada do PACE converge para o mínimo de qualquer perda convexa com gradientes estocásticos não enviesados na taxa padrão de $O(1/\sqrt{T})$, até um fator constante dependente da força da EMA. Isso garante que o método nunca seja significativamente pior que o SGD padrão.
   • Configuração Quadrática: Na configuração quadrática usada para a derivação, a dinâmica controlada melhora estritamente o erro quadrático limite do estimador de média de iterações em comparação com a dinâmica não controlada (SGD padrão com EMA). A melhoria pode ser arbitrariamente grande em certas instâncias quadráticas.
3. Validação Empírica: O método é avaliado em:
   • Fine-tuning: Modelos SmolLM2-1.7B, Qwen3-1.7B e Gemma3-1B no dataset smol-smoltalk.
   • Pré-treinamento: GPT-2 (124M) no FineWeb com orçamento de tokens Chinchilla-optimal.
   • Baselines: O PACE é comparado contra o AdamW vanilla, AdamW com avaliação de EMA e o algoritmo Schedule-Free.

Resultados

• Desempenho: O PACE melhora estritamente tanto o AdamW vanilla quanto o AdamW com avaliação de EMA em todos os modelos testados (1B–2B parâmetros) e datasets.
• Robustez: A melhoria é robusta em uma ampla gama de taxas de aprendizado, esquemas de decaimento (constante, cosseno, Warmup-Stable-Decay) e hiperparâmetros (força de pullback $c$, potência $\kappa$ da EMA).
• Comparação com Schedule-Free: O PACE é competitivo com o algoritmo Schedule-Free (Defazio et al., 2024), que também utiliza a média de iterações para remover a necessidade de decaimento da taxa de aprendizado. No entanto, o PACE alcança isso com uma regra de atualização diferente.
• Decaimento da Taxa de Aprendizado: O PACE funciona bem mesmo com taxas de aprendizado constantes, sugerindo que ele pode estabilizar o treinamento em regimes onde o decaimento da taxa de aprendizado é tipicamente necessário. Em comparações de orçamento de tokens, a trajetória de LR constante do PACE acompanha o "fundo do rio" (river floor) do panorama de perda tão efetivamente quanto os ramos WSD, muitas vezes superando-os.

Significância e Limitações

Significância:
O artigo demonstra que otimizar para o modelo retornado (a média) em vez da iteração final leva a ganhos de desempenho tangíveis. Ao tratar a trajetória de treinamento como um problema de controle, o PACE fornece um método teoricamente fundamentado para estabilizar o treinamento e melhorar a generalização sem exigir ajustes complexos de cronograma (schedule). Os autores observam que o PACE oferece convergência teórica melhorada na configuração quadrática em comparação com métodos padrão, embora isso venha com uma garantia de convergência teórica ligeiramente mais lenta no cenário convexo geral devido ao fator de EMA.

Limitações:

• Suposições Teóricas: A derivação baseia-se em um modelo quadrático simplificado com ruído aditivo. Embora a convergência seja provada para funções convexas gerais, a incorporação rigorosa de momentum e pré-condicionamento adaptativo (centrais no AdamW) na análise teórica permanece uma direção aberta.
• Escala: Os experimentos foram limitados a modelos de até 2B parâmetros para fine-tuning e 124M para pré-treinamento. A eficácia em escalas maiores (ex: 70B+) não foi verificada.
• Sobrecarga de Memória: O PACE requer a manutenção de uma cópia adicional dos pesos do modelo (o estado da EMA) durante o treinamento, o que introduz uma sobrecarga de memória não trivial, embora isso seja mitigado por grandes tamanhos de lote na prática.
• Sensibilidade à Semente (Seed): Devido a restrições de computação, as execuções principais usaram uma única semente, embora estudos de robustez com múltiplas sementes tenham confirmado a estabilidade da ordenação do método.

Os autores concluem que o PACE representa um passo promissor em direção ao alinhamento de algoritmos de otimização com a estratégia real de implantação dos LMs modernos, onde o modelo com média é o produto final.