Towards Engineering Scaling Laws with Pretraining Data Composition

Technische Samenvatting: Naar het Ingenieurswerk van Schalingswetten met Pretraining Datacompositie

Probleemstelling
Neurale schalingswetten beschrijven hoe de prestaties van modellen verbeteren als een machtswet met betrekking tot rekenkracht, modelgrootte en datasetgrootte. Hoewel deze relaties goed gevestigd zijn voor grote taalmodellen (LLM's), komen ze nu ook op in de deeltjesfysica. Een cruciaal onderscheid in de fundamentele fysica is de mogelijkheid om hoogwaardige synthetische data te genereren via simulators tegen relatief lage kosten in vergelijking met de computationele kosten van het trainen van grotere modellen. Dit creëert een unieke kans om de pretraining-dataset zelf te ontwerpen om de schalingsvoordelen te beïnvloeden. De centrale vraag die wordt behandeld is of de compositie van pretraining-data — specifiek de diversiteit en de afstemming op downstream-taken — kan worden ontworpen om het rekenoptimaal schalingsregime te verschuiven van het bevoordelen van grotere modellen naar het bevoordelen van grotere datasets.

Methodologie
De studie richt zich op de taak van het classificeren van hadronische jets geproduceerd in hoogenergetische deeltjesbotsingen. De auteurs maken gebruik van een generieke transformer-architectuur die jet-data verwerkt als een puntenwolk (point cloud), waarbij de modelgroottes variëren van ongeveer 3.000 tot 10,5 miljoen parameters (over drie ordes van grootte heen) terwijl de diepte en de dimensies van de attention heads constant blijven.

Het experimentele ontwerp omvat een tweefasig trainingsprotocol:

1. Pretraining: Modellen worden voorgetraind op subsets van de JetClass-II dataset, die 188 klassen van gesimuleerde jets bevat. De auteurs definiëren vier verschillende pretraining-subsets om diversiteit en afstemming te manipuleren:
   • QCD: Alleen jets geïnitieerd door lichte quarks of gluonen (17 klassen).
   • QCD + res2p: QCD-jets plus jets van twee-lichaam-vervallen van Beyond the Standard Model (BSM) resonanties.
   • QCD + res34p: QCD-jets plus jets van drie- of vier-lichaam-vervallen van BSM-resonanties.
   • QCD + res2p + res34p: De volledige dataset inclusief alle BSM-resonantievervallen.
   • Opmerking: BSM-subsets introduceren een grotere diversiteit (meer procesklassen, bredere fase-ruimte dekking) en een betere afstemming met de downstream-taak (multi-prong topologieën) vergeleken met QCD-only data.
2. Fine-tuning: Voorgetrainde modellen worden gefinetuned op de oorspronkelijke JetClass dataset voor een 10-klasse jet-classificatietaak (het identificeren van lichte quarks/gluonen, top-quarks, W/Z-bosonen en Higgs-deeltjes). Deze taak vereist het identificeren van prong-multipliciteit en massaschalen, die goed vertegenwoordigd zijn in de met BSM verrijkte pretraining-data maar slecht vertegenwoordigd zijn in QCD-only data.

De schalingsexponenten worden geëxtraheerd door machtswetten te fitten aan de rekenoptimale modelgrootte ($N^*$) en datasetgrootte ($D^*$) als functie van de totale rekenkracht ($C$). De studie vergelijkt deze exponenten over "scratch" training (trainen vanaf nul, zonder pretraining) en de verschillende pretraining-configuraties.

Belangrijkste Resultaten
De studie toont aan dat de compositie van de pretraining-data de rekenoptimale schalingsexponenten significant verandert:

• Scratch Training: Training vanaf nul levert exponenten op van $a \approx 0,52$ (modelgrootte) en $b \approx 0,48$ (datasetgrootte), wat wijst op een min of meer even evenredige allocatie van rekenbronnen tussen modelgrootte en data, consistent met bevindingen in LLM's.
• QCD-Only Pretraining: Pretraining uitsluitend op QCD-jets resulteert in een marginale verschuiving ($a \approx 0,53, b \approx 0,47$), wat suggereert dat pretraining alleen, zonder specifieke afstemming of diversiteit, het fundamentele schalingsregime niet fundamenteel verandert.
• BSM-Augmented Pretraining: Het opnemen van BSM-resonantievervallen in de pretraining-corpus veroorzaakt een dramatische verschuiving. Met de volledige BSM-verrijkte dataset verschuiven de exponenten naar $a \approx 0,22$ en $b \approx 0,78$.
  • Dit geeft aan een regime waarin de rekenoptimale strategie zwaar de voorkeur geeft aan het vergroten van de dataset boven het vergroten van de modelgrootte.
  • De verschuiving vertegenwoordigt een factor-2,3 reductie in de schalingsexponent voor modelgrootte vergeleken met de scratch-baseline.
  • Fine-tuning loss-curves bevestigen dat BSM-verrijkte pretraining de loss consistent verlaagt voor alle modelgroottes, waarbij de voordelen toenemen voor grotere modellen.

Belangrijkste Bijdragen

1. Engineering van Schalingswetten: Het artikel biedt de eerste systematische studie die aantoont dat de compositie van pretraining-data kan worden gebruikt om de schalingsexponenten in de fundamentele fysica te manipuleren. Het toont aan dat diversiteit en downstream-afstemming in de pretraining-corpus een schalingsregime kunnen verschuiven naar een data-bevoordelende strategie.
2. Kwantitatieve Verschuiving: Het werk kwantificeert de verschuiving van een gebalanceerd schalingsregime ($a \approx b \approx 0,5$) naar een sterk data-bevoordelend regime ($a \approx 0,22, b \approx 0,78$) door BSM-fysica in de pretraining te integreren.
3. Implicaties voor Foundation Models: De resultaten suggereren dat foundation-modellen die zijn voorgetraind op diverse en afgestemde synthetische data, optimale prestaties kunnen bereiken met kleinere parametersaantallen, waardoor bespaarde rekenbudgetten kunnen worden hergebruikt voor het genereren van extra trainingsdata.

Betekenis en Claims
De auteurs beweren dat dit werk een nieuwe ontwerpplek identificeert voor wetenschappelijke machine learning: de fysieke inputs voor het trainen van foundation models. In tegenstelling tot natuurlijke taal of beelddomeinen, waar datacuratie beperkt wordt door beschikbaarheid, kan de fundamentele fysica gebruikmaken van goedkope, hoogwaardige simulators om pretraining-corpora te construeren die de schalingswetten actief vormgeven.

Het artikel concludeert bescheiden dat hoewel pretraining op goed samengestelde corpora een regime mogelijk maakt waarin de downstream-rekenkracht het beste aan meer data dan aan grotere modellen kan worden besteed, verder werk nodig is om te verifiëren of deze resultaten generaliseren naar verschillende fine-tuning taken, grotere model schalen en verschillende dataset groottes. De studie beweert niet dat het alle schalingsuitdagingen heeft opgelost, maar benadrukt dat pretraining-compositie-engineering een onderbelicht hefboommechanisme is voor het maximaliseren van het ontdekkingspotentieel van wetenschappelijke foundation models.