How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

Resumo Técnico: Leis de Escalonamento com Passos de Treinamento e Tamanho de Lote

Enunciado do Problema
As leis de escalonamento existentes para Grandes Modelos de Linguagem (LLMs), como a lei Chinchilla, modelam a perda de teste do modelo como uma função do tamanho do modelo ($N$) e do total de dados de treinamento ($D$). Embora eficazes para determinar a alocação ideal de computação entre $N$ e $D$, essas leis não contabilizam explicitamente os hiperparâmetros de treinamento, especificamente o tamanho do lote ($b$) e o número de passos de treinamento ($K$). Tentativas anteriores de modelar o escalonamento de hiperparâmetros frequentemente falham ao não integrar-se com as leis de escalonamento baseadas em perda ou ao depender do conceito de "tamanho de lote crítico" que, em algumas formulações teóricas, implica um tamanho de lote ideal de um — uma contradição aos achados empíricos onde tamanhos de lote ideais escalam com o orçamento de tokens. Além disso, ajustar leis de escalonamento para hiperparâmetros ideais normalmente requer varreduras exaustivas de tamanhos de lote para vários orçamentos de tokens, incorrendo em custos computacionais massivos.

Metodologia
Os autores propõem uma lei de escalonamento de três termos (3TL) que decompõe explicitamente o orçamento de dados de treinamento $D$ em tamanho de lote em tokens ($M = b \times s$, onde $s$ é o comprimento da sequência) e passos de treinamento ($K$), tal que $D = MK$. A forma funcional proposta para a perda $L$ é:

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

onde $(E, A, B, C, \alpha, \beta, \gamma)$ são parâmetros ajustados a partir de execuções de treinamento.

Os autores comparam esta 3TL contra uma abordagem de lei de dois termos (2TL), onde o tamanho do modelo $N$ é fixo, e a lei é ajustada independentemente para cada $N$ usando a forma $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Principais etapas metodológicas incluem:

1. Derivação do Tamanho de Lote Ideal: Ao minimizar a 3TL em relação a $M$ sujeito à restrição $D=MK$, os autores derivam um tamanho de lote ideal $M^\star$ que escala como uma lei de potência de $D$: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Utilização de Dados: Diferente de métodos anteriores que exigem a identificação do único tamanho de lote ideal para cada configuração (via varreduras completas), a 3TL é ajustada usando dados de tamanhos de lote subótimos. Isso permite que a lei seja robustamente ajustada com significativamente menos execuções de treinamento.
3. Modelagem de Regimes Subótimos: Para lidar com restrições de hardware onde o tamanho de lote teoricamente ideal é inviável, os autores propõem um procedimento de ajuste em duas etapas. Primeiro, eles usam a 3TL para informar uma forma funcional simplificada $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ (assumindo $\gamma \approx \beta$). Segundo, eles ajustam uma lei de potência ao intervalo de tamanhos de lote $\epsilon$-subótimos (onde $\epsilon$ representa um limiar de desperdício de computação de 5%).
4. Análise do Tamanho de Lote Crítico: Os autores analisam o comportamento do número de passos $K_{\bar{L}}(b)$ necessários para atingir uma perda alvo $\bar{L}$ como uma função do tamanho do lote, comparando as previsões da 3TL com achados empíricos sobre o tamanho de lote crítico.

Resultados Principais
Os autores avaliaram sua abordagem em dois conjuntos de dados: Li (Li et al., 2025a) e OpenEuroLLM.

• Escalonamento do Tamanho de Lote Ideal: A 3TL recupera com sucesso o escalonamento do tamanho de lote ideal $M^\star$. No conjunto de dados Li, o escalonamento derivado $M^\star \approx 0.667 \cdot D^{0.566}$ alinha-se de perto com a "Lei de Passo" ($M^\star \approx 0.58 \cdot D^{0.571}$) de Li et al. (2025a), que foi derivada de uma varredura completa de tamanho de lote.
• Eficiência Computacional: Um achado primário é que a 3TL pode ser ajustada usando apenas 28% das execuções de treinamento exigidas pelos métodos tradicionais. Ao mascarar o conjunto de dados para incluir apenas 2 tamanhos de lote por combinação de $(N, D)$ (em vez de uma varredura completa), a 3TL ainda recupera a lei de escalonamento correta para $M^\star$. Em contraste, ajustar diretamente uma lei de potência para os poucos "melhores" tamanhos de lote disponíveis (sem modelar a superfície de perda) resulta em leis de escalonamento altamente instáveis e imprecisas.
• Tamanhos de Lote Subótimos: O procedimento de ajuste de duas etapas modela com sucesso o intervalo de tamanhos de lote subótimos que incorrem em no máximo 5% de desperdício de computação. Os autores descobrem que este intervalo tem uma largura de aproximadamente $2^2$ (um fator de 4) no tamanho do lote, e seu comportamento de escalonamento é consistente entre diferentes tamanhos de modelo e conjuntos de dados.
• Tamanho de Lote Crítico: A 3TL modela corretamente o fenômeno onde o número de passos para atingir uma perda alvo torna-se menos sensível ao aumento do tamanho do lote além de um ponto crítico. Crucialmente, o modelo mostra que este comportamento crítico escala com o orçamento de tokens $D$, mas é amplamente invariante ao tamanho do modelo $N$, correspondendo às observações empíricas de Zhang et al. (2025).
• Comparação com Chinchilla: Quando a 3TL é reduzida a uma forma Chinchilla padrão (ao substituir o $M^\star$ ideal de volta na equação de perda), o escalonamento de computação ideal implícito para o tamanho do modelo $N$ difere de um ajuste direto da lei Chinchilla. A 3TL sugere um expoente de escalonamento menor para $D$, destacando que a escolha da formulação da lei de escalonamento impacta as configurações ideais derivadas.

Significância e Alegações
O artigo afirma que a lei de três termos oferece um framework unificado que faz a ponte entre as leis de escalonamento estilo Chinchilla (tamanho do modelo vs. dados) e as leis de escalonamento de hiperparâmetros (tamanho de lote). Sua significância reside em três áreas principais:

1. Robustez com Dados Esparsos: A capacidade de ajustar a lei usando tamanhos de lote subótimos reduz drasticamente o custo computacional de derivar leis de escalonamento, tornando viável determinar tamanhos de lote ideais sem varreduras exaustivas.
2. Consistência Teórica: A formulação é estruturalmente semelhante aos limites de convergência em otimização estocástica, fornecendo uma ponte teórica entre a análise de escalonamento empírico e a teoria de otimização.
3. Aplicabilidade Prática: Diferente de modelos anteriores que implicam um tamanho de lote ideal de um ou falham em descrever regimes subótimos, a 3TL fornece um tamanho de lote ideal não trivial e modela explicitamente o desempenho em regimes subótimos, o que é crítico para o treinamento prático sob restrições de hardware.

Os autores observam limitações, incluindo possíveis inconsistências nos valores dos parâmetros entre diferentes conjuntos de dados (por exemplo, o termo de perda irredutível $E$) e o fato de que a lei ainda requer taxas de aprendizado ideais para cada configuração. Eles também reconhecem que o escalonamento de computação ideal implícito pela 3TL difere dos ajustes diretos de Chinchilla, sugerindo que ela pode não ser o único instrumento para determinar tamanhos de modelos ideais, mas é superior para entender a interação entre passos, tamanho de lote e dados.