Resumo Técnico: Sobre a Não Linearidade da Escala da Taxa de Aprendizado para Treinamento de LLMs

Declaração do Problema

Treinar grandes modelos de linguagem (LLMs) de forma eficiente exige um ajuste cuidadoso de hiperparâmetros, mas realizar varreduras diretas de taxas de aprendizado em escalas alvo é frequentemente proibitivamente caro. Consequentemente, profissionais dependem da transferência de hiperparâmetros, prevendo configurações ideais para modelos de grande escala com base em execuções de menor escala. Uma suposição prevalente neste domínio é que a taxa de aprendizado ideal ($\eta^*$) segue uma relação de escalonamento log-linear (lei de potência) em relação à escala de dados ($D$) e ao tamanho do modelo ($N$):
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
No entanto, estudos existentes frequentemente falham em avaliar rigorosamente se essa suposição log-linear é válida globalmente ou apenas localmente, e frequentemente omitem o custo computacional necessário para ajustar essas leis de escala. Este artigo investiga a validade dessa lei de escala e a eficiência de diferentes estratégias de extrapolação (eixo de dados vs. eixo de modelo) e parametrizações (taxa de aprendizado vs. taxa de aprendizado efetiva).

Metodologia

Os autores conduziram um estudo empírico sistemático utilizando arquiteturas do tipo GPT-2 (variando de 22M a 707M de parâmetros) treinadas no conjunto de dados FineWeb-100B com orçamentos de tokens de 5B a 100B.

1. Design Experimental:

   • Os modelos foram treinados usando AdamW com um tamanho de lote fixo (0,52M de tokens) e um cronograma de warmup-estabilidade-decaimento (WSD).
   • As taxas de aprendizado foram varridas sobre uma grade logarítmica ($\log_2 \eta \in [-16, -7]$).
   • Taxa de Aprendizado Ideal ($\eta^*$): Estimada através do ajuste de um polinômio cúbico à perda de validação como uma função de $\log \eta$.
   • Taxa de Aprendizado Efetiva ($\eta_{eff}$): Definida como o tamanho do passo no espaço de pesos normalizado: $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, onde $\hat{w} = w/\|w\|_2$. A taxa de aprendizado efetiva ideal ($\eta^*_{eff}$) foi similarmente estimada.

2. Métricas de Avaliação:

   • $R^2$ Fora da Distribuição ($R^2_{OOD}$): Mede a precisão preditiva das leis de escala em escalas não vistas.
   • Razão de Computação Extra (ECR): Quantifica o custo computacional adicional incorrido ao usar uma taxa de aprendizado prevista versus a verdadeira ótima para atingir a mesma perda de validação.
   • Extrapolação Restrita por Computação: Os autores compararam a extrapolação ao longo do eixo de dados (variando $D$ com $N$ grande e fixo) versus o eixo do modelo (variando $N$ com $D$ grande e fixo) sob idênticos orçamentos de computação.

3. Análise Teórica:

   • O artigo analisa a dinâmica implícita do AdamW, especificamente a relação entre a taxa de aprendizado bruta $\eta$ e a taxa de aprendizado efetiva $\eta_{eff}$, considerando o comportamento transitório das normas de peso.
   • Experimentos foram repetidos usando o AdamH, um otimizador que controla explicitamente $\eta_{eff}$, para isolar os efeitos da dinâmica da norma de peso.

Principais Resultados

1. Não Linearidade da Taxa de Aprendizado Ideal

O estudo descobre que a taxa de aprendizado ideal $\eta^*$ não segue uma lei de escala log-linear global.

• Curvatura Ascendente: À medida que a escala de dados $D$ e o tamanho do modelo $N$ aumentam, $\eta^*$ exibe uma curvatura ascendente sistemática. A extrapolação log-linear a partir de escalas menores consistentemente subestima a taxa de aprendizado ideal para escalas maiores.
• Linearidade da Taxa de Aprendizado Efetiva: Em contraste, a taxa de aprendizado efetiva ($\eta^*_{eff}$) exibe um comportamento de escalonamento substancialmente mais linear através dos mesmos intervalos, com coeficientes de correlação de Pearson mais altos.

2. Superioridade da Extrapolação do Eixo de Dados

Ao comparar estratégias de extrapolação sob orçamentos de computação fixos:

• Eixo de Dados vs. Eixo do Modelo: O ajuste de leis de escala ao longo do eixo do tamanho dos dados ($D$) produz previsões significativamente mais precisas do que ao longo do eixo do tamanho do modelo ($N$).
• Eficiência de Computação: O uso da extrapolação do eixo de dados com a taxa de aprendizado efetiva incorre em apenas ~2% de computação adicional em relação ao uso da verdadeira taxa de aprendizado ideal. Isso é menos de 20% do computação necessária para uma varredura completa de hiperparâmetros na escala alvo.

3. Explicação Mecanística

O artigo atribui a não linearidade em $\eta^*$ à dinâmica transitória das normas de peso no AdamW:

• Dois Regimes: A relação entre $\eta$ e $\eta_{eff}$ depende de se a norma de peso $\|w_t\|$ atingiu o equilíbrio.
  • Regime de Equilíbrio: $\eta_{eff} \propto \sqrt{\eta}$ (inclinação 2 no espaço log-log).
  • Regime Pré-equilíbrio: $\eta_{eff} \propto \eta$ (inclinação 1 no espaço log-log).
• Impacto da Escala: À medida que a taxa de aprendizado ideal diminui para modelos/dados maiores, o horizonte de treinamento frequentemente permanece no regime pré-equilíbrio (onde as normas de peso ainda não relaxaram). Isso altera a dinâmica do tamanho de passo efetivo, causando a curvatura ascendente observada em $\eta^*$.
• Validação via AdamH: Ao usar o AdamH para controlar explicitamente $\eta_{eff}$, a lei de escala torna-se estritamente log-linear, confirmando que a não linearidade no AdamW padrão surge do mapeamento implícito e variante no tempo entre $\eta$ e $\eta_{eff}$.

Significância e Alegações

O artigo reivindica três contribuições primárias:

1. Correção Empírica: Demonstra que a lei de escala log-linear padrão para taxas de aprendizado é apenas uma propriedade local e falha em escalas maiores devido à curvatura ascendente, levando a extrapolações imprecisas.
2. Estratégias de Transferência Práticas: Identifica duas escolhas específicas que melhoram substancialmente a transferência de hiperparâmetros:
   • Parametrizar a lei de escala usando a taxa de aprendizado efetiva ($\eta_{eff}$) em vez da taxa de aprendizado bruta.
   • Extrapolar ao longo do eixo de dados ($D$) em vez do eixo do tamanho do modelo ($N$).
3. Insight Teórico: Fornece uma explicação mecanística para a não linearidade observada, ligando-a ao comportamento transitório das normas de peso no AdamW. Os autores argumentam que $\eta_{eff}$ é a quantidade mais fundamental que governa a dinâmica de otimização em arquiteturas invariantes de escala.

Os autores observam limitações, incluindo o foco em arquiteturas GPT-2 e AdamW, a dependência de ajuste cúbico para a estimativa de $\eta^*$, e o fato de a análise estar limitada a escalas de até 100B de tokens e 707M de parâmetros. Eles não afirmam que essas descobertas se generalizam automaticamente para todas as arquiteturas ou otimizadores sem verificação adicional.