Towards Engineering Scaling Laws with Pretraining Data Composition

Sintesi Tecnica: Verso l'Ingegneria delle Leggi di Scalabilità con la Composizione dei Dati di Pretraining

Definizione del Problema
Le leggi di scalabilità neurale descrivono come le prestazioni dei modelli migliorino secondo una legge di potenza rispetto al calcolo (compute), alla dimensione del modello e alla dimensione del dataset. Sebbene siano ben consolidate per i grandi modelli linguistici (LLM), queste relazioni stanno emergendo anche nella fisica delle particelle. Una distinzione fondamentale nella fisica fondamentale è la capacità di generare dati sintetici ad alta fedeltà tramite simulatori a un costo relativamente basso rispetto alla spesa computazionale per l'addestramento di modelli più grandi. Ciò crea un'opportunità unica per ingegnerizzare il dataset di pretraining stesso per influenzare il comportamento della scalabilità. La domanda centrale affrontata è se la composizione dei dati di pretraining — specificamente la loro diversità e l'allineamento con i task a valle (downstream) — possa essere ingegnerizzata per spostare il regime di scalabilità computazionalmente ottimale dal favorire modelli più grandi al favorire dataset più grandi.

Metodologia
Lo studio si concentra sul compito di classificare i jet adronici prodotti in collisioni di particelle ad alta energia. Gli autori utilizzano un'architettura transformer generica che elabora i dati dei jet come una nuvola di punti (point cloud), variando le dimensioni del modello da circa 3.000 a 10,5 milioni di parametri (coprendo tre ordini di grandezza) mantenendo fissi la profondità e le dimensioni delle teste di attenzione.

Il design sperimentale prevede un protocolo di addestramento in due fasi:

1. Pretraining: I modelli vengono preaddestrati su sottoinsiemi del dataset JetClass-II, che contiene 188 classi di jet simulati. Gli autori definiscono quattro distinti sottoinsiemi di pretraining per manipolare diversità e allineamento:
   • QCD: Solo jet iniziati da quark leggeri o gluoni (17 classi).
   • QCD + res2p: Jet QCD più jet derivanti da decadimenti a due corpi di risonanze oltre il Modello Standard (BSM).
   • QCD + res34p: Jet QCD più jet derivanti da decadimenti a tre o quattro corpi di risonanze BSM.
   • QCD + res2p + res34p: Il dataset completo che include tutti i decadimenti di risonanze BSM.
   • Nota: I sottoinsiemi BSM introducono una maggiore diversità (più classi di processo, copertura più ampia dello spazio delle fasi) e un migliore allineamento con il task a valle (topologie multi-prong) rispetto ai dati solo-QCD.
2. Fine-tuning: I modelli preaddestrati vengono sottoposti a fine-tuning sul dataset originale JetClass per un compito di classificazione di jet a 10 classi (identificazione di quark leggeri/gluoni, quark top, bosoni W/Z e particelle Higgs). Questo compito richiede l'identificazione della molteplicità dei prong e delle scale di massa, che sono ben rappresentate nei dati preaddestrati aumentati con BSM ma scarsamente rappresentate nei dati solo-QCD.

Gli esponenti di scalabilità vengono estratti adattando leggi di potenza alla dimensione del modello computazionalmente ottimale ($N^*$) e alla dimensione del dataset ($D^*$) in funzione del calcolo totale ($C$). Lo studio confronta questi esponenti tra l'addestramento "da zero" (scratch) e le varie configurazioni di pretraining.

Risultati Chiave
Lo studio dimosta che la composizione dei dati di pretraining altera significativamente gli esponenti di scalabilità computazionalmente ottimali:

• Addestramento da zero (Scratch): L'addestramento partendo da zero produce esponenti di $a \approx 0,52$ (dimensione del modello) e $b \approx 0,48$ (dimensione del dataset), indicando un'allocazione delle risorse computazionali approssimativamente bilanciata tra dimensione del modello e dati, coerentemente con quanto riscontrato negli LLM.
• Pretraining Solo-QCD: Il pretraining esclusivamente su jet QCD produce uno spostamento marginale ($a \approx 0,53, b \approx 0,47$), suggerendo che il pretraining da solo, senza specifico allineamento o diversità, non cambia fondamentalmente il regime di scalabilità.
• Pretraining Aumentato con BSM: L'inclusione di decadimenti di risonanze BSM nel corpus di pretraining causa uno spostamento drammatico. Con il dataset completo aumentato con BSM, gli esponenti passano a $a \approx 0,22$ e $b \approx 0,78$.
  • Ciò indica un regime in cui la strategia computazionalmente ottimale favorisce pesantemente l'aumento della dimensione del dataset rispetto all'aumento della dimensione del modello.
  • Lo spostamento rappresenta una riduzione di un fattore di 2,3 dell'esponente di scalabilità per la dimensione del modello rispetto al baseline da zero.
  • Le curve di perdita (loss) del fine-tuning confermano che il pretraining potenziato con BSM riduce costantemente la perdita per tutte le dimensioni del modello, con benefici che aumentano per i modelli più grandi.

Contributi Chiave

1. Ingegneria delle Leggi di Scalabilità: Il documento fornisce il primo studio sistematico che mostra come la composizione dei dati di pretraining possa essere ingegnerizzata per spostare gli esponenti di scalabilità nella fisica fondamentale. Dimostra che la diversità e l'allineamento a valle nel corpus di pretraining possono spostare il regime di scalabilità ottimale verso strategie che favoriscono i dati.
2. Spostamento Quantitativo: Il lavoro quantifica lo spostamento da un regime di scalabilità bilanciato ($a \approx b \approx 0,5$) a un regime fortemente orientato ai dati ($a \approx 0,22, b \approx 0,78$) incorporando la fisica BSM nel pretraining.
3. Implicazioni per i Modelli di Fondazione: I risultati suggeriscono che i modelli di fondazione preaddestrati su dati sintetici diversificati e allineati possono raggiungere prestazioni ottimali con un numero inferiore di parametri, permettendo di reindirizzare i budget computazionali risparmiati verso la generazione di ulteriori dati di addestramento.

Significato e Rivendicazioni
Gli autori affermano che questo lavoro identifica un nuovo spazio di progettazione per il machine learning scientifico: gli input fisici per l'addestramento dei modelli di fondazione. A differenza dei domini del linguaggio naturale o delle immagini dove la cura dei dati è limitata dalla disponibilità, la fisica fondamentale può sfruttare simulatori ad alta fedeltà ed economici per costruire corpora di pretraining che modellano attivamente le leggi di scalabilità.

Il documento conclude con modestia che, sebbene il pretraining su corpora ben composti permetta un regime in cui il calcolo a valle sia meglio speso in più dati piuttosto che in modelli più grandi, sono necessari ulteriori lavori per verificare se questi risultati si generalizzano attraverso diversi task di fine-tuning, scale di modello più grandi e diversi set di dati. Lo studio non rivendica di aver risolto tutte le sfide di scalabilità, ma evidenzia l'ingegneria della composizione del pretraining come una leva sottoutilizzata per massimizzare il potenziale di scoperta dei modelli di fondazione scientifici.