Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs

Título: Previsão de Curvas de Reavaliação de Treinamento (TREC) Habilita Currículos de Dados Eficazes para LLMs

Autores: Shane Bergsma, Nolan Dey e Joel Hestness (Cerebras Systems)

1. O Problema

O treinamento de Grandes Modelos de Linguagem (LLMs) frequentemente incorpora fases de "currículo de dados" ou "annealing" (resfriamento), onde dados de alta qualidade (HQ), específicos de domínio (ex: matemática, código) ou recentes são inseridos nas etapas finais do pré-treinamento.

• Hipótese Comum: Acredita-se que colocar esses dados no final do treinamento, quando a taxa de aprendizado (Learning Rate - LR) está próxima de zero, maximiza sua eficácia.
• O Desafio: Não há princípios claros sobre a colocação ótima de dados. A prática atual depende de heurísticas (colocar no final) ou ablações caras e demoradas para encontrar o melhor momento de inserção. Além disso, a ordem dos dados durante o treinamento influencia o que o modelo retém e esquece, um fenômeno relacionado ao "esquecimento catastrófico".

2. Metodologia e Conceitos Fundamentais

Os autores introduzem uma nova ferramenta diagnóstica chamada Curva de Reavaliação de Treinamento (TREC - Training Re-evaluation Curve) e um método para prevê-la proativamente.

A. Definição de TREC

A TREC mede o desempenho de um modelo totalmente treinado (com pesos finais $\theta_T$) quando avaliado sequencialmente nos batches de treinamento originais ($B_1, \dots, B_T$).

• Cálculo: $L_{re}(t) = L(B_t; \theta_T)$.
• Interpretação: Se o modelo retivesse todos os dados igualmente, a curva seria plana. Na prática, a curva varia, mostrando em quais momentos do treinamento o modelo "lembra" melhor dos dados. Um valor de perda (loss) mais baixo na TREC indica que os dados daquele momento foram bem retidos pelo modelo final.
• Observação Chave: Em muitos casos, a TREC atinge seu mínimo (melhor retenção) antes do final do treinamento, especialmente sob esquemas de LR com "step decay" (queda abrupta), contradizendo a intuição de colocar dados HQ apenas no final.

B. A Hipótese Principal

Colocar dados de alta qualidade (HQ) nos pontos onde a TREC está no seu mínimo (onde a retenção é máxima) maximiza o desempenho em tarefas alvo, superando a colocação no final do treinamento.

C. Previsão da TREC (Sem Treinar o Modelo Duas Vezes)

Como calcular a TREC requer treinar o modelo e depois reavaliá-lo (custo proibitivo), os autores propõem um modelo preditivo baseado na dinâmica do otimizador AdamW:

1. Escala de Tempo do AdamW (Timescale $\tau$): Os parâmetros do AdamW podem ser vistos como uma Média Móvel Exponencial (EMA) das atualizações de peso. O tempo de escala $\tau = \frac{1}{\eta \lambda}$ (onde $\eta$ é a LR e $\lambda$ é o weight decay) controla o quanto os dados antigos influenciam os pesos finais.
2. Deriva do Minimizador (Minimizer Drift): A eficácia de um gradiente diminui se o "mínimo local" da função de perda para aquele batch específico se mover à medida que o treinamento avança.
3. Fórmula Preditiva: A forma da TREC pode ser prevista combinando os coeficientes da EMA do AdamW ($c(\hat{t})$) com um termo de ajuste baseado na fração de treinamento ($\hat{t}^m$):
   $$\hat{L}_{re}(\hat{t}) = 1 - c(\hat{t})^p \cdot \hat{t}^m$$
   Onde $m$ é um expoente que depende do Tokens-Per-Parameter (TPP) e da escala de tempo $\tau$.

3. Contribuições Principais

1. Introdução da TREC: Uma métrica diagnóstica que revela a relação entre a posição dos dados no treinamento e a retenção pelo modelo final.
2. Estudo em Grande Escala: Análise de 600 curvas TREC em modelos de 111M a 3.9B de parâmetros e datasets de 20 a 1280 TPP.
3. Modelo Preditivo: Demonstração de que a forma da TREC é governada principalmente pela escala de tempo do AdamW ($\tau$) e pode ser prevista analiticamente antes do treinamento.
4. Aplicações Práticas:
   • Explicação de resultados de trabalhos anteriores (ex: por que o Llama-3 405B não se beneficiou de annealing no GSM8k).
   • Melhoria no Continual Pre-Training (CPT) de um modelo de 3.9B parâmetros.
   • Análise de arquiteturas MoE (Mixture-of-Experts), mostrando que a TREC reflete o TPP efetivo dos especialistas.

4. Resultados Experimentais

• Validação da Hipótese de Posicionamento: Em experimentos com modelos de 610M parâmetros, inserir dados de código (HQ) na posição onde a TREC era mínima resultou na melhor perda de validação, superando tanto a colocação no final quanto a mistura uniforme.
• Generalização de Escala: A forma da TREC alinha-se através de diferentes escalas de modelo (de 111M a 3.3B) quando a escala de tempo $\tau$ e o TPP são mantidos constantes.
• Previsão Acurada: O modelo preditivo (baseado em $\tau$ e TPP) consegue prever a forma da TREC com alta correlação (Pearson $r_p \approx 98\%$) em modelos maiores, mesmo sendo ajustado em modelos menores.
• Caso de Uso em CPT (3.9B): Ao aplicar dados matemáticos de alta qualidade no meio do CPT (onde a TREC predita estava baixa) em vez de no final, houve melhoria consistente no desempenho, independentemente da taxa de aprendizado usada no CPT (desde que não fosse excessivamente alta).
• Análise de Receitas Existentes:
  • Llama-3 405B: A TREC predita mostrou que a fase de annealing no final do treinamento ocorria quando a taxa de aprendizado era tão baixa que os coeficientes da EMA eram quase zero, resultando em pouca retenção dos dados inseridos.
  • OLMo-2 e Feng et al.: As posições de inserção de dados nessas obras alinham-se com os vales preditos pela TREC, validando a abordagem.

5. Significado e Impacto

• Fim das Heurísticas Cegas: A TREC oferece uma base teórica e prática para substituir a regra de "colocar dados no final" por uma estratégia baseada na dinâmica real do otimizador.
• Economia de Recursos: Permite o design proativo de currículos de dados sem a necessidade de treinar múltiplos modelos para ablações.
• Compreensão de Dinâmicas de Treinamento: Revela que a retenção de dados é uma função complexa da taxa de aprendizado, weight decay, tamanho do lote e escala de tempo do otimizador, e não apenas do momento cronológico.
• Otimização de CPT: Fornece diretrizes claras para Continual Pre-Training, indicando que inserir dados novos no momento de maior sensibilidade do modelo (vale da TREC) é mais eficaz do que no final.

Conclusão

O trabalho demonstra que a Curva de Reavaliação de Treinamento (TREC) é uma ferramenta poderosa para entender e otimizar a colocação de dados em LLMs. Ao prever onde o modelo retém melhor os dados com base nos parâmetros do AdamW, os pesquisadores podem posicionar dados de alta qualidade estrategicamente, melhorando o desempenho final de forma mais eficiente e fundamentada do que as práticas atuais baseadas em heurísticas.