Efficient Construction of Model Family through Progressive Training Using Model Expansion

Probleemstelling

Met de toenemende praktische toepassing van Large Language Models (LLMs) is het standaard geworden om model-families aan te bieden: sets van modellen met dezelfde architectuur maar verschillende parametergroottes (bijv. 1B, 4B, 8B parameters). Deze variatie is nodig om te voldoen aan uiteenlopende computatiebeperkingen, van mobiele apparaten tot grote servers.

De traditionele aanpak om zo'n familie te bouwen, is het onafhankelijk trainen van elk model vanaf nul (van scratch). Dit leidt tot een cumulatieve computekost die lineair toeneemt met het aantal modellen. Het trainen van een volledige familie (bijv. van 1B tot 8B parameters) vereist dus een enorme hoeveelheid GPU-uren, wat een significante financiële en ecologische last vormt voor onderzoekers en bedrijven.

Methodologie: Progressieve Training

De auteurs stellen een nieuwe methode voor: Progressive Training (Progressieve Training). In plaats van elk model onafhankelijk te trainen, worden kleinere modellen geleidelijk uitgebreid (expanded) om grotere modellen te initialiseren.

1. Model Uitbreiding (Model Expansion):

   • Het proces begint met het trainen van het kleinste model ($X_1$) vanaf willekeurige initialisatie.
   • Voor het trainen van het volgende, grotere model ($X_{i+1}$), worden de parameters van het reeds getrainde kleinere model ($X_i$) gebruikt als startpunt.
   • Hiervoor wordt de techniek bert2BERT gebruikt, die zowel de breedte (hidden dimensions) als de diepte (aantal lagen) van een Transformer-architectuur kan vergroten door gewichten te dupliceren en te stacken.

2. Token Allocatie en Kostenbeperking:

   • Het doel is om de totale rekentijd (FLOPs) van het trainen van de hele familie te beperken tot de kosten van het trainen van alleen het grootste model ($X_n$) vanaf nul.
   • De auteurs verdelen het totale token-budget zo dat: $\sum FLOPs(X_i, T^{prog}_i) = FLOPs(X_n, T^{scratch}_n)$.
   • Dit betekent dat de kleinere modellen minder tokens zien dan ze zouden zien bij onafhankelijk trainen, maar de "winst" in efficiëntie komt door het gebruik van de reeds geleerde kennis uit de kleinere modellen.

3. Aanpassing van de Leerkoers (Learning Rate):

   • Een cruciale innovatie is het dynamisch aanpassen van de maximale leerkoers op basis van de modelgrootte.
   • Kleinere modellen kunnen profiteren van hogere leerkoersen voor snellere convergentie, terwijl grotere modellen stabielere, lagere leerkoersen nodig hebben om instabiliteit (zoals loss spikes) te voorkomen.
   • In de experimenten werd de leerkoers lineair verlaagd van $1.5 \times 10^{-3}$ (voor 1B) naar $3.0 \times 10^{-4}$ (voor 8B).

Experimentele Opzet

• Dataset: FineWeb-Edu.
• Architectuur: Llama-achtige architectuur met variërende grootte: [1B, 2B, 4B, 8B] parameters.
• Vergelijking:
  • Independent: Elk model getraind vanaf nul volgens de Chinchilla-wet (of 2x Chinchilla).
  • Progressive: Modellen getraind via uitbreiding.
  • Progressive + Fixed Data: Een controlegroep waarbij het totale aantal unieke tokens wordt beperkt tot dat van het grootste model (om te bewijzen dat de prestatieverbetering niet komt door simpelweg meer unieke data te zien).

Belangrijkste Resultaten

1. Rekenkundige Efficiëntie:

   • De progressieve trainingsmethode reduceert de totale computekost met ongeveer 25% ten opzichte van onafhankelijk trainen, terwijl een volledige model-familie wordt opgeleverd.
   • Bij gebruik van een aangepaste leerkoers (LR adjustment) wordt de FLOPs-reductie zelfs 31%.

2. Prestaties:

   • Modellen getraind via progressieve training presteren gelijkwaardig of beter dan onafhankelijk getrainde modellen op diverse benchmarks (perplexiteit, LAMBADA, WinoGrande, ARC, etc.).
   • Het grootste model (8B) toont zelfs verbeteringen in perplexiteit en downstream-taken, vooral wanneer de leerkoers wordt aangepast.
   • De controlegroep (Progressive + Fixed Data) behaalde vergelijkbare resultaten, wat bevestigt dat de winst niet voortkomt uit het zien van meer unieke data, maar uit de efficiëntere initialisatie.

3. Consistentie tussen Modellen:

   • Een uniek voordeel is de hogere gedragsconsistentie binnen de familie. De auteurs maten de Kullback-Leibler (KL) divergentie tussen de output-verdelingen van opeenvolgende modelgroottes.
   • Progressief getrainde modellen vertoonden een significante lagere KL-divergentie dan onafhankelijk getrainde modellen. Dit betekent dat ze meer vergelijkbare interne mechanismen en besluitvormingsprocessen delen.

4. Praktisch Voordeel: Speculatieve Decoding:

   • De consistentie werd getest in speculatieve decoding (waarbij een klein model "drafts" genereert voor een groot model).
   • Een combinatie van een 4B en 8B model, beide progressief getraind, resulteerde in een acceptatiepercentage van 93,2% (tegenover 87,2% bij onafhankelijke modellen) en een 14,8% snellere inferentie-tijd. Dit toont aan dat de consistentie direct vertaalt naar efficiëntere inferentie.

Bijdragen en Significantie

• Efficiënte Constructie: De paper biedt een bewezen methode om model-families te bouwen met een kwart minder rekenkracht, wat essentieel is voor duurzame AI-ontwikkeling.
• Superieure Prestaties: Het weerlegt de aanname dat "trainen vanaf een bestaand model" noodzakelijkerwijs leidt tot suboptimale resultaten; met de juiste strategie (zoals LR-aanpassing) kunnen deze modellen zelfs beter presteren.
• Gedragconsistentie: Het paper introduceert het belang van consistente gedragingen binnen een model-familie, wat cruciaal is voor:
  • Deployments: Makkelijker schalen tussen apparaten zonder dat gebruikerservaring drastisch verandert.
  • Fine-tuning: Het overdragen van voorkeursdata of veiligheidsfilters tussen verschillende modelgroottes.
  • Inferentie-optimalisatie: Betere prestaties bij technieken zoals speculatieve decoding.

Conclusie:
De auteurs tonen aan dat progressieve training via model-expansie niet alleen een kostenefficiëntere manier is om model-families te creëren, maar ook resulteert in een coherenter en praktischer bruikbare set modellen. Deze aanpak is vooral waardevol in omgevingen met beperkte resources waar het trainen van meerdere schalen wenselijk is.