Systematic selection of surrogate models for nonequilibrium chemistry

¹ Interdisciplinary Center for Scientific Computing, Heidelberg University, Heidelberg, Germany
² Institute of Computer Engineering, Heidelberg University, Heidelberg, Germany

^★ Corresponding author: This email address is being protected from spambots. You need JavaScript enabled to view it.

Received: 12 January 2026
Accepted: 18 February 2026

Abstract

Context. Nonequilibrium chemistry is central to many astrophysical environments, but remains a major computational bottleneck in simulations because solving the associated stiff, coupled, ordinary differential equation systems is expensive. Neural surrogates promise substantial increases in speed, yet most existing studies are limited to proof-of-concept demonstrations and lack rigorous, dataset-grounded comparisons of architectures or systematic optimization toward accuracy and efficiency.

Aims. We aim to establish a principled procedure for optimizing and selecting surrogate models for astrochemistry that would enable representative and quantitative comparisons across architectures. This requires joint optimization for accuracy and efficiency, suitable metrics for performance assessment, and evaluation of surrogate reliability under practical constraints such as uncertainty quantification (UQ) and iterative prediction.

Methods. To this end, we employed CODES, a benchmarking framework that performs multi-objective hyperparameter tuning, trains optimized configurations, and evaluates their behavior across multiple dimensions. We compared four surrogate families, two fully connected models and two latent-evolution models. These surrogates were optimized and trained on four KROME-generated datasets spanning primordial and molecular-cloud chemistry, with up to 287 reactions across 37 species, including parametric variations in radiation field and metallicity. Each model predicted chemical abundances and temperature over a 10 kyr interval for user-specified output times.

Results. Dual-objective optimization reveals pronounced accuracy–efficiency trade-offs for all architectures and enables substantial efficiency gains with minimal accuracy loss. Across datasets, architectures group naturally by inductive bias: fully connected models, which impose minimal structural assumptions, achieve the highest accuracy and the most reliable UQ, but show the characteristic long-term error growth associated with low-bias models. Latent-evolution models – though less accurate – exhibit reduced error accumulation under iterative rollouts.

Conclusions. Our results underscore the importance of systematic optimization and comprehensive architectural comparison to make trade-offs explicit. The datasets, architectures, metrics, and benchmarking procedure are publicly bundled in CODES to support representative and reproducible comparisons.