A Kolmogorov-Arnold Surrogate Model for Chemical Equilibria: Application to Solid Solutions

Deze studie toont aan dat Kolmogorov-Arnold-netwerken, die learnable spline-functies gebruiken in plaats van vaste activatiefuncties, een nauwkeurigere en efficiëntere data-gedreven surrogate zijn dan multilayer perceptrons voor het modelleren van chemische evenwichten, met name bij de berekening van oplosbaarheden van radionuclide-dragende vaste oplossingen voor de veiligheidsbeoordeling van geologische afvalbergingen.

De kern

Stel je voor dat je een gigantische, ingewikkelde chemische machine probeert te simuleren. Deze machine bepaalt hoe chemicaliën in de grond reageren, bijvoorbeeld wanneer radioactief afval wordt opgeslagen in een diepe ondergrondse berging. Om te weten of dit veilig is, moeten wetenschappers deze machine miljarden keren laten draaien in een computer.

Het probleem? De computer is zo traag. Het kost hem vaak duizenden keren langer om de chemische reacties uit te rekenen dan om de stroming van water te simuleren. Het is alsof je een racewagen probeert te besturen, maar de motor (de chemie) is gemaakt van lood en de wielen (de waterstroming) van rubber. Je komt er niet.

De oplossing: Een slimme "tussenpersoon"

In dit onderzoek proberen de auteurs (Leonardo Boledi en zijn team) dit probleem op te lossen met een nieuwe soort kunstmatige intelligentie, genaamd KAN (Kolmogorov-Arnold Netwerk).

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De oude methode: De zware "Rekenmachine" (MLP)

Vroeger gebruikten wetenschappers een standaard type AI, een "Meerlaagse Perceptron" (MLP).

• De analogie: Stel je voor dat je een chef-kok bent die een recept moet onthouden. De oude AI is als een kok die alles uit zijn hoofd moet leren door duizenden keren te oefenen. Hij is goed, maar hij moet heel veel informatie onthouden (parameters) om precies te weten wat er gebeurt als je een beetje meer zout of suiker toevoegt.
• Het nadeel: Hij is traag om te trainen en maakt soms nog steeds fouten bij complexe situaties.

2. De nieuwe methode: De "Slimme Schat" (KAN)

De auteurs testen nu de nieuwe KAN-technologie.

• De analogie: In plaats van alles uit het hoofd te leren, heeft deze nieuwe AI een magisch, rekbaar meetlint (een 'spline') op zijn duim. Als je vraagt: "Wat gebeurt er als ik meer zout toevoeg?", trekt hij niet aan een vaste formule, maar past hij zijn meetlint direct aan de situatie aan.
• Het voordeel: Hij heeft veel minder geheugen nodig (minder parameters) om even goed of zelfs beter te zijn. Hij is flexibeler en sneller in het vinden van de juiste oplossing, net als iemand die een kaart gebruikt in plaats van alles uit het hoofd te kennen.

Wat hebben ze getest?

Ze hebben deze nieuwe AI getest op twee grote uitdagingen:

1. Cement (De "Schoolvoorbeeld" test):
   Ze lieten de AI kijken naar hoe cement veroudert. Het resultaat? De nieuwe KAN-methode maakte 62% minder fouten dan de oude methode, terwijl hij minder "hersencellen" gebruikte. Het was alsof je een slimme student hebt die met een dun boekje beter scoort dan een student met een dik handboek.

2. Radioactief Afval (De "Echte" test):
   Dit is het belangrijkste deel. Ze keken naar hoe radioactief elementen (zoals Radium) zich vastzetten in kristallen in de grond (zoutoplossingen). Dit is extreem complex omdat de temperatuur en de mengsels veranderen.

   • Ze testten dit van simpele mengsels tot zeer ingewikkelde driedelige kristallen.
   • Het resultaat: De KAN-AI was zo nauwkeurig dat hij in bijna 100% van de gevallen de juiste uitkomst gaf. De fouten waren zo klein dat ze nauwelijks meetbaar waren.

Waarom is dit een doorbraak?

Stel je voor dat je een berging voor radioactief afval moet ontwerpen die 100.000 jaar veilig moet blijven.

• Met de oude methode zou het duizenden jaren duren om alle mogelijke scenario's te simuleren.
• Met deze nieuwe KAN-methode kunnen ze diezelfde berekeningen 16 keer sneller doen, met minder fouten.

De prijs:
De enige keer dat de nieuwe methode "traag" is, is tijdens het leren (trainen). Het kost de computer ongeveer 10 minuten langer om te leren dan de oude methode. Maar dat is een eenmalige investering. Zodra hij geleerd heeft, is hij razendsnel en supernauwkeurig.

Conclusie

Dit onderzoek laat zien dat we met deze nieuwe "meetlint-AI" (KAN) veel sneller en veiliger kunnen voorspellen hoe radioactief afval zich in de aarde gedraagt. Het is alsof we van een langzame, zware stoomlocomotief zijn overgestapt op een snelle, wendbare elektrische trein. Dit helpt ons om de aarde veiliger te maken voor de toekomst.

Technische samenvatting

Probleemstelling

Geochemische oplosmiddelen (solvers) vormen een grote rekenkundige uitdaging, vooral binnen reactief transportmodelleren (RTM). In simulaties voor diepe geologische opslag van radioactief afval moeten chemische evenwichtsberekeningen miljarden keren worden uitgevoerd om de interactie tussen vloeistofstromen en geochemische reacties te modelleren.

• Rekenkosten: De chemische solver is vaak de duurste component van de simulatie, met kosten die tot 10.000 keer hoger kunnen zijn dan die van de vloeistofstroom-solver.
• Beperkingen van bestaande methoden: Hoewel er reeds machine learning (ML) benaderingen zijn gebruikt om deze kosten te verlagen (bijv. via Multilayer Perceptrons of MLPs), blijft de zoektocht naar de meest efficiënte en accurate datagedreven surrogate-modellen (vervangers) actueel. Bestaande diepe neurale netwerken hebben vaak veel parameters nodig en vertonen soms hoge voorspellingsfouten bij complexe niet-lineaire relaties.

Methodologie

De auteurs onderzoeken de toepasbaarheid van Kolmogorov-Arnold Netwerken (KANs) als surrogate-modellen voor chemische evenwichten, in vergelijking met traditionele MLPs.

1. Architectuur van KANs:

   • In tegenstelling tot MLPs, waar vaste activeringsfuncties op de knooppunten (neurons) worden toegepast en de gewichten worden aangepast, plaatsen KANs leerbare spline-gebaseerde functies op de randen (edges) tussen de neurons.
   • Dit is gebaseerd op het Kolmogorov-Arnold representatietheorema, dat stelt dat elke multivariate continue functie kan worden geschreven als een eindige samenstelling van univariate functies.
   • Deze structuur maakt het mogelijk om complexe niet-lineaire relaties met minder trainbare parameters nauwkeuriger te modelleren.

2. Chemische Systemen en Data-Generatie:

   • De trainingdata werd gegenereerd met de solver GEM-Selektor (gebaseerd op Gibbs-energie-minimalisatie).
   • Er werden vier testgevallen onderzocht met toenemende thermodynamische complexiteit:
     • Benchmarks: Hydratatie van een cementstelsel (CaO-SiO2-H2O).
     • Toepassing 1: Mechanische menging van RaSO4 (ideale menging zonder echte vaste oplossing).
     • Toepassing 2: Een niet-ideale binaire vaste oplossing (Ba,Ra)SO4.
     • Toepassing 3: Een niet-ideale ternaire vaste oplossing (Sr,Ba,Ra)SO4, inclusief temperatuurafhankelijkheid.
   • De datasets omvatten Sobol-steekproeven over verschillende inputvariabelen (concentraties, temperatuur, ionsterkte).

3. Training en Vergelijking:

   • De modellen werden getraind in PyTorch Lightning met Optuna voor hyperparameter-optimalisatie.
   • De prestaties werden vergeleken op basis van absolute fout (RMSE), relatieve fout (RRMSE) en trainingsduur.

Belangrijkste Bijdragen

• Eerste toepassing voor Radionucliden: Dit is het eerste werk dat datagedreven surrogate-modellen onderzoekt voor de co-neerslag van radionucliden (Ra) in vaste oplossingen, variërend van eenvoudige mengsels tot complexe ternaire systemen.
• Vergelijking KAN vs. MLP: Het biedt een rigoureuze vergelijking tussen de nieuwe KAN-architectuur en de standaard MLP voor geochemische toepassingen.
• Temperatuurafhankelijkheid: Voor het eerst wordt temperatuur expliciet meegenomen als een variabele in het surrogate-model voor radium-opname, wat essentieel is voor realistische scenario's.
• Open Data: De datasets, inputbestanden en scripts zijn openbaar beschikbaar gesteld.

Resultaten

1. Nauwkeurigheid (Cement Benchmark):

   • KANs presteerden significant beter dan MLPs.
   • Verbetering: De absolute fout (RMSE) werd met 62% verlaagd en de relatieve fout (RRMSE) met 59%.
   • KANs behaalden deze resultaten zelfs met ongeveer een derde van het aantal trainbare parameters vergeleken met de referentie-MLP.

2. Nauwkeurigheid (Radium Vaste Oplossingen):

   • Voor de binaire en ternaire systemen hielden de KANs mediane voorspellingsfouten dicht bij 1 × 10⁻³.
   • Geen enkele voorspelling overschreed de drempel van 10% fout, zelfs niet in het meest complexe ternaire systeem (Sr,Ba,Ra)SO4.
   • KANs vertoonden aanzienlijk minder "hoog-fout" voorspellingen dan MLPs.

3. Rekenefficiëntie:

   • Trainingstijd: KANs zijn trager om te trainen (tot 4x langer dan MLPs), wat echter een eenmalige kostenpost is.
   • Evaluatiesnelheid: Bij het uitvoeren van 5.000 evenwichtsberekeningen boekten de KAN-surrogates een snelheidswinst van 87% tot 93% ten opzichte van de originele GEM-Selektor solver. Dit komt neer op een versnellingfactor van ongeveer 16x.

4. Stabiliteit:

   • In tegenstelling tot eerdere rapporten over instabiliteit bij het oplossen van partiële differentiaalvergelijkingen, vertoonden de KANs in deze geochemische toepassingen geen trainingsinstabiliteit of divergentie.

Betekenis en Conclusie

Dit onderzoek toont aan dat Kolmogorov-Arnold Netwerken een veelbelovende technologie zijn om reactief transportmodelleren te versnellen. De combinatie van superieure nauwkeurigheid, parameter-efficiëntie (minder geheugenvereisten) en extreme snelheidswinst tijdens de evaluatie maakt KANs ideaal voor het vervangen van dure geochemische solvers.

Hoewel de trainingskosten iets hoger zijn, is dit een verwaarloosbare investering binnen de totale modelleringstijd, vooral gezien de enorme besparing bij het uitvoeren van miljoenen simulaties voor veiligheidsbeoordelingen van diepe geologische opslagplaatsen. De volgende stap in het onderzoek zal zijn om deze surrogate-modellen te integreren in tijdsafhankelijke simulaties, waarbij aandacht moet worden besteed aan het handhaven van massabehoud en het voorkomen van foutenaccumulatie over de tijd.