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

Sintesi Tecnica: Leggi di Scaling con Step di Addestramento e Dimensione del Batch

Enunciato del Problema
Le attuali leggi di scaling per i Large Language Models (LLM), come la legge di Chinchilla, modellano la perdita di test (test loss) in funzione della dimensione del modello ($N$) e del totale dei dati di addestramento ($D$). Sebbene efficaci per determinare l'allocazione ottimale del calcolo tra $N$ e $D$, queste leggi non tengono esplicitamente conto degli iperparametri di addestramento, specificamente la dimensione del batch ($b$) e il numero di step di addestramento ($K$). I precedenti tentativi di modellare lo scaling degli iperparametri spesso falliscono nel integrarsi con le leggi di scaling basate sulla perdita o si affidano al concetto di "dimensione critica del batch" che, in alcune formulazioni teoriche, implica una dimensione ottimale del batch pari a uno — una contraddizione rispetto ai risultati empirici dove le dimensioni ottimali del batch scalano con il budget di token. Inoltre, l'adattamento (fitting) delle leggi di scaling per gli iperparametri ottimali richiede solitamente sweep esaustivi di dimensioni del batch per vari budget di token, comportando costi computazionali massicci.

Metodologia
Gli autori propongono una legge di scaling a tre termini (3TL) che decompone esplicitamente il budget dei dati di addestramento $D$ in dimensione del batch in token ($M = b \times s$, dove $s$ è la lunghezza della sequenza) e step di addestramento ($K$), tale che $D = MK$. La forma funzionale proposta per la perdita $L$ è:

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

dove $(E, A, B, C, \alpha, \beta, \gamma)$ sono parametri derivati dai cicli di addestramento.

Gli autori confrontano questa 3TL con un approccio a due termini (2TL), dove la dimensione del modello $N$ è fissa e la legge viene adattata indipendentemente per ogni $N$ utilizzando la forma $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Le fasi metodologiche chiave includono:

1. Derivazione della Dimensione Ottimale del Batch: Minimizzando la 3TL rispetto a $M$ sotto il vincolo $D=MK$, gli autori derivano una dimensione ottimale del batch $M^\star$ che scala come una legge di potenza di $D$: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Utilizzo dei Dati: A differenza dei metodi precedenti che richiedono l'identificazione dell'unica dimensione del batch ottimale per ogni configurazione (tramite sweep completi), la 3TL viene adattata utilizzando dati provenienti da dimensioni del batch subottimali. Ciò consente di adattare la legge in modo robusto con significativamente meno cicli di addestramento.
3. Modellazione del Regime Subottimale: Per affrontare i vincoli hardware dove la dimensione teorica ottimale del batch è impraticabile, gli autori propongono una procedura di fitting a due stadi. Primo, utilizzano la 3TL per informare una forma funzionale semplificata $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ (assumendo $\gamma \approx \beta$). Secondo, adattano una legge di potenza all'intervallo di dimensioni del batch $\epsilon$-subottimali (dove $\epsilon$ rappresenta una soglia di spreco computazionale del 5%).
4. Analisi della Dimensione Critica del Batch: Gli autori analizzano il comportamento del numero di step $K_{\bar{L}}(b)$ necessari per raggiungere una perdita target $\bar{L}$ in funzione della dimensione del batch, confrontando le previsioni della 3TL con i risultati empirici riguardanti la dimensione critica del batch.

Risultati Chiave
Gli autori hanno valutato il loro approccio su due dataset: Li (Li et al., 2025a) e OpenEuroLLM.

• Scaling della Dimensione Ottimale del Batch: La 3TL recupera con successo lo scaling della dimensione ottimale del batch $M^\star$. Sul dataset Li, lo scaling derivato $M^\star \approx 0.667 \cdot D^{0.566}$ si allinea strettamente con la "Step-Law" ($M^\star \approx 0.58 \cdot D^{0.571}$) di Li et al. (2025a), che è stata derivata da uno sweep completo della dimensione del batch.
• Efficienza Computazionale: Un risultato primario è che la 3TL può essere adattata utilizzando solo il 28% dei cicli di addestramento richiesti dai metodi tradizionali. Mascherando il dataset per includere solo 2 dimensioni del batch per ogni combinazione $(N, D)$ (invece di uno sweep completo), la 3TL recupera comunque la corretta legge di scaling per $M^\star$. Al contrario, adattare direttamente una legge di potenza ai pochi "migliori" batch disponibili (senza modellare la superficie della perdita) risulta in leggi di scaling altamente instabili e imprecise.
• Dimensioni del Batch Subottimali: La procedura di fitting a due stadi modella con successo l'intervallo di dimensioni del batch subottimali che comportano al massimo il 5% di spreco computazionale. Gli autori rilevano che questo intervallo ha una ampiezza di circa $2^2$ (un fattore 4) nella dimensione del batch, e il suo comportamento di scaling è coerente attraverso diverse dimensioni del modello e dataset.
• Dimensione Critica del Batch: La 3TL modella correttamente il fenomeno per cui il numero di step per raggiungere una perdita target diventa meno sensibile all'aumento della dimensione del batch oltre un certo punto critico. Fondamentalmente, il modello mostra che questo comportamento critico scala con il budget di token $D$ ma è ampiamente invariante rispetto alla dimensione del modello $N$, concordando con le osservazioni empiriche di Zhang et al. (2025).
• Confronto con Chinchilla: Quando la 3TL viene ridotta a una forma standard di Chinchilla (sostituendo la $M^\star$ ottimale nell'equazione della perdita), lo scaling computazionale ottimale per la dimensione del modello $N$ implicito differisce da un adattamento diretto della legge di Chinchilla. La 3TL suggerisce un esponente di scaling minore per $D$, evidenziando come la scelta della formulazione della legge di scaling impatti le configurazioni ottimali derivate.

Significato e Rivendicazioni
Il documento afferma che la legge a tre termini offre un quadro unificato che colma il divario tra le leggi di scaling in stile Chinchilla (dimensione del modello vs dati) e le leggi di scaling degli iperparametri (dimensione del batch). La sua importanza risiede in tre aree principali:

1. Robustezza con Dati Sparsi: La capacità di adattare la legge utilizzando dimensioni del batch subottimali riduce drasticamente il costo computazionale per derivare le leggi di scaling, rendendo fattibile determinare le dimensioni ottimali del batch senza sweep esaustivi.
2. Consistenza Teorica: La formulazione è strutturalmente simile ai limiti di convergenza nell'ottimizzazione stocastica, fornendo un ponte teorico tra l'analisi empirica dello scaling e la teoria dell'ottimizzazione.
3. Applicabilità Pratica: A differenza dei modelli precedenti che implicano una dimensione ottimale del batch pari a uno o non riescono a descrivere i regimi subottimali, la 3TL fornisce una dimensione del batch ottimale non banale e modella esplicitamente le prestazioni nei regimi subottimali, il che è critico per l'addestramento pratico sotto vincoli hardware.

Gli autori segnalano limitazioni, tra cui potenziali inconsistenze nei valori dei parametri tra diversi dataset (ad esempio, il termine di perdita irriducibile $E$) e il fatto che la legge richiede ancora tassi di apprendimento (learning rates) ottimali per ogni configurazione. Riconoscono inoltre che lo scaling computazionale ottimale implicito dalla 3TL differisce dai fit diretti di Chinchilla, suggerendo che essa potrebbe non essere l'unico strumento per determinare le dimensioni ottimali del modello, ma è superiore per comprendere l'interazione tra step, dimensione del batch e dati.