Resumo Técnico: Seesaw: Acelerando o Treinamento ao Equilibrar o Escalonamento da Taxa de Aprendizado e do Tamanho do Lote

Declaração do Problema

O pré-treinamento de Grandes Modelos de Linguagem (LLMs) é cada vez mais limitado pelo tempo de relógio (wall-clock time), muitas vezes estendendo-se por vários meses para modelos de última geração. Uma estratégia primária para reduzir essa duração é aumentar o tamanho do lote (batch size), o que aproveita o cálculo paralelo para reduzir o número total de etapas de otimização sequenciais. No entanto, simplesmente aumentar o tamanho do lote indefinidamente leva a retornos decrescentes além de um "tamanho de lote crítico" (CBS), onde a eficiência de amostragem cai.

Embora execuções de treinamento de larga escala recentes (ex: LLaMA, Nemotron, OLMo) tenham empregado esquemas de "rampa de lote" (batch ramp)—aumentando gradualmente o tamanho do lote ao longo do treinamento—essas estratégias são atualmente ajustadas de forma heurística. Há uma falta de justificativa teórica sobre como equilibrar otimamente a queda da taxa de aprendizado com a rampa de aumento do tamanho do lote, particularmente para otimizadores adaptativos como o Adam. A questão central abordada é: Qual é o cronograma ideal de tamanho de lote para minimizar o tempo de execução serial sem sacrificar o desempenho do modelo?

Metodologia

Estrutura Teórica

Os autores desenvolvem um framework principista baseado em regressão linear ruidosa para estabelecer uma equivalência entre a queda da taxa de aprendizado e a rampa de aumento do tamanho do lote.

1. Equivalência de SGD: O artigo fornece a primeira prova de amostra finita (não assintótica) demonstrando que, para o Gradiente Descendente Estocástico (SGD) em regressão linear ruidosa, um processo onde o tamanho do lote dobra a cada fase enquanto a taxa de aprendizado permanece fixa é equivalente (até um fator constante no risco excessivo) a um processo onde o tamanho do lote é fixo e a taxa de aprendizado é reduzida pela metade a cada fase, desde que ambos os processos consumam o mesmo número total de pontos de dados.
2. Extensão para SGD Normalizado (NSGD): Para preencher a lacuna para otimizadores adaptativos como o Adam, os autores analisam o SGD Normalizado, um substituto analítico tratável para o Adam. Eles introduzem a Assunção 2, que postula que, no regime dominado pela variância observado na prática, as normas dos quadrados esperados dos gradientes são dominadas por ruído aditivo (escalando como $O(1/B)$). Sob esta assunção, o NSGD efetivamente se reduz ao SGD com uma taxa de aprendizado reescalonada.
3. A Condição de Equivalência: A análise deriva uma relação específica necessária para manter a dinâmica de perda ao ajustar ambos os parâmetros. Para um fator de queda da taxa de aprendizado $\alpha$ e um fator de aumento do tamanho do lote $\beta$, os processos permanecem equivalentes se o produto $\alpha\sqrt{\beta}$ for mantido constante.

O Algoritmo Seesaw

Baseado na equivalência teórica, os autores propõem o Seesaw, um escalonador que equilibra dinamicamente a taxa de aprendizado e o tamanho do lote:

• Mecanismo: Sempre que um escalonador padrão (ex: decaimento de cosseno) reduziria a taxa de aprendizado por um fator $\alpha$, o Seesaw, em vez disso, reduz a taxa de aprendizado por $\sqrt{\alpha}$ e aumenta o tamanho do lote por $\alpha$.
• Restrição: Os autores identificam um limite teórico para evitar a divergência. O esquema de rampa mais agressivo que mantém a estabilidade é definido por $\alpha = \sqrt{\beta}$. Exceder isso (ou seja, aumentar o tamanho do lote de forma muito agressiva em relação à queda da taxa de aprendizado) leva à divergência.
• Implementação: O Seesaw é projetado como um substituto direto para escalonadores existentes. Na prática, ele aproxima o decaimento de cosseno usando o decaimento por etapas (step decay), disparando aumentos no tamanho do lote nos contagens específicas de tokens onde a taxa de aprendizado seria, de outra forma, cortada.

Principais Contribuições

1. Equivalência Teórica: O artigo estabelece a primeira equivalência não assintótica entre a queda da taxa de aprendizado e a rampa de aumento do tamanho do lote para SGD e estende isso para o SGD Normalizado sob um regime dominado pela variância.
2. O Escalonador Seesaw: Introdução de um algoritmo principista que substitui o padrão de decaimento de taxa de aprendizado por um cronograma acoplado de taxa de aprendizado e tamanho de lote, teoricamente fundamentado para preservar a dinâmica de perda.
3. Análise de Aceleração: Os autores derivam um limite superior teórico para a redução do tempo de execução serial. Ao utilizar a rampa de aumento mais agressiva e estável ($\alpha = \sqrt{\beta}$) contra um baseline de decaimento de cosseno, a aceleração máxima teórica é calculada como $(1 - 2/\pi) \approx 36,3\%$.

Resultados Empíricos

Os autores avaliaram o Seesaw em modelos de 150M, 300M e 600M de parâmetros treinados na escala Chinchilla ($D=20N$) utilizando o codebase do OLMo.

• Paridade de Desempenho: O Seesaw iguala a dinâmica de perda de validação dos esquemas de recozimento de cosseno (cosine annealing) padrão em todos os tamanhos de modelo quando treinado no tamanho de lote crítico.
• Redução de Tempo de Execução: O Seesaw alcança uma redução significativa no tempo de execução serial, reduzindo o tempo de relógio em aproximadamente 36%, aproximando-se de perto do limite teórico derivado na análise.
• Robustez: O método funciona efetivamente com AdamW e várias configurações de decaimento de peso (weight decay).
• Limites de Agressividade: Experimentos variando os parâmetros $\alpha$ e $\beta$ confirmam a restrição teórica. Esquemas que são mais agressivos do que $\alpha = \sqrt{\beta}$ (por exemplo, dobrando o tamanho do lote enquanto mantém a taxa de aprendizado constante ou decaindo-a muito lentamente) resultam em divergência ou desempenho subótimo.
• Limite do Tamanho de Lote Crítico: Os autores observam que a estratégia Seesaw é eficaz até o tamanho de lote crítico. Além deste ponto, onde o termo de ruído não domina mais a norma do gradiente, a equivalência falha e o Seesaw performa pior do que o decaimento de cosseno padrão.

Significância e Alegações

O artigo alega fornecer uma fundamentação teórica rigorosa para a prática de "rampa de lote" (batch ramp), movendo-a de um ajuste heurístico para uma estratégia de otimização principista. Ao demonstrar que a queda da taxa de aprendizado e a rampa de aumento do tamanho do lote são intercambiáveis sob condições específicas, o Seesaw oferece um método prático para acelerar o pré-treinamento de LLMs sem comprometer a qualidade do modelo.

Os autores enfatizam que o Seesaw serve como um substituto direto para escalonadores existentes em otimizadores adaptativos. A principal significância reside na capacidade de reduzir o tempo de execução de treinamento de grandes modelos em cerca de um terço, abordando o gargalo de hardware onde os requisitos de treinamento estão superando as melhorias de hardware. O trabalho também esclarece os limites desta estratégia, observando que ela depende do regime dominado pela variância e pode não se aplicar uma vez que os tamanhos de lote excedam o limiar crítico onde os termos de gradiente médio dominam.