Resumo Técnico: Adam-Decay com Portão Cosseno (CGAD) para DiLoCo Desacoplado

1. Declaração do Problema

Em sistemas assíncronos de DiLoCo Desacoplado (DDiLoCo), os trabalhadores realizam múltiplos passos de otimização local antes de trocar pseudo-gradientes com um servidor de parâmetros global. Um desafio crítico nesta arquitetura é a atraso (staleness): os pseudo-gradientes aplicados na rodada externa $t$ podem ter sido computados $\tau$ rodadas anteriormente devido à latência de rede, trabalhadores lentos (stragglers) ou orçamentos de assincronia.

O otimizador externo padrão para DiLoCo é o SGD com momento de Nesterov. Este método aplica uma etapa de atualização uniforme a cada pseudo-gradiente recebido, sem condicionar à sua idade ($\tau$). À medida que $\tau$ aumenta, o buffer de momento acumula uma visão desatualizada da paisagem de perda. Consequentemente, o otimizador dá passos que ultrapassam direções que não descem mais a perda atual, levando a instabilidade ou divergência, particularmente à medida que a escala do modelo aumenta. Análises teóricas existentes de SGD com momento assíncrono sugerem que as taxas de convergência degradam com um fator proporcional a $\tau_{max}^2$ (o quadrado do atraso máximo).

2. Metodologia: Adam-Decay com Portão Cosseno (CGAD)

Os autores propõem o Adam-Decay com Portão Cosseno (CGAD), uma substituição direta para o otimizador externo de Nesterov que condiciona explicitamente as atualizações à idade do pseudo-gradiente.

Mecanismo Central

O CGAD modifica a atualização padrão do Adam escalando cada pseudo-gradiente recebido $g_f$ para o fragmento $f$ por um fator $\sigma(\tau)$ antes de entrar nos buffers de primeiro e segundo momento do Adam:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Onde:

• Decaimento Exponencial ($e^{-\alpha \tau}$): Modela o decaimento de informação de gradientes desatualizados. À medida que $\tau$ cresce, a contribuição do gradiente para os buffers de momento é reduzida exponencialmente.
• Portão Cosseno ($\gamma(\tau)$): Uma função de corte suave definida como:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  Este portão é igual a 1 para gradientes frescos, decai suavemente para 0 em um corte escolhido $\tau_{cut}$, e zera estritamente contribuições para $\tau > \tau_{cut}$.

Principais Escolhas de Design

1. Portão Pré-Buffer: A escala $\sigma(\tau)$ é aplicada ao gradiente antes de ele atualizar os buffers de primeiro momento ($m$) e segundo momento ($v$). Isso impede que gradientes desatualizados inflacionem o segundo momento ($v$), o que de outra forma atenuaria atualizações frescas subsequentes.
2. Variante por Fragmento (PA-CGAD): Para agendadores de sincronização parcial (por exemplo, Streaming DiLoCo), o método rastreia a idade de fragmentos individuais ($a_f$) e porta atualizações com base em $\max(\tau, a_f)$, permitindo um modelo unificado de atraso tanto para atrasos de rede quanto para atrasos de agendamento.
3. Hiperparâmetros: O método introduz dois hiperparâmetros, $\alpha$ (taxa de decaimento) e $\tau_{cut}$ (corte), com padrões ($\alpha=0.2, \tau_{cut}=32$) que transferem entre escalas de modelo.

3. Contribuições Teóricas

O artigo fornece um limite de convergência não assintótico para uma atualização adaptativa com portão idealizada em objetivos não convexos suaves.

• Independência de Viés de Atraso: O termo de viés de atraso do limite de convergência depende do hiperparâmetro de portão $\alpha$ (especificamente $O(1/e^\alpha)$) em vez do atraso máximo realizado $\tau_{max}$.
• Contraste com Análise Padrão: Análises padrão de SGD com momento assíncrono produzem um termo de complexidade escalando com $\tau_{max}^2$. Ao substituir isso por um termo dependente de $\alpha$, o mecanismo de portão desacopla teoricamente a penalidade de convergência da distribuição de pior caso de atraso, desde que $\tau_{cut} \ge 1/\alpha$.
• Escopo: Os autores observam que isso é uma motivação para o mecanismo de portão, e não um teorema completo de convergência para Adam empírico, pois assume passos normalizados limitados em vez de derivá-los de estatísticas empíricas.

4. Resultados Experimentais

Experimentos foram conduzidos no pré-treinamento de modelos de linguagem estilo Llama em três escalas: 25 M, 1 B e 7 B parâmetros, usando atrasos controlados até $\tau=16$.

Estabilidade e Convergência

• Escala 1 B: Em $\tau=16$, a receita padrão de Nesterov divergiu (perda $\approx 319$), enquanto o CGAD alcançou uma perda de 7.44, comparável ao seu desempenho em $\tau=0$ (7.48). O método de Nesterov perdeu a convergência totalmente em $\tau \ge 8$.
• Escala 25 M: O CGAD e outras bases conscientes de atraso (Adam-Decay, SDM, Nesterov Atrasado) permaneceram estáveis, enquanto Nesterov divergiu em $\tau \ge 4$. Nesta escala, as diferenças entre métodos conscientes de atraso foram marginais.

Seguro de Escala e Robustez

A descoberta mais significativa emerge na escala 7 B sob condições de memória limitada (pesos bf16, AdamW de 8 bits):

• Variância de Semente para Semente: A base mais próxima, Adam-Decay (CGAD sem o corte cosseno), exibiu um aumento de 27 vezes no desvio padrão de semente para semente ($\sigma$) à medida que o modelo escalava de 25 M para 7 B (de 0.07 para 1.89 em $\tau=8$). Em contraste, o $\sigma$ do CGAD permaneceu plano (aprox. 0.38).
• Risco de Disparo Único: Usando "média + $\sigma$" como proxy para risco de implantação de disparo único, o risco do Adam-Decay em 7 B / $\tau=8$ (13.19) cruzou a perda ao acaso (10.83) por 2.4 nats, implicando uma alta probabilidade de falha em uma única execução de produção. O risco do CGAD (11.28) permaneceu bem abaixo deste limite.
• Mecanismo de Falha: Os autores atribuem a falha do Adam-Decay em 7 B ao acúmulo de ruído de arredondamento quadrado no buffer de segundo momento ($v$) de gradientes desatualizados que nunca são totalmente zerados. O corte cosseno no CGAD previne isso ao zerar estritamente contribuições além de $\tau_{cut}$.

Sensibilidade de Ajuste

• Nesterov vs. CGAD: Mesmo quando Nesterov foi reajustado (removendo o momento inteiramente), não pôde igualar o desempenho de uma configuração CGAD ajustada em altos atrasos. Nesterov ajustado exigiu hiperparâmetros diferentes para diferentes regimes de atraso, enquanto o CGAD manteve robustez em toda a faixa.
• Sensibilidade de Hiperparâmetros: O $\alpha$ ótimo muda ligeiramente com a escala (0.2 para 25 M, 0.4 para 1 B em $\tau=8$), mas o padrão $\alpha=0.2$ permanece robusto em todo o espectro de atraso.

5. Significado e Alegações

O artigo alega que o CGAD é a escolha de escalonamento mais robusta para treinamento assíncrono de DiLoCo, abordando especificamente a fragilidade do momento de Nesterov sob grandes atrasos controlados.

• Seguro de Escala: A contribuição primária é o corte cosseno, que atua como "seguro de escala". Enquanto o decaimento exponencial sozinho (Adam-Decay) é suficiente para modelos menores ou configurações de precisão total, o corte rígido é crítico para modelos de fronteira (7 B+) operando com estados de otimizador comprimidos (8 bits/bf16). Isso previne o acúmulo de variância desatualizada que leva a variância catastrófica de semente para semente e falha de disparo único.
• Compatibilidade Direta: O CGAD requer apenas uma modificação de uma linha em implementações existentes do Adam e adiciona dois hiperparâmetros com padrões que transferem entre escalas.
• Impacto Prático: Para treinamento assíncrono de fronteira, o método oferece a diferença entre uma única execução de treinamento que é lançada com sucesso e uma que falha (acima da perda ao acaso).

Os autores permanecem modestos quanto ao escopo teórico, reconhecendo que o limite de convergência se aplica a uma atualização idealizada e que a superioridade empírica é demonstrada especificamente dentro da família DiLoCo e contextos de pré-treinamento estilo Llama. Eles observam que a contagem de passos internos ($H=8$) usada nos experimentos é menor que as configurações de produção de DiLoCo ($H=100-500$), embora estudos de sensibilidade sugiram que a classificação persista em $H$ mais altos.