Mixture of Training: Recombining Small-Scale Scaffolded Pretraining Runs into a Larger Language Model

Sintesi Tecnica: Mixture of Training (MoT)

Problematica
Il pre-addestramento di modelli linguistici di grandi dimensioni (LLM) è attualmente organizzato come un singolo processo di ottimizzazione end-to-end monolitico. Questa architettura crea un collo di bottiglia alla scalabilità in cui tutti gli strati devono essere addestrati simultaneamente, rendendo il processo fragile rispetto ai guasti (dove un singolo errore interrompe l'intera esecuzione) e poco flessibile per la ricerca iterativa (dove i miglioramenti richiedono il riavvio o la continuazione di un sistema massiccio e accoppiato). Il documento si chiede se il pre-addestramento possa essere scomposto in job più piccoli, addestrabili indipendentemente, che possano poi essere ricomposti in un modello più grande e coerente senza sacrificare la qualità.

Metodologia: Mixture of Training (MoT)
Gli autori introducono il Mixture of Training (MoT), una procedura di pre-addestramento modulare e scaffoldata progettata per partizionare un Transformer target in blocchi di strati contigui. Il metodo opera in tre fasi:

1. Fase 0 (Preparazione): Un allineatore (aligner) pre-addestrato ($A$) viene selezionato o preparato. Il Transformer target ($F$) e l'allineatore sono entrambi partizionati in $K$ blocchi di strati contigui. L'allineatore è compatibile per forma con il target, condividendo la larghezza globale, la dimensionalità delle teste di attenzione, la larghezza del feed-forward, gli embedding dei token e le teste di output.
2. Fase 1 (Scaffolded Training): Il modello target viene decomposto in $K$ submodelli ($f_1, \dots, f_K$). Per ogni submodello $f_i$, viene costruito un network scaffoldato $S_i$ inserendo $f_i$ nella struttura dell'allineatore congelata:
   $$S_i = a_K \circ \dots \circ a_{i+1} \circ f_i \circ a_{i-1} \circ \dots \circ a_1$$
   In questa configurazione, solo il blocco target $f_i$ è addestrabile; tutti gli strati dell'allineatore sono congelati. Ogni scaffold è ottimizzato indipendentemente utilizzando lo standard della perdita di predizione del next-token. Fondamentalmente, non avviene alcuno scambio di gradienti tra i blocchi target durante questa fase, ma essi apprendono tutti all'interno del contesto rappresentazionale stabile fornito dall'allineatore condiviso.
   3.Fase 2 (Ricomposizione e Adattamento): L'allineatore viene scartato. I blocchi target addestrati vengono ricomposti nel modello completo $\hat{F} = f_K \circ \dots \circ f_1$. Questo modello "cold-composed" può essere valutato direttamente, oppure può subire un breve passaggio di adattamento end-to-end per risolvere eventuali discrepanze rappresentazionali rimanenti.

Contributi Chiave

1. Scaffolded Modular Training: Una procedura che addestra slice di profondità target indipendentemente, preservando al contempo un'interfaccia rappresentazionale condivisa tramite un allineatore congelato, consentendo la ricomposizione senza necessità di cucire gli strati post-hoc.
2. Prova del Meccanismo: Uno studio su piccola scala che dimostra come slice di profondità addestrate indipendentamente possano essere ricomposte con successo in un modello linguistico da 1,3 miliardi di parametri utilizzabile.
3. Contabilità di Compute e Scheduling: Un'analisi dettagliata dei trade-off tra calcolo (FLOPs), esposizione dei token e il percorso critico equivalente ideale degli strati. Ciò include il trattamento esplicito del costo dell'allineatore, distinguendo tra costi di addestramento "diretti" e costi "completamente addebitati" (assegnando l'intero costo dell'allineatore a un singolo run) rispetto ai costi ammortizzati su più run.
4. Studi di Ablazione: Evidenza empirica che mostra la necessità dell'allineatore, l'impatto di stream di dati disgiunti e il trade-off qualità-efficienza associato all'aumento del numero di suddivisioni ($K$).

Risultati Sperimentali
Gli autori hanno valutato il MoT su un modello in stile Gemma da 12 strati e 1,3 miliardi di parametri, addestrato sul dataset C4.

• Parità di Qualità: Uno schema di "qualità-parità", che ha esteso l'addestramento dei submodelli e incluso un passaggio di adattamento di 30k step, ha raggiunto la stessa perplexity (15,0) del baseline monolitico.
• Trade-off di Calcolo:
  • Lo schema di qualità-parità ha elaborato un numero aggregato di token maggiore (47,1B vs 33,6B) e ha avuto un percorso critico ideale più breve (potenziale di accelerazione di 1,7x), ma ha richiesto più FLOP totali (285,0 EFLOPs vs 268,4 EFLOPs) quando il costo dell'allineatore è stato completamente addebitato a un singolo run.
  • Tuttavia, se l'allineatore viene riutilizzato in $R$ run indipendenti, il costo effettivo per run diminuisce. Il documento nota che per $R \ge 3$, il costo effettivo scende al di sotto del baseline monolitico.
  • Uno schema a minor computazione (MoT + 15k adattamento) ha raggiunto una perplexity di 15,9 a 189,4 EFLOPs (completamente addebitati), valore che è inferiore al budget del baseline, sebbene non sia stato fornito un baseline monolitico a parità di calcolo per un confronto diretto a questo specifico budget.
• Ablazioni:
  • Necessità dell'Allineatore: Senza l'allineatore, la qualità della cold-composition è crollata (PPL 38,9), confermando che l'allineatore previene il mismatch delle interfacce.
  • Stream di Dati: Stream di dati disgiunti per i submodelli hanno migliorato la qualità della cold-composition (PPL 19,3) rispetto a dati condivisi (PPL 20,3) in presenza di un allineatore.
  • Conteggio delle Divisioni: Aumentare le divisioni da $K=2$ a $K=4$ ha ridotto il calcolo ma ha degradato la qualità della cold-composition (PPL 24,8), evidenziando un trade-off tra parallelismo e complessità dell'interfaccia.

Significatività e Rivendicazioni
Il documento posiziona il MoT non come un sostituto generale del pre-addestramento monolitico, ma come un framework su piccola scala per studiare se gli scaffolded sub-run possano fungere da unità di addestramento riutilizzabili.

Gli autori rivendicano tre conclusioni primarie:

1. I submodelli addestrati indipendentamente possono essere ricomposti in un modello linguante coerente se addestrati all'interno di uno scaffold di un allineatore condiviso; senza di esso, la cold-composition fallisce.
2. Il mismatch tra i blocchi addestrati indipendentamente è ampiamente recuperabile; un breve passaggio di adattamento colma la maggior parte del divario, e uno schema più lungo può raggiungere la parità con l'addestramento monolitico.
3. Il MoT altera la forma ingegneristica del pre-addestramento creando job più piccoli, indipendentemente programmabili, riavviabili e ablati. Ciò è particolarmente adatto alla ricerca di addestramento su piccola scala, rendendola potenzialmente più economica, più riproducibile e più facile da iterare.

Gli autori dichiarano esplicitamente che l'evidenza attuale dimostra la composibilità modulare e identifica un regime di scheduling di riutilizzo ammortizzato, piuttosto che stabilire un vantaggio uniforme di calcolo o di velocità hardware. Lo studio è intenzionalmente limitato nell'ambito (una singola famiglia di modelli, un singolo dataset, metriche basate solo sulla perplexity) per servire come prova del meccanismo per un design space concreto di sub-run scaffoldati e riutilizzabili.