How to Train a Shallow Ensemble

Dit artikel presenteert een efficiënt protocol voor het trainen van ondiepe ensemble-modellen voor interatomische potentialen, waarbij volledige fine-tuning van een probabilistisch getraind model de kalibratie van onzekerheid behoudt terwijl de trainingskosten tot 96% worden verlaagd.

De kern

Hoe je een "slimme groep" van AI-modellen traint zonder je computer te laten smelten

Stel je voor dat je een superkrachtige AI hebt die kan voorspellen hoe atomen zich gedragen. Dit is handig voor het ontwerpen van nieuwe medicijnen of batterijen. Maar AI-modellen maken soms fouten. De grote vraag is: hoe weet de AI zelf of ze zeker is van haar antwoord, of dat ze eigenlijk aan het gokken is?

In dit artikel onderzoeken onderzoekers een slimme manier om deze "onzekerheid" te meten. Ze gebruiken een techniek die lijkt op het vragen van advies aan een groep experts in plaats van één persoon.

Hier is de uitleg in gewone taal, met een paar creatieve vergelijkingen:

1. Het Probleem: De "Zekere" Fout

Normaal gesproken trainen we een AI met één doel: zo min mogelijk fouten maken. Maar als je een AI vraagt: "Hoe heet het weer?", en het antwoordt "20 graden", weten we niet of het zeker 20 graden is, of dat het AI-model gewoon gissen doet.

Voor wetenschappers is dit gevaarlijk. Als de AI denkt dat ze zeker is, maar eigenlijk helemaal niet, kunnen ze een dure experimentele batterij bouwen die nooit werkt. Ze hebben een manier nodig om te zeggen: "Ik denk dat het 20 graden is, maar ik heb een 50% kans dat ik het mis heb."

2. De Oplossing: Een "Ondiepe" Groep (Shallow Ensemble)

De traditionele manier om dit op te lossen is door 10 of 20 verschillende AI-modellen te trainen en hun antwoorden te middelen. Als ze allemaal hetzelfde zeggen, ben je zeker. Als ze het oneens zijn, weet je dat er onzekerheid is.

Het probleem: Dit is extreem duur en traag. Het is alsof je 20 verschillende ingenieurs moet inhuren, elk met hun eigen kantoor en gereedschap, om één brug te ontwerpen. Dat kost te veel tijd en geld.

De slimme truc: De onderzoekers gebruiken een "Shallow Ensemble" (een ondiepe groep).

• De Analogie: Stel je voor dat je één groot, slim ingenieursbureau hebt (de "ruggengraat" van het model). Dit bureau kent de basisregels van de natuur perfect.
• In plaats van 20 hele nieuwe bureaus te bouwen, nemen ze één bureau en laten ze slechts de laatste stap (de "uitvoer") variëren.
• Het is alsof je één groot team hebt, maar je laat 5 verschillende managers de uiteindelijke beslissing nemen. Ze gebruiken allemaal dezelfde kennis, maar ze interpreteren het net iets anders. Dit is veel sneller en goedkoper.

3. Het Nieuwe Inzicht: Niet alleen Energie, maar ook Krachten

De onderzoekers ontdekten twee belangrijke dingen:

• Punt 1: Je moet de AI leren om "twijfel" te voelen.
  Als je de AI alleen traint om de energie (hoe stabiel een molecule is) goed te voorspellen, is ze goed in energie, maar slecht in het voorspellen van krachten (hoe atomen bewegen).

  • Vergelijking: Het is alsof je een autochauffeur traint om alleen de snelheid perfect te houden, maar je vergeet te zeggen dat hij ook moet opletten of de remmen werken. Hij rijdt snel, maar als hij moet remmen, is hij paniekerig en onzeker.
  • Oplossing: Je moet de AI expliciet trainen om ook de onzekerheid van de krachten te berekenen.

• Punt 2: Het "Nieuwe" is niet altijd beter.
  Je zou denken: "Laten we gewoon 20 modellen van nul af trainen." Dat werkt het beste, maar het duurt te lang.
  De onderzoekers vonden een hack:

  1. Train eerst één goed model (of een kleine groep) op energie.
  2. Gebruik dit als startpunt.
  3. Fijnafstemmen (Fine-tuning): Laat de AI alleen nog maar een beetje "nadenken" over de onzekerheid.
  • Vergelijking: In plaats van 20 nieuwe ingenieursbureaus te bouwen, neem je één ervaren bureau, laat je ze een paar dagen extra nadenken over de "wat als"-scenario's, en klaar. Je bespaart tot 96% van de tijd, terwijl de kwaliteit bijna hetzelfde blijft.

4. Waarom is dit belangrijk?

Dit onderzoek geeft een handleiding voor wetenschappers:

• Als je tijd en geld hebt: Train een "Shallow Ensemble" met een speciale formule die zowel energie als krachten en hun onzekerheid meeneemt.
• Als je haast hebt (wat vaak het geval is): Neem een bestaand model, laat het een beetje "fijnafstemmen" op onzekerheid, en je hebt een betrouwbare AI die weet wanneer ze moet zeggen: "Ik weet het niet zeker, probeer het maar niet."

Samenvattend

De onderzoekers hebben een manier gevonden om AI-modellen voor atomen slimmer en veiliger te maken zonder dat het je maanden van rekentijd kost. Ze hebben ontdekt dat je de AI moet leren om niet alleen het antwoord te geven, maar ook om te zeggen: "Hoe zeker ben ik hierover?" En dat kun je doen door een slimme groep te laten samenwerken, in plaats van honderden individuen te trainen.

Dit zorgt ervoor dat we in de toekomst sneller nieuwe materialen en medicijnen kunnen vinden, zonder dat we bang hoeven te zijn dat de AI ons een vals vertrouwen geeft.

Technische samenvatting

Titel: Hoe een Shallow Ensemble te trainen (How to Train a Shallow Ensemble)

Auteurs: Moritz Schäfer, Matthias Kellner, Johannes Kästner, Michele Ceriotti
Datum: 18 februari 2026

---

1. Probleemstelling

Machine Learning Interatomic Potentials (MLIPs) bieden een snellere alternatief voor first-principles methoden (zoals DFT), maar introduceren onzekerheid in hun voorspellingen. Deze onzekerheid kan voortkomen uit beperkte kennis (epistemische), ruis in de data (aleatorisch) of modelmisspecificatie. Voor betrouwbare toepassingen, zoals actief leren (active learning) en het propageren van fouten naar thermodynamische grootheden, is het cruciaal dat deze onzekerheid gekalibreerd is (d.w.z. dat de voorspelde onzekerheid overeenkomt met de werkelijke fout).

Traditionele methoden voor onzekerheidskwantificatie (UQ), zoals diepe ensembles (Deep Ensembles) waarbij elk model onafhankelijk wordt getraind, zijn computatieel te duur voor grootschalige atomaire simulaties. Alternatieven zoals Shallow Ensembles (waarbij alleen de laatste laag wordt geënssembleerd) en Last-Layer Prediction Rigidity (LLPR) (een Laplace-benadering) zijn efficiënter, maar er is onduidelijkheid over:

1. Hoe deze methoden het beste getraind moeten worden voor zowel energie- als krachten-onzekerheid.
2. Of bestaande methoden (zoals LLPR) voldoende zijn of dat expliciete training voor onzekerheid nodig is.
3. Hoe men de hoge rekentijd van het trainen van probabilistische modellen voor krachten kan reduceren zonder in te leveren op kwaliteit.

2. Methodologie

De auteurs onderzoeken systematisch trainingsstrategieën voor Shallow Ensembles op basis van neurale netwerken (zoals GMNN, So3krates en EquivMP). Ze vergelijken verschillende benaderingen:

• Shallow Ensembles (DPOSE): Een ensemble van meerdere output-eenheden in de laatste laag, getraind met een gezamenlijke Negative Log-Likelihood (NLL) loss. Dit zorgt ervoor dat onzekerheidsinformatie doorwerkt naar de diepere lagen.
• LLPR (Last-Layer Prediction Rigidity): Een post-hoc methode die een Laplace-benadering toepast op de posterior van de laatste laag van een model dat getraind is met een standaard Mean Squared Error (MSE) loss.
• Trainingsdoelen:
  • Energie alleen: Training met NLL voor energie en MSE voor krachten.
  • Volledig Probabilistisch: Training met NLL voor zowel energie als krachten (vereist het berekenen van Jacobianen, wat duur is).
• Efficiënte Protocollen: Om de hoge kosten van het trainen van probabilistische krachten te omzeilen, testen de auteurs een fine-tuning strategie. Hierbij wordt een ensemble geïnitieerd vanuit een reeds getraind MSE-model (via LLPR of een hybride loss) en vervolgens fijn afgestemd (fine-tuned) met een gezamenlijke NLL-loss voor energie en krachten. Ze vergelijken Last-Layer Fine-Tuning (alleen de laatste laag aanpassen) met Full-Model Fine-Tuning (alle gewichten aanpassen).

3. Belangrijkste Bijdragen en Resultaten

A. Calibratie van Energie- en Krachtenonzekerheid

• Energie: Shallow Ensembles getraind met NLL voor energie leveren over het algemeen betere gekalibreerde onzekerheid op dan LLPR (gebaseerd op MSE). LLPR faalt vaak bij complexe datasets (zoals ionische vloeistoffen) omdat de features van het vooraf getrainde model "stijf" zijn en uitschieters niet goed onderscheiden.
• Krachten: Het blijkt dat het trainen van een model puur op energie-onzekerheid niet leidt tot betrouwbare krachten-onzekerheid.
  • Bij datasets zoals BMIM (ionische vloeistof) treedt er een element-specifieke miscalibratie op: het model is te zelfverzekerd voor bepaalde atomen (bijv. Fluorine en Boron) en niet voor anderen.
  • Conclusie: Het expliciet modelleren van krachten-onzekerheid via een NLL-loss voor krachten is essentieel voor betrouwbare kalibratie.

B. De Kosten van Probabilistische Krachttraining

Het trainen van een ensemble van scratch met een NLL-loss voor krachten is computatieel zeer intensief omdat de Jacobiaan van de energie naar de coördinaten voor elk ensemblelid berekend moet worden. Dit kan de trainingstijd met een factor 17 verhogen (afhankelijk van de architectuur).

C. Het Efficiënte Protocol (Fine-Tuning)

De auteurs introduceren een efficiënt protocol dat de kosten drastisch verlaagt zonder kwaliteit te verliezen:

1. Train een basismodel (of een Shallow Ensemble) met een hybride loss (NLL voor energie + MSE voor krachten) of gebruik een bestaand MSE-model.
2. Initieer het ensemble door de laatste laag te samplen uit de LLPR-posterior of door een Shallow Ensemble te starten.
3. Voer Full-Model Fine-Tuning uit met de gezamenlijke NLL-loss (energie + krachten).

Resultaten van dit protocol:

• Kwaliteit: De kalibratiekwaliteit (gemeten in Relative Log-Likelihood, RLL) is bijna identiek aan het trainen van een ensemble van scratch.
• Efficiëntie: De trainingstijd wordt met tot 96% gereduceerd.
• Robuustheid: Full-Model Fine-Tuning is noodzakelijk om de "feature rigidity" van LLPR te doorbreken. Alleen de laatste laag aanpassen (Last-Layer FT) is vaak onvoldoende voor complexe systemen zoals BMIM, omdat de onderliggende features dan niet aangepast worden aan de nieuwe onzekerheidsverdeling.

4. Significantie en Praktische Richtlijnen

Dit werk biedt praktische richtlijnen voor het trainen van MLIPs met betrouwbare onzekerheidskwantificatie:

1. Voor nieuwe modellen: Train een Shallow Ensemble met een hybride loss (NLL energie + MSE krachten) en voer daarna Full-Model Fine-Tuning uit met een gezamenlijke NLL-loss. Dit is de meest betrouwbare en kostenefficiënte methode.
2. Voor bestaande foundation modellen: Gebruik een LLPR-gebaseerde initialisatie (gebaseerd op het MSE-model) en voer Full-Model Fine-Tuning uit met de gezamenlijke NLL-loss. Dit corrigeert de miscalibratie die vaak optreedt bij post-hoc LLPR-toepassingen.
3. Krachten zijn cruciaal: Het negeren van probabilistische training voor krachten leidt tot systematische fouten in de onzekerheidsschatting, wat kritiek is voor toepassingen zoals actief leren.
4. Implementatie: De auteurs hebben hun methoden geïmplementeerd in de apax en IPSuite softwarestacks en maken de code openbaar beschikbaar.

Conclusie:
De studie toont aan dat "Shallow Ensembles" een krachtige, schaalbare oplossing zijn voor onzekerheidskwantificatie in atomaire machine learning. Door slimme initialisatie en full-model fine-tuning kunnen onderzoekers de hoge kosten van probabilistische training vermijden, terwijl ze toch nauwkeurige en gekalibreerde onzekerheidsschattingen voor zowel energie als krachten behouden. Dit maakt MLIPs veiliger en betrouwbaarder voor productie-simulaties.