Fine-tuning MLIP foundation models: strategies for accuracy and transferability

Dit artikel evalueert zeven fine-tuning-strategieën voor op machine learning gebaseerde interatomaire potentiaalmodellen (MLIP foundation models) over diverse chemische benchmarks, waarbij wordt onthuld dat hoewel randvoorwaarden zoals de kwaliteit van het foundation model en de correcte energie-initialisatie essentieel zijn, naïeve fine-tuning optimaal is voor nauwkeurigheid bij één enkel systeem, terwijl multihead replay uniek de robuustheid buiten de distributie (out-of-distribution) behoudt voor bredere implementatie.

De kern

Stel je voor dat je een meesterkok hebt die jarenlang heeft geleerd om perfecte maaltijden te bereiden met alleen anorganische ingrediënten zoals stenen, metalen en zouten. Deze chef is een "Foundation Model". Nu wil je deze chef een specifiek nieuw gerecht leren maken, zoals een delicate organische soep of een biologische stoofpot, met een zeer kleine hoeveelheid nieuwe recepten.

De grote vraag is: Hoe leer je deze chef het nieuwe gerecht aan zonder dat hij de oude gerechten vergeet, of zonder dat zijn bestaande vaardigheden worden aangetast?

Dit artikel is een enorm keukenexperiment waarbij ze zeven verschillende manieren testen om deze meesterkok te "fine-tunen" (opnieuw te trainen). De onderzoekers ontdekten dat de methode van het onderwijzen minder belangrijk is dan drie cruciale "pre-cooking" stappen: het kiezen van de juiste chef, het instellen van de juiste basis en het afstellen van de hitte.

Hier is de uitslag van hun bevindingen in eenvoudige termen:

1. De Drie "Pre-Flight" Checks (Het Belangrijkste Deel)

Voordat je zelfs maar begint met het aanleren van het nieuwe recept, moet je drie dingen goed krijgen. Als je dit verpest, kan geen enkele onderwijsmethode je redden.

• Kies de Juiste Chef (Kwaliteit van het Foundation Model):

  • De Analogie: Je zou ook geen chef inhuren die alleen weet hoe hij water moet koken om hem te leren hoe hij een soufflé moet bakken.
  • De Bevinding: De kwaliteit van het originele model is belangrijker dan de fine-tuning strategie. Een model dat getraind is op een enorme, diverse dataset van anorganische materialen (zoals het "OMat24"-model) is veel beter in het leren van nieuwe, vreemde chemie dan een ouder, kleiner model. Zelfs als je dezelfde onderwijsmethode gebruikt, zal een "beter" foundation model altijd een beter eindgerecht produceren.

• Stel het Nulpunt In (Atomische Referentie-energie / $E_0$):

  • De Analogie: Stel je voor dat je de hoogte van een gebouw meet. Als je begint te meten vanaf de kelder in plaats van de begane grond, kloppen je cijfers niet en lijkt het gebouw te zweven of begraven te zijn. In de chemie moet je het "gewicht" van de individuele atomen aftrekken, zodat het model alleen leert over hoe ze met elkaar interageren.
  • De Bevinding: De onderzoekers ontdekten dat het gebruik van een slimme, "model-bewuste" manier om dit nulpunt in te stellen cruciaal is. Als je een luie, gemiddelde schatting gebruikt, wordt het model instabiel. Het kan er op papier goed uitzien (lage foutscores), maar het zal instorten wanneer je probeert de echte fysica te simuleren (zoals een gebouw dat bezwijkt in een windtunneltest).

• Zet de Hitte Lager (Hyperparameters):

  • De Analogie: Bij het leren van een nieuwe vaardigheid wil je niet zo snel gaan dat je struikelt, maar je wilt ook niet zo langzaam gaan dat je nooit klaar bent.
  • De Bevinding: Verschillende onderwijsmethoden hebben verschillende "learning rates" nodig. Bijvoorbeeld, een methode genaamd LoRA (die slechts een klein deel van het model verandert) kan een zeer snelle learning rate aan, terwijl een methode die twee dingen tegelijk leert een zeer langzaam, voorzichtig tempo nodig heeft.

2. De Zeven Onderwijsstrategieën

Zodra de drie checks hierboven zijn doorstaan, testten de onderzoekers zeven manieren om het nieuwe recept aan te leren:

1. Naive Fine-Tuning: "Blijf gewoon koken." Je neemt de hele chef en blijft hem trainen op de nieuwe data.
   • Resultaat: Geweldig voor het perfect leren van één specifiek gerecht. Maar als je deze chef later voor een ander type eten wilt gebruiken, kan het zijn dat hij zijn oude vaardigheden is vergeten (een probleem dat "catastrophic forgetting" wordt genoemd).
2. Layer Freezing: "Raak de basis niet aan." Je vergrendelt de kennis van de chef over basisvaardigheden zoals messenwerk en laat hem alleen de nieuwe saus leren.
   • Resultaat: Goed, maar soms te rigide. Het beperkt hoe goed de chef zich kan aanpassen aan de nieuwe ingrediënten.
3. LoRA (Low-Rank Adaptation): "Voeg een spiekbriefje toe." In plaats van het hele kookboek te herschrijven, voeg je een klein, efficiënt notitieblok toe aan de schort van de chef dat alleen de nieuwe regels bevat.
   • Resultaat: Zeer efficiënt en accuraat voor specifieke taken, vergelijkbaar met Naive tuning.
4. Multihead Replay: "De Dual-Head Chef." Je geeft de chef twee hoeden. Eén hoed is voor het nieuwe gerecht, en de andere hoed is voor de oude, vertrouwde gerechten. Hij oefent beide tegelijkertijd.
   • Resultaat: Dit is de winnaar op het gebied van veiligheid. Het is de enige methode die er consequent in slaagt om te voorkomen dat de chef zijn oude vaardigheden vergeet. Het houdt de chef goed in zowel het nieuwe gerecht als de oude gerechten.
5. Pseudolabel Replay: "De Synthetische Chef." In plaats van echte oude recepten te gebruiken, gebruik je de eigen voorspellingen van de chef over oude recepten om te oefenen.
   • Resultaat: Werkt goed en is flexibel omdat je niet de originele oude data nodig hebt, alleen het geheugen van de chef.
6. Replay + LoRA: Het combineren van het spiekbriefje met de twee hoeden.
   • Resultaat: Goed, maar de "Dual Head" alleen was vaak al voldoende.

3. De Belangrijkste Conclusies

• Verander het wiel niet opnieuw uit: Als je een model nodig hebt voor een specifieke, smalle taak (zoals alleen het simuleren van zout water), is Naive Fine-Tuning de snelste en makkelijkste manier om een geweldig resultaat te krijgen.
• Vergeet het verleden niet: Als je een model nodig hebt dat complexe, nieuwe situaties kan aan (zo zoals een nieuw type batterij of een complex biologisch molecuul) zonder zijn oorspronkelijke training te vergeten, moet je Multihead Replay gebruiken. Dit is de enige strategie die het model robuust en veilig hield tegen "vergeten".
• Kwaliteit boven trucjes: Het artikel benadrukt dat het besteden van tijd aan het kiezen van een hoogwaardig foundation model en het correct instellen van de energie-referenties belangrijker is dan het kiezen van de perfecte fine-tuning algoritme. Als de fundering zwak is of de wiskunde verkeerd is ingesteld, zal de beste onderwijsmethode ter wereld niet helpen.

Kortom: Om het beste AI-model voor chemie te krijgen, begin je met een slim fundament, stel je je wiskundige regels correct in, en als je wilt dat de AI veelzijdig is en niet vergeetachtig, leer je hem het "Dual Head" proces aan (Multihead Replay).

Technische samenvatting

Technische Samenvatting: Fine-tuning van MLIP Foundation Models

Probleemstelling
Machine-learned interatomaire potentiaal (MLIP) foundation models hebben het vermogen aangetoond om te transfereren naar diverse chemische systemen, wat een workflow biedt die het arbeidsintensieve proces van het vanaf nul trainen van taakspecifieke potentialen vermijdt. Echter, de gemeenschap mist systematische begeleiding over hoe en wanneer men deze modellen moet fine-tunen. Vroege rapporten suggereerden dat naïeve fine-tuning vaak leidt tot "catastrofale vergetelheid" (catastrophic forgetting), wat de adoptie van beperkende technieken (bijv. layer freezing, Low-Rank Adaptation) stimuleerde die oorspronkelijk voor large language models zijn ontwikkeld. Dit artikel onderzoekt of deze beperkingen noodzakelijk zijn of dat vroege mislukkingen te wijten waren aan andere factoren, zoals zwakkere foundation models, onjuiste initialisatie van de atomaire referentie-energie ($E_0$), of instabiele trainingsprocedures. De studie beoogt de belangrijkste factoren te karakteriseren die de uitkomsten van fine-tuning vormgeven, specifiek de nauwkeurigheid van de doeltaak en de robuustheid buiten de distributie (out-of-distribution, OOD).

Methodologie
De auteurs evalueren zeven verschillende fine-tuning strategieën over vijf chemisch diverse benchmarks, drie generaties foundation models, en trainingssets variërend over vijf ordes van grootte.

• Geëvalueerde Fine-tuning Strategieën:

  1. Naïef: Volledige parameterupdates via voortgezette gradiëntafdaling.
  2. Layer Freezing (Varianten): Het bevriezen van embedding/message-passing lagen terwijl de readouts worden getraind; of het bevriezen van de embedding en de eerste message-passing laag.
  3. Low-Rank Adaptation (LoRA): Het injecteren van trainbare low-rank decomposities in zowel scalaire als equivariante lineaire lagen terwijl de gepreënteerde gewichten bevroren blijven.
  4. Multihead Replay: Gelijktijdige optimalisatie op doeldata en een replay-dataset (uit pretraining of via pseudolabels) met behulp van afzonderlijke readout heads.
  5. Pseudolabel Replay: Een variant van multihead replay waarbij de replay-labels worden gegenereerd door het foundation model zelf, waardoor de bron van de replay wordt losgekoppeld van de oorspronkelijke pretraining corpus.
  6. Replay + LoRA: Het combineren van multihead replay met LoRA.

• Benchmarks: De studie beslaat systemen met een toenemende afwijking van het OMat24 pretraining domein (periodieke anorganische bulk):

  • Lithium argyrodiet elektrolyten (inorganische periodieke vaste stof).
  • Aquatisch NaCl (ionische oplossing).
  • Ijs-polymorfen (moleculaire vaste stof).
  • SN2 reacties (gasfase reactieve chemie).
  • SPICE biomoleculen (organische/biologische conformers).

• Technische Implementaties: De auteurs hebben drie nieuwe mogelijkheden geïmplementeerd in de MACE-codebase:

  • LoRA aangepast voor equivariante message-passing architecturen (zowel voor scalaire als equivariante lineaire lagen).
  • Pseudolabelled replay om de bronnen van de replay-data te ontkoppelen.
  • Model-bewuste her-schatting van de atomaire referentie-energie ($E_0$) om de gepreënteerde baselines af te stemmen op de doeldata.

• Evaluatiemetrieken: Naast standaard puntgewijze energie- en krachtfouten, onderzoekt de studie dynamische en extrapolatieve gedragingen, inclusief radiale distributiefuncties (RDFs) van moleculaire dynamica (MD), Nudged Elastic Band (NEB) reactieprofielen, MD-stabiliteitstests en Random Structure Search (RSS) om falen in kortetermijn-afstoting te detecten.

Belangrijkste Resultaten

1. Voorwaarden domineren de strategiekeuze: De studie vindt dat de kwaliteit van het foundation model, correcte $E_0$ initialisatie en goed gekozen hyperparameters randvoorwaarden zijn waarvan de impact routineus groter is dan die van de specifieke fine-tuning strategie.

   • Kwaliteit van het Foundation Model: Nieuwere foundation models (bijv. gebaseerd op OMat24) presteren consequent beter dan oudere modellen (gebaseerd op MPTraj) in OOD transfer, zelfs met vaste fine-tuning recepten.
   • $E_0$ Initialisatie: Het gebruik van "gemiddelde" $E_0$ waarden leidt tot aanzienlijk hogere fouten en MD-instabiliteit (bijv. ijsmodellen die binnen 50 ps falen). "Her-geschatte" $E_0$ waarden (het afstemmen van het nulpunt van het gepreënteerde model op de doeldata) zijn cruciaal voor stabiliteit en transfereerbaarheid, en leveren vaak betere resultaten op dan de keuze van het fine-tuning algoritme zelf.
   • Hyperparameters: Naïeve fine-tuning vereist lagere leersnelheden en een hogere EMA decay. LoRA verdraagt hogere leersnelheden. Multihead replay vereist substantieel lagere leersnelheden om concurrerende update-signalen te voorkomen. Weight decay moet op nul worden ingesteld om te voorkomen dat parameters weggetrokken worden van de gepreënteerde oplossing.

2. Prestaties per Doelstelling:

   • In-distributie Specialisatie (Enkel Systeem): Voor smalle taken (bijv. SN2 barrières, aquatisch NaCl solvatie) bereiken de meeste strategieën (Naïef, LoRA, Multihead) een sterke nauwkeurigheid en overtreffen zij consistent modellen die vanaf nul zijn getraind. Naïeve fine-tuning biedt de beste convergentie voor single-system toepassingen.
   • Out-of-Distribution Robuustheid: Bij evaluatie van transfer naar gerelateerde maar onbekende composities (bijv. niet-argyrodiet elektrolyten) of verschillende chemieën (bijv. biomoleculen), is Multihead Replay (met ofwel originele of pseudolabelled data) de enige aanpak die consequent OOD robuustheid behoudt. Het behoudt nauwkeurigheid op de pretraining distributie terwijl het de doeltaak leert, wat effectief catastrofale vergetelheid voorkomt.
   • Freezing en LoRA: Hoewel effectief voor parameterefficiëntie, vertoonden layer freezing en LoRA beperkingen in het aanpassen aan solvatie-kenmerken of het behouden van brede chemische robuustheid vergeleken met multihead replay in de geteste scenario's.

Betekenis en Claims
Het artikel claimt dat de waargenomen fragiliteit van naïeve fine-tuning in MLIPs grotendeels een resultaat is van een suboptimale opzet in plaats van een intrinsieke beperking van de methode. De auteurs stellen dat:

• Naïeve fine-tuning een levensvatbaar en vaak superieur startpunt is voor single-system toepassingen, mits het foundation model van hoge kwaliteit is en de $E_0$ waarden correct zijn her-geschat.
• Multihead replay de noodzakelijke strategie is voor bredere inzet waar het behoud van het gedrag van het foundation model buiten de fine-tuning distributie vereist is.
• Pseudolabelled replay een praktisch voordeel biedt door het gebruik van elke structureel diverse dataset voor replay mogelijk te maken, waardoor de afhankelijkheid van toegang tot de oorspronkelijke pretraining corpus wordt weggenomen.

Dit werk stelt vast dat voor praktijkgebruikers, het investeren in het sterkste beschikbare foundation model en het waarborgen van de correcte afstemming van de atomaire referentie-energie, kritiekere ontwerpkeuzes zijn dan het selecteren van een specifiek beperkt fine-tuning algoritme. De studie biedt een systematisch kader voor de inzet van MLIP foundation models, waarbij fine-tuning wordt verschoven van een niche optie naar een standaard startpunt voor systeem-specifieke ontwikkeling.