Resumo Técnico: A Média Dual Otimista Unifica Otimizadores Modernos

Declaração do Problema

A otimização em aprendizado profundo evoluiu ao longo de dois eixos principais: geometria e adaptação (projetar atualizações com normas implícitas ou pré-condicionadores que refletem a estrutura do modelo, por exemplo, Adam, Muon) e a composição de ingredientes (combinar feedback de gradiente, momento e cronogramas para estabilidade, por exemplo, Lion, métodos sem cronograma). Embora muitos otimizadores de última geração (como Muon, Lion, AdEMAMix e NAdam) pareçam distintos, eles carecem de um quadro teórico unificado que explique sua relação e forneça orientação principista para o ajuste de hiperparâmetros.

Um desafio específico e persistente é o ajuste do decaimento de peso. No treinamento de múltiplas épocas, o decaimento de peso atua como um regularizador que restringe as normas das iterações. No entanto, mesmo no treinamento de uma única época, onde o sobreajuste é menos preocupante, o decaimento de peso frequentemente permanece benéfico. Regras recentes de escalonamento sugerem definir o decaimento de peso proporcional a $1/d$ (dimensão do modelo), mas desvendar a dependência da dimensão do modelo $d$ versus o horizonte de treinamento $n$ é difícil. Além disso, os valores ótimos de decaimento de peso variam significativamente entre diferentes otimizadores (por exemplo, o Lion requer um decaimento muito maior que o AdamW), necessitando de um ajuste custoso por otimizador.

Metodologia: SODA

Os autores introduzem o SODA (Média Dual Otimista Estocástica), uma generalização da Média Dual Otimista (ODA) que unifica o processamento dual dos gradientes e o processamento primal das iterações.

Quadro Teórico

O SODA generaliza a atualização padrão de média dual:
$$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$$
introduzindo uma sequência de extrapolação primal $(y_k)$ e uma variável dual otimista $(\bar{m}_{k+1})$. O algoritmo mantém:

1. Processamento Dual: Uma média de gradiente baseada em momento $m_{k+1}$ e uma média otimista $\bar{m}_{k+1}$ que incorpora o gradiente atual $\nabla f(y_k, \xi_k)$.
2. Processamento Primal: Uma atualização de iteração $x_{k+1}$ e um ponto extrapolado $y_{k+1}$, controlados pelos parâmetros $\lambda_k$ e $\bar{\lambda}_k$.

As regras de atualização são:
$$\begin{aligned} m_{k+1} &= (1 - \alpha_k)m_k + \alpha_k \nabla f(y_k, \xi_k) \\ \bar{m}_{k+1} &= (1 - \bar{\alpha}_k)m_{k+1} + \bar{\alpha}_k \nabla f(y_k, \xi_k) \quad (\text{otimismo}) \\ z_{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}_{k+1}) \\ x_{k+1} &= (1 - \lambda_k)x_k + \lambda_k z_{k+1} \\ y_{k+1} &= (1 - \bar{\lambda}_k)x_{k+1} + \bar{\lambda}_k z_{k+1} \quad (\text{extrapolação primal}) \end{aligned}$$
onde $h_k$ é um regularizador (mapa espelho) que define a geometria (por exemplo, $\ell_\infty$ para Lion, espectral para Muon).

Unificação de Otimizadores Existentes

O artigo demonstra que, ao selecionar parâmetros e geometrias específicos, o SODA recupera vários otimizadores modernos como casos especiais:

• NAdam / Lion / Muon: Estes são recuperados sob o regime "Modernizado" onde $\bar{\lambda}_k = 0$ (sem extrapolação primal para o ponto de avaliação do gradiente) e $\bar{\alpha}_k > 0$ (otimismo).
  • Lion: Usa geometria $\ell_\infty$ ($h_k = K^*$) e parâmetros de momento específicos.
  • Muon: Usa geometria espectral e momento de Nesterov ($\bar{\alpha}_k = \alpha_k$).
  • NAdam: Usa uma geometria $\ell_\infty$ suavizada com escalonamento adaptativo.
• A Tabela 1 no artigo mapeia explicitamente esses métodos para suas escolhas correspondentes de média de gradiente ($\alpha_k$), otimismo ($\bar{\alpha}_k$) e geometria ($h_k$).

O Wrapper SODA

Uma contribuição prática chave é um Wrapper SODA (Algoritmo 1) que pode envolver qualquer otimizador base (por exemplo, Adam, Lion, Muon) sem introduzir novos hiperparâmetros.

• Configuração: Define $\bar{\lambda}_k = 0$ (correspondendo ao regime "modernizado"), usa média uniforme de iterações $\lambda_k = 1/(k+2)$ e escala o tamanho do passo $\gamma_k = (k+2)\eta_k$.
• Eliminação do Decaimento de Peso: O wrapper ancora as atualizações à inicialização $z_0$ em vez da origem. A regra de atualização resultante implementa efetivamente um cronograma de decaimento de peso de $\mu_k = 1/(k+2)$.
• Interpretação: Os autores argumentam que, neste quadro, o decaimento de peso não é meramente um regularizador, mas uma forma de média primal. Este decaimento $1/k$ fundamentado teoricamente elimina a necessidade de ajustar um hiperparâmetro separado de decaimento de peso.

Contribuições Principais

1. Unificação via Otimismo e Dualização: O SODA fornece um único modelo que recupera Muon, Lion, NAdam e outros como instâncias otimistas da Média Dual, esclarecendo suas relações estruturais.
2. Garantias Teóricas: O quadro deriva escolhas de hiperparâmetros (especificamente o cronograma de decaimento $1/k$) que são fundamentadas teoricamente na análise de convergência para funções convexas $L$-suaves. Reenquadra o decaimento de peso como um mecanismo para média primal.
3. Wrapper Sem Parâmetros: O wrapper SODA proposto remove a necessidade de ajustar o decaimento de peso para qualquer otimizador base, mantendo ou melhorando o desempenho.
4. Validação Empírica: Experimentos extensivos mostram melhorias consistentes através de tamanhos de modelo e horizontes de treinamento.

Resultados Experimentais

Os experimentos foram conduzidos no nanoGPT usando o conjunto de dados FineWeb100, comparando otimizadores envoltos pelo SODA contra bases ajustadas.

• Transferência através do Horizonte: O decaimento de peso ótimo para o Muon foi encontrado diminuindo à medida que o horizonte de treinamento aumentava, contradizendo regras de escalonamento estáticas (como $1/d$). O decaimento dinâmico $1/(k+2)$ do wrapper SODA superou consistentemente o Muon, mesmo quando o Muon recebeu seu decaimento de peso melhor ajustado (Figura 1).
• Aplicabilidade Geral: O wrapper melhorou o desempenho para Adam, Muon e Scion através de diferentes horizontes de treinamento (1$\times$ e 4$\times$ passos Chinchilla) sem qualquer ajuste adicional (Figura 2).
• Escalonamento: Ao escalar o tamanho do modelo de 64M para 1B de parâmetros (seguindo o escalonamento Chinchilla), o SODA (especificamente a variante otimista SODA$^\dagger$) superou consistentemente os otimizadores base (Scion/uScion), com a lacuna de desempenho ampliando-se à medida que o tamanho do modelo aumentava (Figura 4).
• Simplificação: A configuração SODA$^\dagger$ simplificou a base Muon substituindo Adam por Lion para todas as camadas e usando um único hiperparâmetro de momento, ainda assim alcançando desempenho competitivo ou superior.

Significado e Alegações

O artigo alega que o SODA oferece um mecanismo principista para transferir estratégias de otimização através de diferentes horizontes de treinamento e escalas de modelo sem a necessidade de ajuste de hiperparâmetros.

• Reenquadramento do Decaimento de Peso: Os autores propõem uma mudança de perspectiva: o decaimento de peso deve ser entendido não apenas como um regularizador contra o sobreajuste, mas como uma forma de média primal que ancora a trajetória de otimização à inicialização.
• Impacto Prático: Ao unificar otimizadores modernos sob o quadro SODA, os autores fornecem um wrapper "plug-and-play" que elimina o processo tedioso e caro de ajustar o decaimento de peso, o que é particularmente crítico à medida que os tamanhos dos modelos tornam o ajuste manual proibitivamente caro.
• Modéstia: Os autores reconhecem limitações, notando que sua análise teórica está atualmente restrita ao cenário convexo (embora empiricamente informativa para aprendizado profundo) e foca no treinamento de uma única época. Eles sugerem que o treinamento de múltiplas épocas ainda pode exigir um hiperparâmetro para a força de regularização devido a preocupações com sobreajuste.

Em resumo, o SODA unifica o espaço de design de otimizadores modernos através da lente da média dual otimista e fornece um método fundamentado teoricamente e sem parâmetros para melhorar a estabilidade e o desempenho do treinamento em todas as escalas.