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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 CO2_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.
CategoryHyperparameterApplies to
OptimizationLearning rateAll
Weight decayAll
Batch sizeAll
CapacityBase channel widthAll
Architecture-specificTransformer layersCNN-Transformer
Transformer headsCNN-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 parametersValues
Caprock thickness1905 m
Caprock porosity0.08
Caprock permeability (horizontal and vertical)0.001 mD
Basement thickness500 m
Basement porosity0.09
Basement permeability (horizontal and vertical)2.3 mD
Table 2: Parameters utilized for CO 2 storage simulation.
MethodPressureSaturation
R2R^{2}RMSERel L2L_{2}R2R^{2}RMSERel L2L_{2}
UDeepONet3D0.85342.36970.01140.82490.01390.4144
CNNTransformer3D0.98610.73000.00350.90980.01000.2974
EDConvLSTM3D0.98870.65770.00320.93090.00880.2604
UFNO3D0.99110.58530.00280.91020.01000.2967
RecurrentRUNet3D0.99720.32900.00160.93590.00840.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 R2R^{2} , especially on saturation where they fall below hand-tuned baselines. UDeepONet3D and FNO3D (saturation) are omitted for clarity.
MethodSearch Overhead (min)Best R2R^{2}
PressureSaturationPressureSaturation
All baselines (brute-force)68.7 GPU-hours (16 models to convergence)0.99720.9359
Random Search (120 trials)47.493.60.98200.7593
Optuna TPE (120 trials)45.650.70.98040.8915
Hyperopt TPE (120 trials)85.766.10.99220.8848
AutoSurrogate55.542.30.99720.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 R2R^{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 ( μ=8.2\mu{=}8.2 s), reasoning ( μ=21.7\mu{=}21.7 s), and narrative generation ( μ=34.8\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.
StageAgent ActionObservation / ReasoningOutcome
Data ProfilingExtract dataset statistics via profiling tools“1000 samples, spatial grid 80×80×2080\times 80\times 20 , 31 timesteps. Saturation range [0,0.466][0,0.466] , mean =0.005=0.005 (very sparse: >99%>99\% of voxels near zero). Pressure normalized ( μ=208.14\mu{=}208.14 , σ=6.01\sigma{=}6.01 ); saturation kept in native [0,1][0,1] range.”Profile stored in shared context
Model SelectionEstimate 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 1Bayesian HPO (TPE + median pruner, 15 trials, 5 epochs/trial)Search space: base_channels ∈{8,16,32,64}\in\{8,16,32,64\} , lr ∈[10−5,10−2]\in[10^{-5},10^{-2}] , batch_size ∈{1,2,4}\in\{1,2,4\} , λbce∈[10−4,0.1]\lambda_{\mathrm{bce}}\in[10^{-4},0.1] . Best trial: bc=16 , lr =6.27×10−4=6.27\times 10^{-4} , λbce=4.6×10−4\lambda_{\mathrm{bce}}=4.6\times 10^{-4} , bs=4 . 6 trials completed, 9 pruned.Best val loss: 8.73×10−48.73\times 10^{-4}
Training Round 1Train ResUNet3D with best HPO config (30 epochs, early stopping patience 15)Training converges; test evaluation on 200-sample hold-out set.R2=0.9479R^{2}=0.9479 (below 0.95 target)
Self-Correction: ContinuationResume training for 10 additional epochs from last checkpoint“ R2R^{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 SwitchSwitch 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.
HyperparameterBaselineAutoSurrogate
DefaultPressureSaturation
Architecture(fixed)ResUNet3DUNet3D
Base channels163232
NormalizationGroupInstanceInstance
Learning rate10−410^{-4}4.12×10−34.12\times 10^{-3}7.02×10−37.02\times 10^{-3}
Weight decay10−510^{-5}3.0×10−43.0\times 10^{-4}1.8×10−51.8\times 10^{-5}
Table 6: Key hyperparameters discovered by AutoSurrogate compared with hand-tuned baseline defaults.
TaskRoundStrategyArchitectureR2R^{2}Status
Saturation (target: make R2R^{2} higher than 0.95)
Sat.1Initial HPO + TrainResUNet3D0.9479Below target
Sat.1b(i) ContinuationResUNet3D0.9479Plateau
Sat.2(iii) Arch. switch + HPOUNet3D–Training
Sat.2b(i) ContinuationUNet3D–Improving
Sat.Final–UNet3D0.9532Target 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.
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.
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.
More in this topic

Neural Surrogate Modeling

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

See all 410 papers

Explore similar work

CardsList
Jul 26, 2026

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

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.

Georg Winkler, Martin Stoll

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

Neural Surrogate ModelingPDE Surrogate ModelingNeural OperatorsComputational Fluid DynamicsAutoregressive Modeling
AskPDF
Apr 1, 2026

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

Reservoir simulation workflows face a fundamental data asymmetry: input parameter fields (geostatistical permeability realizations, porosity distributions) are free to generate in arbitrary quantities, yet existing neural operator surrogates require large corpora of expensive labeled simulation trajectories and cannot exploit this unlabeled structure. We introduce PI-JEPA (Physics-Informed Joint Embedding Predictive Architecture), a surrogate pretraining framework that trains without any completed PDE solves, using masked latent prediction on unlabeled parameter fields under per-sub-operator PDE residual regularization. The predictor bank is structurally aligned with the Lie--Trotter operator-splitting decomposition of the governing equations, dedicating a separate physics-constrained latent module to each sub-process (pressure, saturation transport, reaction), enabling fine-tuning with as few as 100 labeled simulation runs. On single-phase Darcy flow, PI-JEPA achieves 1.9×1.9\times lower error than FNO and 2.4×2.4\times lower error than DeepONet at Nℓ=100N_\ell{=}100, with 24% improvement over supervised-only training at Nℓ=500N_\ell{=}500, demonstrating that label-free surrogate pretraining substantially reduces the simulation budget required for multiphysics surrogate deployment.

Brandon Yee, Pairie Koh

Physics Lab, Yee Collins Research Group · Graduate School of Business, Stanford University

Multiphysics SimulationPDE Surrogate ModelingPDE Operator LearningSelf-Supervised Pre-TrainingJoint-Embedding Predictive ArchitecturePhysics-Based Simulation
AskPDF
Apr 21, 2026

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

Deep learning surrogates for CFD flow-field prediction often rely on large, complex models, which can be slow and fragile when data are noisy or incomplete. We introduce FlowForge, a staged local rollout engine that predicts future flow fields by compiling a locality-preserving update schedule and executing it with a shared lightweight local predictor. Rather than producing the next frame in a single global pass, FlowForge rewrites spatial sites stage by stage so that each update conditions only on bounded local context exposed by earlier stages. This compile-execute design aligns inference with short-range physical dependence, keeps latency predictable, and limits error amplification from global mixing. Across PDEBench, CFDBench, and BubbleML, FlowForge matches or improves upon strong baselines in pointwise accuracy, delivers consistently better robustness to noise and missing observations, and maintains stable multi-step rollout behavior while reducing per-step latency.

Xiaowen Zhang, Ziming Zhou, Fengnian Zhao +1

Shanghai Jiao Tong University · University of Michigan

Efficient InferencePDE Surrogate ModelingSpatiotemporal ForecastingComputational Fluid Dynamics
AskPDF