Schedules de learning rate e warmup

01 · Conceito

Conceito

Na metade de uma run longa de pretraining, o cluster cresce. O dobro de nós entra, o batch global dobra de quatro milhões de tokens por passo para oito milhões, o throughput melhora lindamente, e todo mundo fica satisfeito. A run termina no prazo em termos de tokens, a curva de loss parece razoável, e o checkpoint final é pior do que deveria — perceptivelmente pior na avaliação downstream do que um piloto menor previu. A causa não são os dados, nem a arquitetura, nem os nós extras. É que o schedule de learning rate foi escrito em steps.

O optimizer transforma gradientes em updates de parâmetro, mas o learning rate define a escala deles, e um único valor constante raramente serve a uma run inteira. O início do treino tem ativações descalibradas e momentos do optimizer pouco confiáveis; o meio quer progresso substancial; o fim se beneficia de updates pequenos e refinadores. Um schedule codifica esse risco variável.

O warmup começa abaixo do pico pretendido e sobe ao longo de um intervalo inicial. O warmup linear no passo $t$ é simplesmente

para $t\le T_{warmup}$. Os primeiros updates ficam contidos enquanto escalas de normalização, magnitudes de gradiente e as estimativas de momento do Adam se assentam. O warmup não conserta um pico excessivo; ele só muda como esse pico é alcançado.

Depois do warmup, o decaimento cosseno é o padrão moderno. Escrevendo $p$ para a fração do horizonte de decaimento completada,

Avalie isso uma vez à mão com um pico de 3e-4 e um piso de 3e-5, um mínimo típico de dez por cento. Na marca de um quarto, $\cos(\pi/4)=0.7071$, então

Na metade do caminho, $\cos(\pi/2)=0$, então

A curva é suave perto das duas pontas e mais íngreme no meio, que é por que ela passa muito tempo perto do pico e depois muito tempo perto do piso.

Agora volte ao cenário de abertura com esses números na mão. Suponha que o plano fosse um orçamento de dois trilhões de tokens a quatro milhões de tokens por passo, dando 500.000 steps, e que o horizonte do scheduler estivesse configurado como 500.000 steps. Dobrar o batch global para oito milhões de tokens significa que os mesmos dois trilhões de tokens são consumidos em 250.000 steps. A run termina quando os dados terminam — em $p=0.5$ — com o learning rate ainda em 1,65e-4, sem nunca ter entrado na fase de decaimento. Todo o refinamento que a segunda metade de um cosseno proporciona simplesmente não aconteceu, e nada nos logs pareceu anormal, porque a taxa seguiu fielmente a curva configurada.

Vale conhecer as alternativas ao cosseno. O decaimento linear cai a um ritmo constante em direção a um ponto final escolhido e torna trivial ler quanto de progresso resta. O decaimento por raiz quadrada inversa, proporcional a $t^{-1/2}$ após uma construção de warmup, foi a receita original do Transformer e tem a propriedade útil de não exigir um horizonte declarado. Constante com warmup mantém o pico depois da rampa, o que serve a um horizonte total incerto, mas pode deixar os updates tardios mais ruidosos do que um modelo quase convergido gostaria. Os pontos finais importam tanto quanto os formatos: decair exatamente até zero pressupõe que a run termina onde foi planejada, e estendê-la depois não produz aprendizado a menos que o schedule seja reescrito, enquanto um piso não nulo preserva adaptação mas continua perturbando um modelo já assentado. Warm restarts, populares em outros contextos, introduzem subidas deliberadas e raramente são apropriados para uma run cara de pretraining de passagem única.

O AdamW separa a otimização baseada em gradiente do weight decay, mas o scheduler e o decay ainda interagem, e parâmetros como escalas de normalização e biases costumam ser excluídos do decay inteiramente. O gradient clipping limita um update dirigido por um gradiente descomunal antes do passo do optimizer; é um guarda-corpo, não substituto de um schedule sensato. Escalar o batch global pode justificar ajustar o pico, mas regras lineares e de raiz quadrada são heurísticas com fronteiras de regime — batches maiores reduzem o ruído de gradiente e podem querer warmup mais longo ou outra afinação do optimizer, então rode varreduras piloto em vez de extrapolar uma regra por ordens de magnitude.

Registre a taxa realizada em cada passo junto com contagem de tokens, loss, norma de gradiente, estatísticas do optimizer e updates pulados de mixed precision. Na retomada, restaure a posição do scheduler: reiniciar o warmup por acidente depois de um checkpoint muda o caminho de otimização, e saltar para uma taxa errada de fim de run fabrica exatamente o tipo de spike que a lição 5.12 precisa diagnosticar. Sobrepor a taxa à curva de loss ajuda a correlacionar instabilidade com escala de update, embora correlação sozinha nunca prove que o schedule causou um spike.

O modelo duradouro é o de momento controlado para o aprendizado. O warmup evita sacudir uma rede descalibrada, a fase central se move rápido por um espaço de parâmetros útil, e o decaimento estreita o passo conforme a run se aproxima do fim planejado. O schedule é parte da especificação do experimento e parte do estado do checkpoint — não configuração cosmética do optimizer, e não algo a expressar numa unidade que muda quando seu cluster muda.