cs.LGMar 16, 2026

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

Authors: Leonardo Boledi, Dirk Bosbach, Jenna Poonoosamy

Organizations: Institute of Fusion Energy and Nuclear Waste Management (IFN-2), Forschungszentrum Jülich GmbH

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

The computational cost of geochemical solvers is a challenging matter. For reactive transport simulations, where chemical calculations are performed up to billions of times, it is crucial to reduce the total computational time. Existing publications have explored various machine-learning approaches to determine the most effective data-driven surrogate model. In particular, multilayer perceptrons are widely employed due to their ability to recognize nonlinear relationships. In this work, we focus on the recent Kolmogorov-Arnold networks, where learnable spline-based functions replace classical fixed activation functions. This architecture has achieved higher accuracy with fewer trainable parameters and has become increasingly popular for solving partial differential equations. First, we train a surrogate model based on an existing cement system benchmark. Then, we move to an application case for the geological disposal of nuclear waste, i.e., the determination of radionuclide-bearing solids solubilities. To the best of our knowledge, this work is the first to investigate co-precipitation with radionuclide incorporation using data-driven surrogate models, considering increasing levels of thermodynamic complexity from simple mechanical mixtures to non-ideal solid solutions of binary (Ba,Ra)SO4_4 and ternary (Sr,Ba,Ra)SO4_4 systems. On the cement benchmark, we demonstrate that the Kolmogorov-Arnold architecture outperforms multilayer perceptrons in both absolute and relative error metrics, reducing them by 62% and 59%, respectively. On the binary and ternary radium solid solution models, Kolmogorov-Arnold networks maintain median prediction errors near 1×10−31\times10^{-3}. This is the first step toward employing surrogate models to speed up reactive transport simulations and optimize the safety assessment of deep geological waste repositories.

Figures & tables

ModelBaSOX4\text{BaSO}{\vphantom{\text{X}}}_{\smash[t]{\text{4}}} ( \SIUnitSymbolMicromol\mathrm{\SIUnitSymbolMicro mol} )NaCl ( mmol\mathrm{mmol} )RaBrX2\text{RaBr}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}} ( \SIUnitSymbolMicromol\mathrm{\SIUnitSymbolMicro mol} )SrSOX4\text{SrSO}{\vphantom{\text{X}}}_{\smash[t]{\text{4}}} ( mmol\mathrm{mmol} )T ( ∘ C)
(i)50 – 50050 – 50050 – 500-25.0
(ii)50 – 5005050 – 500-20.0 – 90.0
(iii)50 – 50050 – 50050 – 5005 – 5025.0
Table 1: Input variables for the GEM-Selektor simulations. Note that we consider different temperature values only in case (ii).
Figure 1: Sketch of an MLP (left) and a KAN (right) with two hidden layers. In KANs, instead of training for weights and biases, learnable activation functions are placed on each edge.
Figure 2: RMSE on the test set for the cement hydration case. The reference MLP model (blue) is shown against two KANs of different sizes (green and orange).
VariableRMSE MLPRMSE KANImprovement (%)
pH3.3071×10−33.3071\times 10^{-3}1.8907×10−31.8907\times 10^{-3}42.83
MassWater4.1398×10−44.1398\times 10^{-4}1.5640×10−41.5640\times 10^{-4}62.22
Ca(aq)5.3386×10−65.3386\times 10^{-6}1.8701×10−61.8701\times 10^{-6}64.97
Si(aq)2.5183×10−62.5183\times 10^{-6}1.0015×10−61.0015\times 10^{-6}60.23
O(aq)1.0262×10−51.0262\times 10^{-5}3.9045×10−63.9045\times 10^{-6}61.95
H(aq)1.0571×10−51.0571\times 10^{-5}3.9398×10−63.9398\times 10^{-6}62.73
Table 2: Prediction errors on the test set for the cement hydration test case. The MLP and KAN models with 1.5\times105 1.5\text{\times}{10}^{5}\text{\,} parameters are shown, together with the percentage improvement of the RMSE.
Figure 3: RRMSE on the test set for the cement hydration case. The reference MLP model (blue) is shown against two KANs of different sizes (green and orange).
VariableRRMSE MLPRRMSE KANImprovement (%)
pH2.9427×10−42.9427\times 10^{-4}1.7009×10−41.7009\times 10^{-4}42.20
MassWater6.3279×10−36.3279\times 10^{-3}1.9846×10−31.9846\times 10^{-3}68.64
Ca(aq)5.4218×10−35.4218\times 10^{-3}2.2426×10−32.2426\times 10^{-3}58.64
Si(aq)1.1215×10−21.1215\times 10^{-2}5.1705×10−35.1705\times 10^{-3}53.90
O(aq)6.7081×10−36.7081\times 10^{-3}2.3244×10−32.3244\times 10^{-3}65.35
H(aq)5.3405×10−35.3405\times 10^{-3}2.3007×10−32.3007\times 10^{-3}56.92
Table 3: Prediction errors on the test set for the cement hydration test case. The MLP and KAN models with 1.5\times105 1.5\text{\times}{10}^{5}\text{\,} parameters are shown, together with the percentage improvement of the RRMSE.
ModelTrainable parametersTraining time ( s\mathrm{s} )
MLP5.6×1045.6\times 10^{4}153
MLP1.5×1051.5\times 10^{5}162
KAN5.6×1045.6\times 10^{4}600
KAN1.5×1051.5\times 10^{5}968
Table 4: Comparison of training times between KANs and MLPs for the cement hydration test case across different network sizes.
Figure 4: Relative error on the test set for the radium uptake case with mechanical mixing. The errors are plotted for three KANs trained on datasets of different sizes, where mm denotes the exponent of the Sobol sampler. In red, we show the number of predictions with an error above 10%10\% .
Figure 5: Relative error on the test set for the radium uptake case with mechanical mixing. The classical MLP model (blue) is shown against two KANs of different sizes (green and orange). In red, we show the number of predictions with an error above 10%10\% .
Figure 6: Relative error on the test set for the radium uptake case with the two-component SS model.
Figure 7: Relative error on the test set for the radium uptake case with the three-component SS model.
ModelRun time GEMS (s)Run time KAN (s)Improvement (%)
(i) Mech. mix2.450.17692.82
(ii) Binary SS1.290.16287.44
(iii) Ternary SS4.130.26093.71
Table 5: Comparison of run times between GEM-Selektor and the KAN surrogates for the radium uptake case studies. The time required to perform 5,000 equilibria calculations is displayed.

Explore similar work

CardsList

  1. Efficient Neural Surrogates for Linear Radiation Transport on the Lattice and Hohlraum benchmarks

    Oct 3, 2026Carmelo Gonzales, Steffen Schotthöfer, Cory D. HauckNeural Surrogate ModelingPDE Surrogate Modeling

    NVIDIA, Santa Clara, US · Seamless Labs, Mountain View, US · Oak Ridge National Laboratory

  2. A numerical study into neural network surrogate model performance for uncertainty propagation

    May 15, 2026Noah Wade, Kirubel TeferraNeural Surrogate ModelingUncertainty Quantification

    ASEE Postdoctoral Associate, U.S. Naval Research Laboratory, Washington, DC, 20375, USA · U.S. Naval Research Laboratory, Washington, 20375, DC, USA

  3. SechKAN: Kolmogorov-Arnold Networks with Hyperbolic Secant Functions

    Jun 30, 2026Hoang-Thang TaNeural Surrogate ModelingNeural Network Activation Functions

    University of Information Technology, Ho Chi Minh City, Vietnam · Vietnam National University Ho Chi Minh City, Ho Chi Minh City, Vietnam