AutoSurrogate: An LLM-Driven Multi-Agent Framework for Autonomous Construction of Deep Learning Surrogate Models in Subsurface Flow

Jiale Liu; Nanzhe Wang

cs.LGApr 13, 2026

AutoSurrogate: An LLM-Driven Multi-Agent Framework for Autonomous Construction of Deep Learning Surrogate Models in Subsurface Flow

Authors: Jiale Liu, Nanzhe Wang

Organizations: School of Physics and Astronomy, The University of Edinburgh, Edinburgh, EH9 3FD, United Kingdom · Institute of GeoEnergy Engineering, School of Energy, Geoscience, Infrastructure and Society, Heriot-Watt University, Edinburgh, EH14 4AS, United Kingdom · Subsurface Energy Transition and Innovation Centre, Heriot-Watt University, Edinburgh, EH14 4AS, United Kingdom

Abstract

High-fidelity numerical simulation of subsurface flow is computationally intensive, especially for many-query tasks such as uncertainty quantification and data assimilation. Deep learning (DL) surrogates can significantly accelerate forward simulations, yet constructing them requires substantial machine learning (ML) expertise - from architecture design to hyperparameter tuning - that most domain scientists do not possess. Furthermore, the process is predominantly manual and relies heavily on heuristic choices. This expertise gap remains a key barrier to the broader adoption of DL surrogate techniques. For this reason, we present AutoSurrogate, a large-language-model-driven multi-agent framework that enables practitioners without ML expertise to build high-quality surrogates for subsurface flow problems through natural-language instructions. Given simulation data and optional preferences, four specialized agents collaboratively execute data profiling, architecture selection from a model zoo, Bayesian hyperparameter optimization, model training, and quality assessment against user-specified thresholds. The system also handles common failure modes autonomously, including restarting training with adjusted configurations when numerical instabilities occur and switching to alternative architectures when predictive accuracy falls short of targets. In our setting, a single natural-language sentence can be sufficient to produce a deployment-ready surrogate model, with minimum human intervention required at any intermediate stage. We demonstrate the utility of AutoSurrogate on a 3D geological carbon storage modeling task, mapping permeability fields to pressure and CO $_2$ saturation fields over 31 timesteps. Without any manual tuning, AutoSurrogate is able to outperform expert-designed baselines and domain-agnostic AutoML methods, demonstrating strong potential for practical deployment.

Figures & tables

Figure 1: Three approaches to constructing deep learning surrogates for subsurface flow. Approach 1: manual, expert-driven pipeline relying on trial-and-error design and tuning. Approach 2: domain-agnostic AutoML, which performs brute-force search without exploiting physical priors. Approach 3: the proposed AutoSurrogate , an LLM-driven multi-agent framework that takes a dataset and a natural-language instruction and autonomously delivers a deployment-ready surrogate through physics-aware reasoning and closed-loop self-correction.

Figure 2: Overall architecture of the AutoSurrogate framework. Given simulation data and a natural-language instruction, four specialized agents collaborate through a shared memory context to autonomously produce deployment-ready surrogate models. The HPO & Training Agent incorporates a closed-loop self-correction mechanism that handles training instabilities and suboptimal convergence through continuation, stability-constrained restart, or architecture switching.

Category	Hyperparameter	Applies to
Optimization	Learning rate	All
	Weight decay	All
	Batch size	All
Capacity	Base channel width	All
Architecture-specific	Transformer layers	CNN-Transformer
Architecture-specific	Transformer heads	CNN-Transformer

Table 1: Searchable hyperparameters exposed by the AutoSurrogate model zoo. The common set applies to all eight architectures in the zoo, whereas architecture-specific hyperparameters are additionally exposed only for the indicated model family. Concrete search ranges are not fixed a priori and are instead determined by the HPO & Training Agent based on the task profile and the selected architecture.

Figure 3: Schematic of the geological model for CO 2 storage, showing the full domain and the central storage aquifer.

Rservoir parameters	Values
Caprock thickness	1905 m
Caprock porosity	0.08
Caprock permeability (horizontal and vertical)	0.001 mD
Basement thickness	500 m
Basement porosity	0.09
Basement permeability (horizontal and vertical)	2.3 mD

Table 2: Parameters utilized for CO 2 storage simulation.

Method	Pressure			Saturation
Method	$R^{2}$	RMSE	Rel $L_{2}$	$R^{2}$	RMSE	Rel $L_{2}$
UDeepONet3D	0.8534	2.3697	0.0114	0.8249	0.0139	0.4144
CNNTransformer3D	0.9861	0.7300	0.0035	0.9098	0.0100	0.2974
EDConvLSTM3D	0.9887	0.6577	0.0032	0.9309	0.0088	0.2604
UFNO3D	0.9911	0.5853	0.0028	0.9102	0.0100	0.2967
RecurrentRUNet3D	0.9972	0.3290	0.0016	0.9359	0.0084	0.2508

Table 3: Quantitative comparison of surrogate models for pressure and CO 2 saturation prediction. Best results are highlighted in bold .

Figure 4: Performance–efficiency Pareto frontier for pressure and CO 2 saturation prediction. The step line connects Pareto-optimal baseline configurations. AutoSurrogate @1 (star) and @3 extend the frontier beyond all baselines. AutoML methods are fast but achieve substantially lower

R^{2}

, especially on saturation where they fall below hand-tuned baselines. UDeepONet3D and FNO3D (saturation) are omitted for clarity.

Method	Search Overhead (min)		Best $R^{2}$
Method	Pressure	Saturation	Pressure	Saturation
All baselines (brute-force)	68.7 GPU-hours (16 models to convergence)		0.9972	0.9359
Random Search (120 trials)	47.4	93.6	0.9820	0.7593
Optuna TPE (120 trials)	45.6	50.7	0.9804	0.8915
Hyperopt TPE (120 trials)	85.7	66.1	0.9922	0.8848
AutoSurrogate	55.5	42.3	0.9972	0.9532

Table 4: Search efficiency comparison. “Search overhead” reports HPO trial time for AutoML methods; for AutoSurrogate , it additionally includes data profiling, model selection, and LLM reasoning. “Best

R^{2}

” is the test-set metric of the final deployed model after full training.

Figure 5: Search efficiency analysis. (a) Search overhead comparison: time spent on HPO trials (and, for AutoSurrogate , data profiling + LLM reasoning). For saturation, AutoSurrogate ’s LLM-guided search is faster than all AutoML methods because it focuses trials on two pre-selected architectures. (b) LLM call profile for a representative pipeline run (26 calls, 7.2 min total). Each bar represents one LLM invocation, colored by type: tool calls (

\mu{=}8.2

s), reasoning (

\mu{=}21.7

s), and narrative generation (

\mu{=}34.8

s). Background shading marks the three pipeline phases. Across five runs, total LLM inference time is 5.8–8.2 min (

{\sim}

2–5% of wall time), confirming that the framework’s “intelligence” is nearly free in computational terms.

Stage	Agent Action	Observation / Reasoning	Outcome
Data Profiling	Extract dataset statistics via profiling tools	“1000 samples, spatial grid $80\times 80\times 20$ , 31 timesteps. Saturation range $[0,0.466]$ , mean $=0.005$ (very sparse: $>99\%$ of voxels near zero). Pressure normalized ( $\mu{=}208.14$ , $\sigma{=}6.01$ ); saturation kept in native $[0,1]$ range.”	Profile stored in shared context
Model Selection	Estimate GPU memory for 8 candidate architectures; select based on task physics	“Residual skip connections suit 3D spatiotemporal mapping. ResUNet3D at base_channels=16 : ${\sim}$ 4.6M params, ${\sim}$ 609 MB GPU memory at batch 1—appropriate for 700 training samples. Supports BCE auxiliary loss for sparse front reconstruction.”	Selected: ResUNet3D
HPO Round 1	Bayesian HPO (TPE + median pruner, 15 trials, 5 epochs/trial)	Search space: base_channels $\in\{8,16,32,64\}$ , lr $\in[10^{-5},10^{-2}]$ , batch_size $\in\{1,2,4\}$ , $\lambda_{\mathrm{bce}}\in[10^{-4},0.1]$ . Best trial: bc=16 , lr $=6.27\times 10^{-4}$ , $\lambda_{\mathrm{bce}}=4.6\times 10^{-4}$ , bs=4 . 6 trials completed, 9 pruned.	Best val loss: $8.73\times 10^{-4}$
Training Round 1	Train ResUNet3D with best HPO config (30 epochs, early stopping patience 15)	Training converges; test evaluation on 200-sample hold-out set.	$R^{2}=0.9479$ (below 0.95 target)
Self-Correction: Continuation	Resume training for 10 additional epochs from last checkpoint	“ $R^{2}$ unchanged at 0.9479 after additional epochs—learning has plateaued. The model’s capacity or hyperparameter regime cannot close the remaining gap.”	Plateau confirmed; no improvement
Self-Correction: Architecture Switch	Switch to UNet3D; launch fresh HPO	“ResUNet3D has saturated. Switching to UNet3D: simpler encoder–decoder without residual blocks may benefit from a different loss landscape under aggressive HPO.”	New HPO initiated

Table 5: End-to-end pipeline trace for the saturation task. Each row corresponds to one stage of the autonomous pipeline, showing the agent’s action, its reasoning based on observed data, and the quantitative outcome. The trace demonstrates the full decision chain from raw data profiling through self-correction to target achievement.

Hyperparameter	Baseline	AutoSurrogate
Hyperparameter	Default	Pressure	Saturation
Architecture	(fixed)	ResUNet3D	UNet3D
Base channels	16	32	32
Normalization	Group	Instance	Instance
Learning rate	$10^{-4}$	$4.12\times 10^{-3}$	$7.02\times 10^{-3}$
Weight decay	$10^{-5}$	$3.0\times 10^{-4}$	$1.8\times 10^{-5}$

Table 6: Key hyperparameters discovered by AutoSurrogate compared with hand-tuned baseline defaults.

Task	Round	Strategy	Architecture	$R^{2}$	Status
Saturation (target: make $R^{2}$ higher than 0.95)
Sat.	1	Initial HPO + Train	ResUNet3D	0.9479	Below target
Sat.	1b	(i) Continuation	ResUNet3D	0.9479	Plateau
Sat.	2	(iii) Arch. switch + HPO	UNet3D	–	Training
Sat.	2b	(i) Continuation	UNet3D	–	Improving
Sat.	Final	–	UNet3D	0.9532	Target met

Table 7: Self-correction timelines for two representative AutoSurrogate pipeline runs. Strategy labels refer to Algorithm 2 : (i) continuation, (ii) stability-constrained restart, (iii) architecture switching, (iv) global-best fallback.

Panel of Figure 6 — Figure 6: CO 2 saturation predictions and absolute errors for Samples 822 and 851 at 30th year. Each sample uses two rows, and the rightmost column shows the shared colorbars.

Panel of Figure 7 — Figure 7: Pressure predictions and absolute errors for Samples 820 and 815 at 30th year. Each sample uses two rows, and the rightmost column shows the shared colorbars.

Neural Surrogate Modeling

Neural surrogate modeling uses learned models to approximate costly simulations, enabling faster prediction, design, control, and inference.

See all 410 papers

No Free Lunch in Flow Surrogates under Time-Varying Boundary Conditions: A Two-Regime Study
Jul 26, 2026Georg Winkler, Martin StollNeural Surrogate Modeling PDE Surrogate Modeling
Department of Mathematics, Chemnitz University of Technology, 09111, Chemnitz, Germany
PI-JEPA: Label-Free Surrogate Pretraining for Coupled Multiphysics Simulation via Operator-Split Latent Prediction
Apr 1, 2026Brandon Yee, Pairie KohMultiphysics Simulation PDE Surrogate Modeling
Physics Lab, Yee Collins Research Group · Graduate School of Business, Stanford University
FlowForge: A Staged Local Rollout Engine for Flow-Field Prediction
Apr 21, 2026Xiaowen Zhang, Ziming Zhou, Fengnian Zhao +1Efficient Inference PDE Surrogate Modeling
Shanghai Jiao Tong University · University of Michigan

Jul 26, 2026J/K move · Enter open · S save

No Free Lunch in Flow Surrogates under Time-Varying Boundary Conditions: A Two-Regime Study

Georg Winkler, Martin Stoll

Department of Mathematics, Chemnitz University of Technology, 09111, Chemnitz, Germany

Neural Surrogate Modeling PDE Surrogate Modeling Neural Operators Computational Fluid Dynamics Autoregressive Modeling

Open paper AskPDF arXiv

We test whether an architecture that succeeds on a simple flow regime also succeeds on a richer one, with each trained separately on each regime. We explore two transient flows under time-varying boundary conditions: the three-dimensional slurry film in chemical-mechanical planarisation (CMP), central to semiconductor manufacturing, and the two-dimensional Kármán vortex street (KVS). Eight surrogate models on one shared pipeline differ in whether they learn the full field or a latent representation, and in whether they predict in one shot or step by step. No single architecture wins both regimes. On the film, a one-shot full-field model reconstructs the cumulative wall shear stress to 2.7% relative error. On the wake, a latent autoregressive DeepONet retains 90% of the shedding power that direct and one-shot models damp to almost zero. The treatment of time decides the outcome. The self-sustained wake calls for autoregressive feedback and the boundary-driven film for a direct map. Pointwise RMSE hides the damped oscillation on the wake, compresses the sixfold lead on the film's process target, and picks the damped model under wake extrapolation. The evaluation scores five physical questions. Trained surrogates answer queries 10^3 to 10^4 times faster than the finite-element solver and pay off from the first query beyond the training set on the film and from the third on the wake. Neither the winning architecture nor its validation holds across regimes. The choice of surrogate should follow the dynamical character of the target flow, and its validation should resolve the failure modes.

Abstract

Neural Surrogate Modeling

Explore similar work

No Free Lunch in Flow Surrogates under Time-Varying Boundary Conditions: A Two-Regime Study

PI-JEPA: Label-Free Surrogate Pretraining for Coupled Multiphysics Simulation via Operator-Split Latent Prediction

FlowForge: A Staged Local Rollout Engine for Flow-Field Prediction