🤖 machine learning

Distillation of Foundation Models for Time-dependent PDEs

Het artikel stelt TREX voor, een kennisdistillatie-framework dat synthetische trajecten met een lange horizon genereert vanuit gefinetunede fundamentele modellen om compacte, efficiënte student-netwerken voor tijdafhankelijke PDE's te trainen, waarbij significante reducties in parameters en inferentielatentie worden bereikt terwijl de nauwkeurigheid van het oorspronkelijke model wordt behouden of overtroffen.

Oorspronkelijke auteurs: Daniel Musekamp, Boshra Ariguib, Andrei Manolache, Mathias Niepert

Gepubliceerd 2026-08-13

📖 6 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Daniel Musekamp, Boshra Ariguib, Andrei Manolache, Mathias Niepert

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je de toekomst van een chaotisch systeem probeert te voorspellen, zoals een kolkende storm, een stromende rivier of de manier waarop warmte zich door een metalen plaat verspreidt. In de wereld van de natuurkunde worden deze beschreven door vergelijkingen die partiële differentiaalvergelijkingen (PDE's) worden genoemd. Het oplossen van deze vergelijkingen is als het navigeren door een doolhof waar de muren constant bewegen; het is ongelooflijk moeilijk en vereist meestal supercomputers om de getallen stap voor stap te verwerken. Onlangs hebben wetenschappers "Foundation Models" gebouwd: enorme, alwetende AI-hersenen die getraind zijn op duizenden verschillende natuurkundige problemen. Deze reuzen kunnen nieuwe taken snel leren, maar ze zijn zo groot en traag dat ze niet gebruikt kunnen worden voor real-time zaken, zoals het aansturen van een robotarm of het voorspellen van het weer in de volgende seconde. Ze zijn als een briljante maar trage bibliothecaris die elk boek in de bibliotheek moet lezen voordat hij één vraag kan beantwoorden.

De grote vraag is: kunnen we een kleine, snelle en behendige student-AI leren om net als de enorme bibliothecaris te denken, zonder dat de reus telkens al het werk hoeft te doen? Dit is de uitdaging van "knowledge distillation" (kennisdestillatie). Meestal heb je veel voorbeelden nodig om een student te onderwijzen. Maar in de echte wereld hebben we vaak slechts een paar momentopnames van een systeem (zoals een paar seconden video van een sensor) en hebben we geen idee wat de begincondities waren voor andere scenario's. De paper die je zojuist hebt gelezen, pakt precies dit probleem aan: hoe neem je een trage, krachtige AI-leraar en comprimeer je zijn wijsheid in een snelle, kleine student, zelfs wanneer we niet genoeg data hebben om de student vanaf nul te trainen.

Het verhaal van TREX: Een kleine robot leren denken als een reus

Maak kennis met TREX (Teacher Rollout Extension). Het is een slimme nieuwe methode, uitgevonden door onderzoeker Daniel Musekamp en zijn team, om het "trage reus"-probleem op te lossen. Stel je voor dat je een meesterkok hebt (het Foundation Model) die miljoenen gerechten heeft bereid en precies weet hoe smaken in de loop van de tijd veranderen. Je wilt een jonge leerling (het Student Model) leren om op dezelfde manier te koken, maar je hebt slechts drie recepten en een zeer korte tijd om te leren. Als je de leerling alleen laat oefenen op die drie recepten, kan hij vast komen te zitten of fouten maken wanneer de ingrediënten licht veranderen.

Het probleem is dat de meesterkok te traag is om elk gerecht te koken voor de leerling om naar te kijken. En je kunt de chef ook niet vragen om nieuwe ingrediënten te "verbeelden", omdat je niet weet welke ingrediënten de leerling in de echte wereld nodig zal hebben.

Hier wordt TREX creatief. In plaats van te wachten tot de meesterkok nieuwe maaltijden vanaf nul bereidt, laat TREX de meesterkok de paar recepten die je wel hebt nemen en deze "uitrollen" (rollout) naar de toekomst. Denk er zo over na: de meesterkok neemt jouw ene recept voor een cake en zegt: "Oké, laten we die cake bakken, en daarna bakken we hem weer, en weer, en weer, voor honderd stappen." Dit creëert een lange, denkbeeldige tijdlijn van hoe die cake evolueert.

Maar hier komt de twist: soms wordt de echte wereld rommelig. Er waait een windvlaag, een lepel valt op de grond, of de temperatuur schiet omhoog. Om de leerling hierop voor te bereiden, gooit TREX af en toe een beetje "ruis" (zoals een zachte schudbeweging of een snufje willekeurige kruiden) op de cake terwijl de meesterkok de cake aan het uitrollen is. De meesterkok moet dan uitzoeken hoe hij de cake kan herstellen en het koken kan voortzetten vanuit die rommelige staat. Dit leert de leerling niet alleen hoe hij een perfecte cake bakt, maar ook hoe hij moet herstellen wanneer dingen misgaan.

Het magische resultaat:
Door deze methode te gebruiken, ontdekten de onderzoekers dat ze een kleine student-model konden trainen dat 3.604 keer kleiner is dan de reusachtige leraar. Toch kan deze kleine student de toekomst van het natuurkundige systeem net zo nauwkeurig voorspellen als de reus. Sterker nog, in sommige tests was de student zelfs nauwkeuriger bij langetermijnvoorspellingen, omdat het leerde stabieler te zijn en minder in de war raakte door kleine fouten dan de reus.

Waarom is dit belangrijk?
De reusachtige modellen zijn als een Ferrari die slechts 2 mijl per gallon rijdt — het is snel in een rechte lijn, maar te duur om voor alledaagse taken te gebruiken. De student-modellen getraind met TREX zijn als een brandstofefficiënte elektrische scooter. Ze zijn:

Snel: Ze kunnen voorspellingen doen die meer dan 10 keer sneller zijn dan die van de reus.
Licht: Ze gebruiken een fractie van het computergeheugen (ongeveer 3,2 keer minder).
Slim: Ze kunnen worden gebouwd met speciale "natuurkundige regels" die in hun brein zijn ingebakken (zoals weten dat als je een vloeistof roteert, de fysica mee moet roteren), wat de reusachtige modellen niet strikt volgden.

Wat de onderzoekers hebben uitgesloten:
Het team heeft gecontroleerd of ze echt de "begincondities" (de exacte begintoestand van het universum) nodig hadden om dit te laten werken. Ze ontdekten dat je de volledige distributie van beginstaten niet hoeft te kennen. Je kunt gewoon beginnen met de paar datapunten die je hebt en de leraar de rest laten "uitrollen". Ze toonden ook aan dat het simpelweg kopiëren van de antwoorden van de leraar niet genoeg is; de "noisy rollouts" (de rommelige, geschudde oefenruns) zijn cruciaal voor de student om te leren hoe hij met de chaos van de echte wereld moet omgaan.

Hoe zeker zijn ze?
De onderzoekers hebben niet zomaar gegokt; ze hebben de cijfers gecontroleerd op verschillende natuurkundige problemen, waaronder vloeistofdynamica (zoals stromend water) en samendrukbaar gas (zoals lucht in een straalmotor). Ze hebben dit getest op twee verschillende reusachtige AI-leraren (genaamd Poseidon en Walrus) en vonden dat de TREX-methode consequent studenten produceerde die de nauwkeurigheid van de leraren evenaarden of zelfs overtroffen. Ze testten zelfs wat er gebeurt als je de "ruis" of de "ground truth"-data verwijdert, en de methode hield stand, hoewel het het beste werkte met beide aanwezig.

Kortom, TREX is een manier om de superintelligentie van een enorme AI te distilleren in een kleine, supersnelle tool die op gewone computers kan draaien, waardoor real-time natuurkundige simulaties mogelijk worden voor zaken als weersvoorspelling, technisch ontwerp en het aansturen van robots. Het is als het nemen van een bibliotheek vol genie en dat inkrimpen tot één enkele, superintelligente gids voor in je broekzak.

Technische Samenvatting: Distillatie van Foundation Models voor Tijdafhankelijke PDE's

1. Probleemdefinitie

Partiële differentiaalvergelijkingen (PDE's) zijn fundamenteel voor het modelleren van spatiotemporele fysische processen zoals vloeistofdynamica. Hoewel neurale netwerken zijn opgekomen als efficiënte surrogaatmodellen voor PDE's, worden ze geconfronteerd met een kritieke afweging:

Foundation Models (FMs): Grootschalige modellen die zijn voorgetraind op diverse fysica-datasets (bijv. Poseidon, Walrus) vertonen een sterke generalisatie naar downstream-taken met beperkte data. Ze zijn echter computationeel duur en vertrouwen vaak op grote transformer-architecturen die een gebrek hebben aan specifieke inductieve biases (bijv. equivariantie) en aanzienlijke inferentietijd en geheugen vereisen. Dit beperkt hun bruikbaarheid in real-time toepassingen zoals controle en planning.
Compacte Surrogaten: Kleinere modellen (bijv. Fourier Neural Operators) zijn efficiënt, maar vereisen doorgaans grote hoeveelheden taakspecifieke trainingsdata om een hoge nauwkeurigheid te bereiken, wat vaak niet beschikbaar is in experimentele of gedeeltelijk geobserveerde systemen.

De kernuitdaging is het overdragen van de generalisatiecapaciteiten van grote, voorgetrainde PDE foundation models naar compacte, efficiënte student-modellen zonder dat toegang nodig is tot de volledige distributie van de initiële condities van het doeldomein, die vaak onbekend of duur is om te samplen.

2. Methodologie: Teacher Rollout Extension (TREX)

De auteurs stellen Teacher Rollout Extension (TREX) voor, een kennisdistillatie-framework ontworpen om de kloof tussen grote foundation models en efficiënte downstream surrogaten te overbruggen.

2.1 Kernconcept

TREX adresseert de beperking van standaard distillatie, die vaak vereist om nieuwe initiële condities te samplen uit een bekende distributie $p_0$ . In veel downstream-scenario's (bijv. experimentele data) is $p_0$ onbekend. In plaats daarvan vergroot TREX de beperkte beschikbare grondwaarheid-trajecten uit door lange synthetische trajecten te genereren via teacher rollouts.

2.2 De TREX-procedure

Fine-tuning: Een voorgetraind foundation model (Teacher, $T$ ) wordt eerst gefinetuned op een kleine set beschikbare downstream-trajecten.
Rollout Generatie: Startend vanaf de initiële toestanden van de weinige beschikbare grondwaarheid-trajecten, wordt de gefinetunede teacher autoregressief uitgerold om lange synthetische trajecten te genereren.
- Occupancy Measure: Dit proces samplet uit de door de teacher geïnduceerde rollout-distributie (toestanden die worden tegengekomen tijdens autoregressieve voorspelling) in plaats van alleen de korte geobserveerde trajecten.
- Ruisinjectie: Om de dekking van de toestandsruimte te verbreden en de student te leren hoe het systeem evolueert vanuit verstoorde toestanden, wordt er periodiek Gaussische ruis geïnjecteerd in de toestand van de teacher tijdens de rollouts. De teacher voorspelt vervolgens de volgende toestand op basis van deze verstoorde toestand.
Student Training: Een compact student-model ( $S$ $S$ ) wordt getraind op een mix van:
- Originele grondwaarheid-transities.
- Door de teacher gegenereerde transities (zowel ruisvrije als met ruis voorzien segmenten).
- De student loss combineert grondwaarheid-supervisie ( $L_{GT}$ ) en teacher-distillatie loss ( $L_{TREX}$ ).

2.3 Architecturale Inductieve Biases

In tegenstelling tot de teacher, die doorgaans een grote, flexibele transformer is, wordt de student-architectuur gekozen om taakspecifieke inductieve biases te incorporeren. De auteurs maken gebruik van een Tensorized Fourier Neural Operator (TFNO), die equivariantie naar translaties afdwingt en hoogfrequente ruis onderdrukt, eigenschappen die niet expliciet door de teacher worden afgedwongen.

3. Belangrijkste Bijdragen

Eerste Systematische Studie: Het artikel presenteert de eerste systematische onderzoek naar het distilleren van algemene voorgetrainde PDE foundation models naar compacte downstream surrogaatmodellen.
TREX Framework: Introductie van een nieuwe distillatiemethode die supervisiedata genereert via ruizige teacher rollouts, waardoor de noodzaak voor een expliciete sampler voor initiële condities wordt geëlimineerd.
Integratie van Inductieve Bias: Demonstratie dat het distilleren van een algemeen foundation model naar een kleine architectuur met geschikte inductieve biases (bijv. equivariantie) de fysische eigenschappen kan respecteren terwijl de voorspellende nauwkeurigheid van de teacher behouden blijft.

4. Experimentele Resultaten

De auteurs evalueerden TREX met behulp van Poseidon en Walrus als teachers over diverse vloeistofdynamica-benchmarks (Navier-Stokes en Compressible Euler vergelijkingen).

Nauwkeurigheid: TREX-getrainde studenten matchen of overtreffen de nauwkeurigheid van de grote foundation model teachers in low-data regimes (4 tot 32 trajecten). In sommige gevallen presteerde de student zelfs beter dan de teacher wanneer deze uitsluitend werd getraind op door de teacher gegenereerde data (zonder grondwaarheid).
Efficiëntie:
- Parameters: De student-modellen verminderden het aantal parameters met meerdere grootteordes (bijv. van ~628M parameters in Poseidon-L naar ~174K in TFNO, een reductie van ~3600x).
- Inferentietijd: De inferentietijd werd met meer dan een orde van grootte verminderd (bijv. van ~270ms naar ~22ms per stap).
- Geheugen: Het GPU-geheugengebruik voor een forward pass werd verminderd met een factor van ~3.2.
Stabiliteit en Equivariantie:
- Studenten vertoonden stabielere long-horizon rollouts met een tragere foutgroei vergeleken met de teachers.
- De TFNO-student demonstreerde aanzienlijk lagere translatie-equivariantie fouten vergeleken met de niet-equivariante teacher, waardoor consistentie onder ruimtelijke verschuivingen tijdens autoregressieve rollouts werd behouden.
Ablaties:
- De methode is robuust tegenover rollout-lengte en ruis-hyperparameters.
- Ruisinjectie tijdens rollouts bleek gunstig, met name bij hogere amplitudes.
- Distillatie bleek effectief zelfs wanneer de student-architectuur (UNet, standaard FNO) de teacher niet kon overtreffen met enkel grondwaarheid-data.

5. Betekenis en Claims

Het artikel claimt dat TREX effectief de afweging tussen de sterke generalisatie van grote foundation models en de efficiëntievereisten van praktische surrogaatmodellering verzoent. Door beperkte downstream-data uit te breiden via teacher rollouts, stelt TREX compacte modellen in staat om long-horizon dynamica en fysische symmetrieën te leren die anders enorme datasets of complexe architecturen zouden vereisen.

De auteurs positioneren dit werk als een stap naar het inzetbaar maken van foundation-model-gebaseerde surrogaten voor real-time wetenschappelijke en technische toepassingen (bijv. manufacturing control, energiesystemen) waar computationele kosten en inferentiesnelheid kritieke beperkingen zijn. Zij merken op dat hoewel de methode afhankelijk is van de kwaliteit van de gefinetunede teacher, het succesvol kennis overdraagt zelfs wanneer de teacher een gebrek heeft aan specifieke inductieve biases, waardoor de student fysische beperkingen zoals equivariantie kan afdwingen.