On the Value of Tokeniser Pretraining in Physics Foundation Models

Dit onderzoek toont aan dat het vooraf trainen van een tokeniser met een auto-encoding-doelstelling de rekenefficiëntie en nauwkeurigheid van fysische foundation-modellen aanzienlijk verbetert, vooral wanneer de vooraftraining op hetzelfde fysische systeem als de downstream-taak wordt uitgevoerd.

De kern

Hier is een uitleg van het onderzoek in eenvoudig Nederlands, met behulp van alledaagse vergelijkingen.

De Kernvraag: Hoe leer je een computer om de natuur te begrijpen?

Stel je voor dat je een superintelligente robot wilt bouwen die kan voorspellen hoe het weer verandert, hoe vloeistoffen stromen of hoe sterrenstelsels bewegen. Deze robot moet enorme hoeveelheden data verwerken: miljoenen pixels die elke seconde veranderen.

Het probleem is dat deze data te groot en te complex is om direct te "eten" door de hersenen van de robot (het AI-model). Het is alsof je iemand probeert te leren zwemmen door hem direct in de diepste, woeligste oceaan te gooien zonder eerst te oefenen in het bad.

De Oplossing: Twee Stappen in plaats van Eén

De auteurs van dit paper ontdekten dat het slim is om de robot in twee stappen te trainen, net zoals een mens eerst een taal leert voordat hij poëzie schrijft.

1. De "Vertaler" (De Tokeniser)
De eerste stap is het bouwen van een vertaler. Deze vertaler kijkt naar de ruwe, chaotische beelden (de pixels) en vat ze samen in compacte, begrijpelijke "woorden" of symbolen (tokens).

• Vergelijking: Denk aan een samenvatting van een dik boek. In plaats van alle 500 pagina's te lezen, schrijft de vertaler een samenvatting van 10 regels die de essentie bevat. Dit maakt het voor de robot veel makkelijker om te begrijpen wat er gebeurt.

2. De "Denker" (Het Dynamics Model)
De tweede stap is de robot die deze samenvattingen leest en voorspelt wat er als volgt gebeurt.

• Vergelijking: Nu leest de robot de samenvattingen en probeert hij te raden hoe het verhaal verder gaat. Omdat hij niet meer hoeft na te denken over elke individuele pixel, kan hij veel sneller en slimmer redeneren.

Het Grote Geheim: Oefen eerst! (Pretraining)

In het verleden probeerden wetenschappers vaak om de "vertaler" en de "denker" tegelijkertijd te leren, vanaf nul. Dit is als proberen een auto te bouwen terwijl je tegelijkertijd leert rijden. Het gaat vaak moeizaam.

Dit paper toont aan dat het veel beter werkt als je de vertaler eerst apart traint (pretraining) voordat je de rest van de robot in werking stelt.

De drie belangrijkste ontdekkingen:

1. Oefen in dezelfde omgeving (In-Domain):
   Als je de vertaler traint op dezelfde soort data als waar hij later voor gebruikt wordt (bijvoorbeeld: eerst oefenen met windstoten, en dan echt voorspellen van windstoten), werkt het fantastisch.

   • Het resultaat: De robot leert 64% sneller en maakt veel minder fouten dan als je alles vanaf nul doet. Het is alsof je een speler eerst laat oefenen op een trainingsveld voordat hij de grote finale speelt.

2. Oefen in een andere omgeving (Out-of-Domain):
   Als je de vertaler traint op iets heel anders (bijvoorbeeld: eerst oefenen met vloeistoffen, en dan voorspellen van sterren), helpt het ook, maar minder.

   • Het resultaat: Het helpt nog steeds, maar niet zo enorm als bij optie 1. Het is alsof je eerst basketbal hebt gespeeld; dat helpt je wel bij voetbal (je bent fit en hebt coördinatie), maar je bent niet direct een topvoetballer.

3. Vries de vertaler in (Freezing):
   Een verrassende vondst is dat je de "vertaler" na het oefenen kunt bevriezen. Je laat hem niet meer veranderen tijdens de echte training.

   • Waarom? Als je de vertaler te veel laat aanpassen tijdens het lange proces van voorspellen, begint de robot te twijfelen en maken de fouten zich op (zoals een dominosteen-effect). Door de vertaler vast te zetten, blijft de basis stabiel.
   • Het voordeel: Dit maakt de training niet alleen beter, maar ook veel goedkoper en sneller, omdat je minder rekenkracht nodig hebt.

Waarom is dit belangrijk?

De natuur is complex en de data is enorm. Door eerst een goede "samenvatter" (tokeniser) te bouwen en deze goed te oefenen voordat we de grote AI aan de slag laten, besparen we enorme hoeveelheden tijd en energie.

Samengevat in één zin:
Net zoals een goede vertaler die de taal al beheerst de communicatie veel soepeler maakt, helpt een vooraf getrainde "samenvatter" AI-modellen om de complexe wetten van de natuur veel sneller en nauwkeuriger te leren begrijpen, vooral als we ze eerst laten oefenen op soortgelijke taken.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "On the Value of Tokeniser Pretraining in Physics Foundation Models", geschreven in het Nederlands.

Probleemstelling

Moderne wetenschappelijke simulaties genereren enorme hoeveelheden data met hoge ruimtelijke en temporele resolutie, die diverse fysieke regimes en schalen omvatten. Het trainen van foundation modellen (zoals transformers) om de onderliggende dynamica van deze data direct in de pixelruimte te leren, is computationally prohibitief (te duur).

De huidige aanpak voor physics foundation modellen traint vaak twee taken gelijktijdig vanaf nul:

1. Het extraheren van compacte representaties (tokens) van hoge-resolutie spatiotemporele data.
2. Het modelleren van de globale fysieke dynamica op basis van deze tokens.

Het paper stelt dat het gelijktijdig leren van beide taken vanaf nul de effectiviteit van beide processen kan belemmeren. In de computer vision is het standaard om tokenizers vooraf te trainen, maar dit is nog niet systematisch onderzocht voor fysica-modellen. De centrale vraag is: Kunnen physics foundation modellen profiteren van het vooraf trainen van de tokenizer, en hoe beïnvloedt de domein-overeenkomst (domain alignment) tussen de voor- en nascholing de prestaties?

Methodologie

1. Dataset en Experimentopzet
De auteurs gebruiken "The Well", een collectie van diverse 2D-fysicasimulaties, waaronder:

• Euler multiquadrants (doelwit voor de downstream taak).
• Rayleigh-Bénard convectie, shear flow en actieve materie (gebruikt voor out-of-domain pretraining).
  De data wordt opgedeeld in sequenties van 10 frames.

2. Architectuur
Het model bestaat uit twee componenten:

• Tokenizer: Een autoencoder gebaseerd op causale convoluties (afgeleid van MAGVIT-2), maar zonder vector quantisatie of adversarial loss. Het leert continue latente representaties. De auteurs introduceren hier flexibele spatiotemporele compressie, waardoor de compressieverhouding tijdens runtime kan worden aangepast.
• Processor: Een transformer-architectuur (van Walrus) met gefactoriseerde ruimtelijke en temporele attention, die de tokens verwerkt om de volgende frame te voorspellen.

3. Training Strategieën
De auteurs vergelijken drie scenario's voor de downstream taak (autoregressieve voorspelling van fysieke dynamica):

• Vanaf nul (No Pretraining): Zowel tokenizer als processor worden samen getraind.
• In-domain Pretraining: De tokenizer wordt eerst getraind op dezelfde dataset (Euler) als de downstream taak.
• Out-of-domain Pretraining: De tokenizer wordt getraind op een mengsel van andere datasets (Rayleigh-Bénard, etc.) voordat de downstream taak op Euler wordt getraind.

Binnen deze scenario's worden twee parameterstrategieën getest:

• Fully Trainable: Alle parameters van de tokenizer worden bijgewerkt tijdens de downstream training.
• Mostly Frozen: De kern van de tokenizer is bevroren; alleen de lagen die direct interface met de pixelruimte en de latente ruimte worden getraind (ongeveer 2% van de tokenizer parameters).

4. Evaluatiemetrics

• VRMSE (Variance-Normalised Root Mean Squared Error): Meet reconstructiefout ten opzichte van de variabiliteit van het doelveld.
• NEPS (Normalised Error Power Spectrum): Analyseert de fout op verschillende frequenties (laag, midden, hoog) om te zien hoe goed structuren op verschillende schalen worden vastgelegd.

Belangrijkste Bijdragen

1. Eerste systematisch onderzoek: Dit is het eerste werk dat systematisch de waarde van tokenizer pretraining voor physics foundation models onderzoekt.
2. Flexibele compressie: Introductie van spatiotemporele compressieoperaties die causale convoluties uitbreiden om runtime-aanpasbare compressieverhoudingen te ondersteunen. Dit maakt aanpassing mogelijk aan verschillende downstream taken zonder hertraining.
3. Strategische preselectie van data: Het paper levert praktische richtlijnen over hoe pretraining data moet worden geselecteerd om de efficiëntie te maximaliseren.

Resultaten

1. Domein-uitlijning is cruciaal

• In-domain pretraining: Leidt tot enorme verbeteringen. Na 10.500 trainingstappen is de VRMSE met 64% verminderd vergeleken met training vanaf nul (0.158 vs 0.439 voor fully trainable modellen). Dit voordeel treedt vroeg op en blijft bestaan.
• Out-of-domain pretraining: Levert bescheiden verbeteringen op (ongeveer 19% verbetering) als de tokenizer nog getraind wordt. Als de tokenizer in dit scenario bevroren wordt, verslechtert de prestatie zelfs onder het niveau van "zonder pretraining". Dit suggereert dat de fysieke velden te verschillend zijn voor directe transfer zonder aanpassing.

2. Invloed op leergedrag per schaal

• Lage frequenties: In-domain modellen bereiken een uitstekende prestatie (NEPS ~0.001), terwijl modellen zonder pretraining veel hogere fouten hebben.
• Midden frequenties: In-domain modellen behouden consistent lagere fouten.
• Hoge frequenties: Alle modellen hebben moeite met hoge frequenties (NEPS ≈ 1.0), maar in-domain pretraining zorgt voor continue verbetering, terwijl modellen zonder pretraining verslechteren tijdens training.

3. Het effect van "Freezing" (Bevriezen)

• Voor korte voorspellingen (volgende frame) presteren "fully trainable" en "mostly frozen" in-domain modellen gelijk.
• Voor lange autoregressieve rollouts (meerdere stappen vooruit) presteert de bevroren tokenizer echter beter. De fouten stapelen zich minder snel op. Dit suggereert dat het bevriezen van de pre-getrainde tokenizer fungeert als regularisatie, wat de stabiliteit van lange voorspellingen verbetert.
• Praktisch voordeel: Bevriezen reduceert het aantal trainbare parameters in de tokenizer met 98% (van 5M naar 85k), wat de rekenefficiëntie enorm verhoogt.

Betekenis en Conclusie

Dit paper demonstreert dat het vooraf trainen van een tokenizer een praktische en effectieve methode is om de training van physics foundation modellen te versnellen en te verbeteren, mits er sprake is van domein-uitlijning tussen de voor- en nascholing.

De belangrijkste inzichten zijn:

• Efficiëntie: Pretraining reduceert de rekentijd om een bepaald prestatieniveau te bereiken aanzienlijk.
• Stabiliteit: Het bevriezen van een in-domain pre-getrainde tokenizer verbetert de stabiliteit bij lange voorspellingen en verlaagt de trainingskosten.
• Data-keuze: De keuze van de pretraining data is kritisch; pretraining op een ander fysiek systeem zonder voldoende overlap in fysieke velden kan zelfs schadelijk zijn als de tokenizer niet goed wordt aangepast.

De auteurs concluderen dat tokeniser pretraining een eenvoudige maar krachtige strategie is voor het bouwen van efficiënte foundation modellen voor de wetenschap, en dat toekomstig werk moet focussen op het optimaliseren van data-compositie en het ontwikkelen van fysica-bewuste evaluatiemetrics.