Model-Agnostic Knowledge Guided Correction for Improved Neural Surrogate Rollout
Authors: Bharat Srikishan
Daniel O'Malley
Mohamed Mehana
Nicholas Lubbers
Nikhil Muralidhar
Venue: ICLR 2025 Poster (2025)
Keywords: deep learning
knowledge guided machine learning
scientific machine learning
computational fluid dynamics
reinforcement learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a novel model-agnostic, cost-aware method which combines a neural surrogate, decision model, and simulator to significantly reduce rollout error when performing time-series PDE prediction.

Abstract: Modeling the evolution of physical systems is critical to many applications in science and engineering. As the evolution of these systems is governed by partial differential equations (PDEs), there are a number of computational simulations which resolve these systems with high accuracy. However, as these simulations incur high computational costs, they are infeasible to be employed for large-scale analysis. A popular alternative to simulators are neural network surrogates which are trained in a data-driven manner and are much more computationally efficient. However, these surrogate models suffer from high rollout error when used autoregressively, especially when confronted with training data paucity. Existing work proposes to improve surrogate rollout error by either including physical loss terms directly in the optimization of the model or incorporating computational simulators as `differentiable layers' in the neural network. Both of these approaches have their challenges, with physical loss functions suffering from slow convergence for stiff PDEs and simulator layers requiring gradients which are not always available, especially in legacy simulators. We propose the Hybrid PDE Predictor with Reinforcement Learning (HyPER) model: a model-agnostic, RL based, cost-aware model which combines a neural surrogate, RL decision model, and a physics simulator (with or without gradients) to reduce surrogate rollout error significantly. In addition to reducing in-distribution rollout error by **47%-78%**, HyPER learns an intelligent policy that is adaptable to changing physical conditions and resistant to noise corruption. Code available at https://github.com/scailab/HyPER.

Published as a conference paper at ICLR 2025
MODEL-AGNOSTIC KNOWLEDGE GUIDED CORREC- TION FOR IMPROVED NEURAL SURROGATE ROLLOUT
Bharat Srikishan'; Daniel O'Malley?, Mohamed Mehana', Nicholas Lubbers?,
Nikhil Muralidhar'
Stevens Institute of Technology, Los Alamos National Laboratory (barikish, wwwraailliiiimmma www.aaillesteeeeeeeeee @stevens edu, fomalled, mzm, www.mmmmenbbbrbbmnmn..
ABSTRACT
Modeling the evolution of physical systems is critical many applications in science and engineering. As the evolution of these systems is governed by par- tial differential equations (PDEs), there are a number of computational simula- tions which resolve these systems with high accuracy. However, as these sim- ulations incur high computational costs, they are infeasible to be employed for large-scale analysis. A popular alternative to simulators are neural network sur- rogates which are trained in data-driven manner and are much more compu- tationally efficient. However, these surrogate models suffer from high rollout error when used autoregressively. especially when confronted with training data paucity. Existing work proposes to improve surrogate rollout error by either in- cluding physical loss terms directly in the optimization of the model or incorporat- ing computational simulators 'differentiable layers' in the neural network. Both these approaches have their challenges, with physical loss functions suffering from slow convergence for stiff PDEs and simulator layers requiring gradients which are not always available, especially in legacy simulators. We propose the Hybrid PDE Predictor with Reinforcement Learning (HyPER) model: a model- agnostic, RL based, cost-aware model which combines neural surrogate, RL decision model, and a physics simulator (with or without gradients) reduce surrogate rollout error significantly. In addition to reducing in in-distribution roll- out error by 47%.78% % HyPER learns an intelligent policy that is adaptable to changing physical conditions and resistant to noise corruption. Code available at httts::/githube com hcaps/a/yaaabaPRRRR
INTRODUCTION
Scientific simulations have been the workhorse enabling novel discoveries across many scientific disciplines. However, executing fine-grained simulations a scientific process interest is a costly undertaking requiring large computational resources and long execution times. In the past decade, the advent of low-cost, efficient GPU architectures has enabled the re-emergence of a powerful func- tion approximation paradigm called deep learning (DL). These powerful DL models, with the ability to represent highly non-linear functions can be leveraged as surrogates to costly scientific simula- tions. Recently, the rapid progress of DL has greatly impacted scientific machine learning (SciML. with the development of neural surrogates in numerous application domains. Some highly-impactful applications include protein structure prediction, molecular discovery [Schauperl & Denny) (2022); Smith et al.] (2018) and domains governed by partial differential equations (PDE) Brunton & Kutz (2024); Raissi et al.) (2019): [Lu al.) (2021b). Neural surrogates have also been successfully em- ployed for modeling fluid dynamics in laminar regimes like modeling blood flow in cardiovascular systems Kissas al. (2020) and for modeling turbulent DDurassay et al. (2019) and multi-phase
flows Muralidhar et al. (202T): Raj et al. (2023): Siddani et al. (2021).
Neural Surrogates are Data Hungry. Although neural surrogates are effective modeling com- plex functions, this ability is usually conditioned upon learning from large trove of representative data. The data-hungry nature popular neural surrogates like neural operators well known in
ccoreesponding author

---

Published as a conference paper at ICLR 2025
existing work (Li et al. 2020: Lu et al. 202 Tripura et al. 2024 Howard al. 2023; Lu et al. 2022). However, many scientific applications suffer from data paucity due to the high cost of the data collection process (i.e., primarily due to high cost of scientific simulations). Hence, neural surrogates employed to model a scientific process of interest, need to address the data paucity bot- tleneck by learning effectively with low volume of training data.
Rollout Errors in Neural Surrogates. Although computational simulations have been designed for modeling various types of physical systems, those exhibiting transient dynamics are espe- cially challenging to model. Solutions to systems exhibiting transient dynamics are usually obtained by discrete-time evolution of the dynam- ics. Simulators used to model such systems are in- voked autoregressively and thereby encounter nu-
merical instability and error buildup over long es-
0.8
HyPER (ours)
Cumulative MSE
timation horizons. Such error buildup during au-
0.6-
toregressive invocation is termed rollout error. Ef-
fective techniques have been developed to reduce
0.4
rollout error of computational simulations and in-
crease their numerical stability over long rollouts. Although autoregressive rollout of neural surro- gates is also affected by rollout error, solutions to minimize this eπror buildup have not been widely investigated. Recently, (Margazoglou & Magri 2023] has inspected the stability of echo-state net- works during autoregressive rollout and List et al. (2024); Carey al.) (2024); [ippe et al. (2024) have characterized rollout errors more general neural surrogates. However, systematic solution to alleviate rollout eπror in neural surrogates for modeling transient dynamics still lacking.
0.2
20%
30%
70%
90%
100%
Trajectory Progression
Figure 1: Cumulative MSE depicting rollout error for a single trajectory of HyPER vs sur- rogate only methods. x's mark the timesteps during the trajectory where our RI policy calls the simulator.
Sewllede-Guiiee Neural Surrogates. One popular method of addressing errors due data paucity in neural surrogates is to leverage knowledge of the theoretical model governing the scien- tific process. Previous efforts have incorporated domain knowledge (as ODEs, PDEs) while training the DL surrogate to develop knowledge-
knowledge-guided guided learning pipelines (Raissi et al. 2019; Karpatne al. 2022; Rackauckas al., 20200 Gao al. 202T). Most of these approaches incorporate the PDE governing the system dynamics as soft regularizers while training the neural surrogates. A majority of such approaches exhibit slow convergence and catastrophic failures in challenging, stiff PDE conditions (Krishnapriyan et al. [202]] Wang et al. 2022).
Hybrid-Modellnng All approaches discussed thus far are so-called surrogate-only (SUG) ap- proaches. Here, the computational simulator is employed only as a means of generating data to train the neural surrogate and discarded post the training. SUG approaches employ only the pre-trained neural surrogate during inference. Although SUG provide instantaneous responses relative to com- putational simulations, they generally have limited generalization ability outside the domain of the training data. An effective complement to SUG approaches are hybrid-modeling approaches (Kurz) ctt al., 2022] Karpatne et al. 2022), that jointly resolve a query by incorporating surrogates in con- junction with computational solvers. Otherwise stated, hybrid-modeling pipelines employ a 'solver- in-the-loop' (Um et al.) 2020) approach. In addition to neural surrogates, there exist a number of classical hybrid modeling techniques which combine full order model (FOM) with reduced order model (ROM) such as proper orthogonal decomposition (Willcox & Peraire 2002), dynamic model decomposition (Kutz al. 2016), and multi-scale methods. While these methods are used accel- erate scientific simulations. they are often limited in expressivity, especially in modeling complex non-lineat dynamics. Recent hybrid models (Suh et al. 2023) generally have a static coupling be- tween the components in the model and require a static interaction/transiee the ROM, while our method proposes an adaptable and learnable interaction between the neural sur- rogate and the simulator. Our proposed method allows dynamic integration of an FOM (fine-grained simulator) with ROM (neural surrogate) using reinforcement learning.
between the FOM and
Bhowledge-Guidanee with Hybrid-Modeling knowledge-guiee guided. A majority of the recent hybrid-modeling approaches are based on directly incorporating PDE solvers as additional layers in the deep learning architecture of neural surro-
Hybrid-modeling approaches are inherently
2

---

Published as a conference paper at ICLR 2025
B
S
Surregost
Initial
CS
CS
HyPER Components
HyPER
T
Decision Function
f Neural Surrogate
Als O.d
Abs Ei 0.000
Abo E.r0
Aoe Fur
S
PDE Simulator
HyPER Roliout Errora
7-5e reduction
1-10.5
1 -Arrror &Computation
Figure 2: Overview of Hybrid PDE Predictor with RL (HyPER) with example rollout. Here we is the decision model, fo is the surrogate, and the simulator. At t = 9.0 in the above trajectory, the decision policy invokes the simulator, correcting the trajectory to reduce rollout error. The effect of this snowledge-guideds correction can be observed by reduction absolute error (dotted green circle) in the figure.
gates (Chen et al. 2018 Belbute-Peres et al. 2020 Donti et al. 2021 Pachalieva et al. 2022), While such approaches address the issues of large errors under data paucity and during rollout, inherent in SUG approaches, they impose the crippling restriction of differentiability on the com- putational solvers to be incorporated as part of the DL. pipeline. Most solvers and computational simulations are NOT differentiable out-of-the-box and hence imposing such differentiability con- straints drastically curtails the applicability of current mybrid-modeling approaches.
To address the existing challenges with surrogate-only and hybrid modeling approaches, we propose the Hybrid PDE Predictor with RL (HyPER) framework. HyPER a model-agnostic,, simulator- agnostic framework that learns invoke the costly computational simulator (in cost-aware man- ner) as knowledge sooweddge-giddd correction to alleviate the effects of rollout errors in surrogates trained with low volumes of training data. Fig. 2 depicts the proposed HyPER framework. Our contribu- tions are as follows.
HyPER is first of its kind knowledge- guided correction mechanism that incorporates simulators in the loop without the requirement of the simulators to be differentiable.
HyPER is model agnostic (i.e., functions with any neural surrogate, scientific simulator) and trained in a cost-aware manner, to intelligently invoke the simulator to correct the rollout error of the neural surrogates.
We demonstrate through rigorous experimentation on in-distribution, out-of-distribution and noisy data that HyPER significantly reduces rollout error relative to SUG approaches by comparing with state-of-the-ae thr-a neural surrogates.
2
METHOD
2.1 PDE PREDICTION
We aim to solve PDEs involving spatial dimensions x = [x1, Ø2 Im] E RT and scalar time t e [0, T". These PDEs relate solution function u ((x,t) t) 12 x [0 - R²" to its partial derivatives over the domain. We assume æn, 1) which define the field values time and at the boundaries of the domain respectively. These time-dependent PDEs can be generally defined as:
have initial conditions u(æ,
0) and boundary conditions uz(
Ou 0t
F(x u, Ou Su
Ox 0x? "......))
We focus autoregressive rollout of these PDEs over time, which can be defined with function g taking current state and time as inputs and producing the next state:
g t). At) u(x, At)

---

Published as a conference paper at ICLR 2025
Note that the equation above can be applied autoregressively any number of timesteps to unroll a PDE trajectory. The cumulative error of this autoregressive process is defined as the rollout enoor, which we aim minimize.
2.2 HYPER COMPONENTS
Surrogate ML Model For the sake of conciseness, we shorten u(x, t) to uy We begin with machine learning model fo(u,), which we denote as the surrogate. This model can be any deep learning model with parameters & that predicts next state u4++1 given starting state and time u4.
Simulator We also define a PDE simulator S(u)) which numerically solves the PDE to find the state u4+11 Crucially, this simulator is only required to return the next state without any gradient information during training and inference.
Decision Model
Finally, we have decision model dg (u,) which takes current state at time and outputs next action: either call the surrogate fo or the simulator S. In HyPER we implement our decision model as a learned policy Ty which we train using reinforcement learning (RL).
We formalize our decision model using a Markov Decision Process (MDP) which is a tuple (S, A, P. r). Our states S consist of current state and current time 6. Our action space is binary at each timestep and defined as {0- call surrogate, = call simulator). Our reward function for trajectory of length 7 is:
R(a)= = Σ - Lifo, ue, unss, as) + b(fa, - C(a,
(3)
Here L represents an error function, b is baseline function, and C is cost function. The baseline function stabilizes the reward values and leads better policy learning. The baseline function b we use is the mean squared eтror of randomly calling the simulator the same number of times as our current HyPER policy. Our error function is defined as the mean squared error of either calling the surrogate simulator according to the policy action:
C(fo, S, the 04+1, (G(u)) ux+1)²
(4)
g(u,, fo(uu) S(w) when de
when a =
(5)
T the length of the trajectory and À is hyperparameter set by the user to specify what percentage of the trajectory to call the simulator. For example, if X = 0.5 then the cost function will penalize the reward if the simulator is not called 50% of the time. Our cost function defined as the absolute difference between À and the percent of the trajectory that our policy called the simulator:
|a|| T
C(a,
(6)
Intuitively, reward function B optimizes the RL decision model to minimize mean squared error while calling the simulator for proportion of the trajectory. To learn a policy To ( (t)) we use the REINFORCE policy gradient algorithm (Sutton et al. 1999) training with the update:
VoJ(0) E
Σ
Vo log (o u.) R(a)
t=
For the full training algorithm and details of HyPER see Appendix A.I
2.3 EXPERIMENTS
2D Navier Stokes Dataset. We create 2D incompressible Navier Stokes fluid flow dataset using 4Flow (Holl & Thuerey (224)) and follow the example of Gupta Brandstetteri (2022). We use the Navier-Stokes equations in vector velocity form along with an additional scalar field representing
particle concentration:
dv
= Vo+pViv Vp
(8)
0t
4

---

Published as a conference paper at ICLR 2025
f (0,0.5c) V v =
(10)
dc = n Ve
(11)
0t
Eq. 8 comprises a convection term - Ve, diffusion term V where μ indicates kinematic
viscosity, a pressure term Vy, and external force term f. Eq. 9 defines the buoyant force which applied in the y-direction. The velocity divergence term Eq. T0 enforcess conservation of mass. Eq. II models the particle concentration field c that advected by the velocity vector field v. Note that influuenees through force f and v affects through Eq. II creating complex dynamics between fields. For our experiments, we generate 1,000 trajectories of 20 timesteps with At = 1.5s at a grid size of 64x64, and diffusion coefficient μ 0.01. We set our velocity boundary condition to 0 (Dirichlet) and our concentration boundary condition to 0c/0x = (Neumann). While we simulate all the fields above, our experiments focus predicting the particle concentration field We split our 1000 trajectories into 3 sets: 400 for surrogate training, 400 for RL training, and 200 for testing. Our RL model selects 4 actions at time, however HyPER puts no restrictions on action selection.
Subsurface Flow Dataset. To evaluate HyPER's ability to work with different PDEs and problem scales, we generate a subsurface flow dataset using the Julia-based DPFEHM (Pachalieva et al.) 2022) simulator. We use DPFEHM to generate 2D dataset of subsurface fluid flows modeled by the Richards equation:
00 - K(h)(VA V2) -T
(12)
at
Eq. 12 models fluid flow underground in an unsaturated medium. & represents the volumetric fluid content, K(h) is the unsaturated hydraulic conductivity, h is the pressure, V2 the geodetic head gradient, and T-1 is the fluid sink term. We generate 500 trajectories of 100 timesteps each with a grid size of 50x50 and timestep size of 10 seconds. Out of the 500 trajectories, 200 are used for surrogate training, 200 are used for RL training, and 100 are used for testing. We use Dirchlet boundary conditions (0) for volumetric flux on all sides of our domain but inject fluid at the top center of the domain at rate of 0.01m/s. ss. In experiments, prediet the pressure field h.
RESULTS AND DISCUSSION
In this section, investigate the effectiveness of HyPER model transient PDE systems effi ciently and with minimal rollout error compared to surrogate-only (SUG) rollout. We compare with state-of-the-ar neural surrogates including UNet and Fourier Neural Operator (FNO), MPP, and PDE-Refiner. Specifically, we investigate the effectiveness of HyPER rollouts under chang- ing PDE dynamics and in noisy data settings. Further, we also demonstrate the surrogate-agnostic and aanulator-agnostc nature of HyPER. Our experiments seek to answer the following research questions:
RQ1: How effective are HyPER rollouts compared to SUG rollouts for transient PDE systems?
RQ2: Are HyPER rollouts effective under changed physical conditions?
RQ3: Are HyPER rollouts effective under noisy data conditions?
RQ4: How crucial is the intelligent decision model for effective HyPER rollouts?
RQS: What is the error/efficiency trade-off between SUG, HyPER, and simulator-only paradigms?
3.1 RQI: HYPER ROLLOUTS VS SUG
We begin by comparing HyPER to six SUG baselines: UNet, FNO, UNet-Multistep. MPP-7S, MPP, and PDE-Ro Reff-eeieer. All baselines except UNet-Multistey given the current state of the resolved field they are trained to predict the next state. Mean squared the full resolution trajectory. Note that all these
trained in one-step-ahead manner i.c.,
error (MSE) loss is used to train all surrogates baselines except MPP-ZS are trained with the same dataset that HyPER is trained with. For full training details Appendix A. and for baseline details see Appendix A.6
The UNet model is based on the modern UNet architecture in PDEArena GGupta & Brandstetter 2022] while the FNO model is built with the neuraloperator library (Kovachki et al. 2021 Li al.
5

---

Published as a conference paper at ICLR 2025
Table 1: 2D Navier-Stokes results with best results in bold. Notice that HyPER rollout incurs significantly lower rollout (cumulative) error compared to SUG models.
2D Navier-Stokes
UNet
FNO
UNet-Multistep
MPP-7S
MPP
PDF-Refiner
HyPER
Final MSE
0.019
0.053
0.016 0.311
0.22 4.495
0.036 0.747
0.024 0.405
0.011
Cumulative MSE
0.312
0.882
0.164
Ground Truth
HyRER Prediction
HyPER Enror
Ground Truth
HyPER Prediction
HyPER Abs
1-15
655
1-- 28.56
(a)
(b)
Figure 3: Predictions and absolute error snapshots of HyPER rollout for two distinct trajectories.
2020). UNet-Multistep is trained to specialize for rollout trajectory prediction, i.e. it minimizes MSE for 20-step autoregressive prediction. MPP-ZS (zero-shot) and MPP (McCabe al.) (023) large transformer based multi-task models which have been pretrained diction tasks. MPP-ZS is the off-the-shelf pretrained model, while MPP finetuned on our data. PDE-Refiner (Lippe et al. 2024) is modern diffusion based model which both predicts the next timestep while adding multi-step noising/denoising process and has been shown to reduce rollout error HyPER utilizes UNet model that is pretrained 400 trajectories while the RL model of HyPER is trained on a separate set of 400 trajectories with À = 0.3. All SUG baselines are trained using 800 trajectories (same HyPER RL plus surrogate dataset) for 200 epochs. shows the MSE of the final trajectory state (Final MSE) well as the aggregated MSE (Cumulative MSE) over all trajectory states rollout. Both metrics are calculated as an average across 200 test trajectories for both SUG and HyPER rollouts. HyPER outperforms all baselines significantly, with an average improvement in cumulative rollout error of 68.30%. While UNet-Multistep seems to the best performing SUG method, because it trains across the full trajectory it learns smoothed prediction that has poor high frequency detail (suffering from spectral bias) in comparison to UNet (see A..4 for examples). By incorporating the simulator, HyPER yields significantly lower cumulative roll- out error while only invoking the simulator during ~ 30% of the trajectory for knowledge-guide guided correction of surrogate rollout.
variety of fluid pre-
In Figure we plot the rollout MSE performance of HyPER versus the surrogate only (SUG) meth- ods for single sample trajectory. Note that HyPER outperforms both SUG models by large margin while only invoking the simulator six times in the 20-step trajectory. Also see Figures a)) and 3(b) which depict qualitative snapshots from two distinct HyPER rollouts.
3.2 RQ2: HYPER ROLLOUTS WITH CHANGING PHYSICAL CONDITIONS
Pretrained neural surrogates are often confronted with modeling contexts with similar PDE dynam- ics but varied initial boundary conditions. Effective rollouts under changing physical conditions is hence crucial requirement for the effective use of neural surrogates in computational science. To this end, we inspect rollout errors of SUG and HyPER, under changing boundary conditions. Specif- ically, test the adaptability of HyPER, we investigate HyPER rollouts with changing boundary conditions in our Savier-Stokes experiment. We generate a separate Navier-Stoker dataset which follows the same initial conditions of previous experiment, but comprises different velocity

---

Published as a conference paper at ICLR 2025
boundary condition the top boundary of the domain. Specifically, the velocity at the top boundary of the domain is increased from 0.0 to 0.5 m/s (imposing an intermittent external forcing effect caus- ing the fluid to escape from the top of the domain) for four intermediate time-steps of the trajectory (timesteps 12-16). Following this, the boundary condition is reverted back to 0.0 for the rest of the trajectory. This results in the fluid escaping out of the top during these time steps changing the PDE dynamics significantly.
Table Results depict rollout error for 2D Navier-Stokes with changing boundary conditions. No- tice that HyPER accumulates lower rollout error compared to SUG approaches.
Changing Boundary
UNet-P
FNO-P
UNet
FNO
HyPER
Final MSE Cumulative MSE
0.362
0,412
0.014
0.027
0.027
1.84
2.538
0.92
0.969
0.527
In this case, surrogates UNet-P and FNO-P are not re-trained with the changed boundary condition and hence show poor performance accumulating significant rollout error. UNet and FNO are trained on the changing boundary data (the same dataset as HyPER), but still suffer from poor rollout emror We train our HyPER RI. policy with a pretrained surrogate (not trained on to the changed boundary setting) and a simulator (fully aware of the changed boundary condition). The RL policy of HyPER trained with 400 these changed trajectories to learn the optimal policy of (frugally) invoking the simulator apply knowledge-guided correction to reduce surrogate rollout errors. As shown in Table 2, HyPER reduces cumulative error in this scenario by 71.36 % and 79.20 relative to UNet- P and FNO-P rollout errors respectively. Even when compared surrogates trained the changing boundary data UNet and FNO, HyPER reduces rollout error by 42.72% and 45.61 % respectively. In Figure 4 demonstrate that HyPER has much lower rollout error than the UNet-P and FNO- models. Figure 4 shows sample qualitative predictions of HyPER and SUG rollouts under the changing boundary scenario of interest, to further reinforce our point. As illustrated, our model prediction is much closer to the ground truth after the boundary condition has changed and fluid has escaped the box owing to the appropriately invoked mnowledge-guided correction by the learned RL. policy in response to increasing surrogate rollout error. Fig. 4h shows the significantly lower rollout error for HyPER rollout relative UNet-P and FNO-P models.
HyPER Policy
HyFER
(iii)
(iv)
Cumulasiee
(vi)
(vii)
Trajectory Proge
(a)
(b)
Figure 4: Predictions and absolute error of HyPER vs UNet and FNO for a single trajectory and timestep. Fig. a shows the rollout error accumulation over the trajectory with 'x' marking the times the simulator is called. Fig. 4 shows the resolved system state for a the same trajectory at a single timestep. Fig. 4b(i) shows the ground truth field while Fig. (b(ii)-(iv)) indicate the result of UNet, FNO and HyPER rollouts respectively. Fig. Fig.4p(y ((vv--vvi)) depict the corresponding absolute errors. We notice that only HyPER rollouts capture the correct characteristics relative to the ground truth owing to the knowledge-guide correction while SUG models are unable to faithfully resolve the trajectory under changing physical conditions.
3.3
RQ3: HYPER ROLLOUTS WITH NOISY DATA
Neural surrogates, in addition to being data hungry and accumulating rollout error, have also been known to exhibit catastrophic failures in challenging PDE conditions (e.g., stiff PDEs) Krish napriyan et al. (202T). As neural surrogates are crucial part of HyPER rollouts, it is imperative
7

---

Published as a conference paper at ICLR 2025
Table Results depicting rollout error for the 2D-Navier-Stokes experiment, with random Gaussian noise added to inputs, at timesteps 12-16. o² is the scale of the noise. The percentage reduction of cumulative MSE rollout error by HyPER, over the best performing SUG model is in parentheses.
Cumulative MSE
Experiment
UNet-P
FNO-P
UNet
FNO
HyPER
Unimodal 8² 1.0
2.931
2.523
1.61
1.713
1.291 (19.81%)
Unimodal o² 0.75
1.746
1.916
1,013
1.291
0.8 (21.03%)
Unimodal o² 0.5
1.009
1.487
0,642
1.137
0.49 (23.68%)
Unimodal o² 0.25
0.612
1.246
0.519
0.862
0.26 (49.90%)
to investigate whether HyPER is capable of adapting to such failures by invoking the knowledge- guided correction (i.e., the simulator) minimize the propagation (and accumulation) of such local failures over the remaining trajectory rollout. Further, another crucial property to investigate the robustness of the intelligent decision mechanism in HyPER, noisy, low-quality surrogate predic- tions. To jointly investigate both goals, we consider simple (contrived) context with noisy inputs supplied the HyPER RL policy. This additive noise, injected at specific steps corrupt the surro- gate output in the trajectory, serves to mimic low-quality surrogate predictions. Hence, experiments with such noisy inputs help characterize the ability of HyPER to adapt to sudden local changes during rollout (like catastrophic surrogate failure) in addition to demonstrating its ability for robust decision-making under noisy data conditions.
To carry out this investigation, we add random Gaussian noise with mean 0 at four separate variance (o²) scales at fixed timesteps of our trajectory. The UNet-P and FNO-P surrogate models are never trained with this noisy data and therefore perform poorly when encountering it. The UNet and FNO models are trained with the noisy data and show improvement. We train the RL policy of HyPER with small set (400) of these noisy trajectories while keeping our surrogate model static (using the UNet-P model). We test two different noise corruption scenarios added to a 20 step trajectory rollout. (i) unimodal noise: a where noise is added at timesteps [12-16] and (ii) bimodal noise: a more sophisticated case where noise added at two different time windows of [2-4] and [15-16]. The unimodal noise results are presented Table 3 where we see that HyPER rollouts outper- form sfate-of-the-art SUG approaches. In this case we notice reduction in cumulative MSE by 19.81 %%.49.9% across the four noise scales. Note that HyPER achieves this significant reduction in rollout error while using the un-specialized UNet-P surrogate, demonstrating that HyPER succeeds in adapting pretrained models to noisy data.
3.4 RQ4: INVESTIGATING HYPER ROLLOUTS VS. RANDOM POLICY ROLLOUTS
To evaluate whether HyPER learns an effective RL policy, we compare to Random Policy base- line. This baseline is designed to invoke the simulator the same number of times as the HyPER RI. policy for a particular rollout, but with uniform random probability over each timestep. By com- paring the HyPER policy rollout Random Policy rollout with the same budget, we show that HyPER learns a superior performing policy in our experimental scenarios.
Table HyPER versus random policy which calls the simulator the same number of times.
Final MSE
Cumulative MSE
Cumulative Wins % HyPER
Experiment
Random Policy HyPER
Random Policy НyPER
Noise Free
0.011
0.011
0.186
0.164
57.00%
Bimodal ²² 1.0
0.121
0.07
2.443
1.706
73.00%
Bimodal x² 0.75
0.055
0.025
1.305
0.739
84,50%
Bimodal o² 0.5
0.037
0.027
0.891
0.683
67.00%
Bimodal x² 0.25
0.022
0.019
0.495
0.424
61.00%
Changing Boundary
0.149
0.027
0.865
0.527
80,50%
8

---

200 test Published as a conference paper at ICLR 2025
We summarize these results in Table 4 which also shows the percentage of trajectories over which HyPER reduces MSE compared to Random Policy (Cumulative Wins %). 'win' is characterized by a HyPER rollout that yields a lower cumulative MSE than the corresponding Random Policy roll- out. In the NNose Free' trajectories, HyPER only wins for 57% of the trajectories because the MSEs each timestep in the Noise Free case have low variance, fairly well. However, in the bimodal noise scales (with fast and large rollout error accumulation), precise invocation of the simulator for knowledge-guided correction imperative prevent error accumulation. Hence, we see HyPER owing to its intelligent (RL-based) decision policy, has higher percentage of wins 70%) and significantly lower cumulative MSE at higher noise scales. We also see the strength of HyPER's ' learned policy when considering the changing boundary condi- tion experiment which has a 80.50% win rate over the Random Policy and reduces cumulative MSE by 39.08% %. This result demonstrates the effectiveness of HyPER in realistic physical scenarios and noisy conditions.
uniform Random Policy performs
3.5 RQ5: COST VERSUS ACCURACY TRADE-OFF
101
Time (secones)
12
MSE Serand
-.0
FNO
Ulit
HyPER
Si Only
FNO UNet
HYPER
(a)
(b)
Figure 5: Figure shows the average time of PDE prediction of a full trajectory for each method. Figure 5/b) illustrates the error per unit time (lower is better) for each method. We do not show the Sim-Only case here as we assume error is effectively zero for the simulator.
natural question is how much of cost we pay in wall clock time when utilizing HyPER compared to baselines? We evaluate this by measuring the average time of each method trajectories and show results Figure3[a), While the wall clock time of SUG methods is lower than HyPER, our method has much lower rollout error, which we illustrate in Figure Here we show the error per unit time (lower is better) and we see that HyPER has lower rollout error per second by large margin in comparison to the basclines (note this plot on log scale). While calling the simulator incurs time cost, the user can specify the À parameter to adjust to their requirements allowing HyPER be flexible in varying settings.
3.6 SURROGATE-AGNOSSS AND SIMULATOR- AGNOSTIC DESIGN OF HYPER
To demonstrate the agnel-agnostic capability of HyPER, we train it on the subsurface flow task (dataset details in Sec. 2.3 S.3. comprising longer (i.e., 100 step) trajectories. In this case, we train Hy- PER using the Julia simulator DPFEHM, which simulates underground fluid flow in porous media. As we see in Table HyPER outperforms both SUG baselines by cumulative MSE per timestep
Table 5: 2D Subsurface flow experiment results for a 100 timestep trajectory.
Cumulative MSE per timestep UNet-P
Experiment
FNO-P
HyPER
Subsurface
4.345
10.938
0.271
reduction of 93.76% and 97.52%. The integration of a distinct Julia-based simulator in a scenario with larger physical scales and times, shows that HyPER surrogate and simulator agnostic and can perform well and reduce rollout error in multiple problem settings.

---

Published as a conference paper at ICLR 2025
4
RELATED WORK
SUG approaches circumvent the use of computational simulations during inference and only em- ploy simulations to generate training data. The U-Net (Ronneberger al. 2015) model a popular approach owing to its ability to capture spatial and temporal dynamics at multiple scales. This model has recently (Gupta & Brandstetter 2022) demonstrated state of the art performance on var- ious fluid dynamics benchmarks. Operator learning (Kovachki et al. (202T) approaches that learn function families of PDES rather than single PDE instances have also been investigated to be ef- fective neural surrogates. Two notable operator learning models are the deep operator network (Lu itt al. 2021a) and the Fourier neural operator (FNO) (Li al., 2020) models. Multiple investigations have been carried out employing operator learning techniques including extending them to multi- resolution (Howard et al. 2023) Lu et al. 2022) settings. Recently TTaammoo et al. (0222 have also demonstrated that FNOs yield state-of-the-ar results on benchmark tasks.
approaches like the popular Physics-informed neural network (PINN) (Raissi al. 2019 Jagtap & Karniadakis 2020; Cuomo al., 2022) have also been ef- fectively employed for improved generalization. A related paradigm of Universal Differential Equa- tions (Rackauckas et al.) (020)) utilizes data-driven surrogates estimate focused sub-components of governing equations for better process modeling. Such approaches assume end-to-end gradient based training in a inhoiies-iformeed manner and are known converge slowly and trivial solu- tions under data paucity and stiff PDE conditions Krishnapriyan et al. (202T) owing to catastrophic loss terms.
Knowledge-guided SUG
gradient imbalance Wang et al. (021)) between data-driven and physics-
SUG for Transient PDE Dynamics All SUG approaches struggle gle to model transient PDE systems an autoregressive manner and incur rollout error (Carey al. (024).. In the work of Lippe et al. 2024), it is demonstrated how the spectral bias of traditional neural surrogates leads to significant error accumulation and they propose an initial refinement solution inspired by the process in diffu- sion modeling, Separately, List al. (2024) have conducted characterization of SUG rollout error and comment about the significant improvement obtainable by incorporating simulators in-the-loop.
for Transient PDE Dynamics. Hybrid-modeling approaches retain the simula- tion and resolve each query employing the neural surrogate and the simulator 'in-the-loop', One recent example is Zhang al. (2022) which selects between simulator and surrogate (in pre- defined rule-based maner) to resolve a PDE trajectory. In their work, the simulator and surrogate each invoked a pre-fixed number of times. In contract, HyPER learns dynamic policy, capa- ble of adapting to scenarios such as changing boundary conditions. Another major drawback with many such approaches (Chen al., 2018 Belbute-Peree et al.) 2020 Um et al./ 2020 Donti al.) 2021 Pachalieva et al. (022) is the requirement of simulators to be differentiable as they are mostly employed as additional layers in the neural network architecture, to be trained end-to-end with the neural surrogates.
5 CONCLUSION AND FUTURE WORK
This work presents first of its kind knowledge-guided correction mechanism reduce rollout errors in neural surrogates that model transient PDE systems. Our proposed method Hybrid PDE Predictor with RL (HyPER) learns a reinforcement learning based cost-aware control policy to par- simoniously invoke (costly) simulation steps to correct erroncous surrogate predictions. In contrast to existing approaches that employ simulators 'in-the-loop' with neural surrogates, HyPER does not impose any differentiability restrictions on the computational simulations. Further, HyPER is sur- rogate and simulator agnostic and applicable to any neural surrogate and off-the-shelf simulator capable of resolving transient PDE systems.
We have demonstrated the effectiveness of our proposed HyPER model in traditional in in-distribution rollouts, under changing physical conditions and under noisy data conditions. Overall HyPER yields significant improvements of cumulative rollout error over sate-offtthear surrogate-only approaches with average 68.30% improvement for in-distribution rollouts, 75.28% improvement for rollouts under changing physical conditions and 28.61% improvement for rollouts under noisy data condi- tions. In the future, we will investigate more sophisticated actor-critic based RL policies based to further improve the sample efficiency of HyPER. We will also explore extensions of HyPER to more challenging problems in multi physics contexts well multi-phase flows.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
DO gratefully acknowledges support from the Department of Energy, Office of Science, Office of Basic Energy Sciences, Geoscience Research program under Award Number LANLECAI.
REFERENCES
Filipe De Avila Belbute-Peree Thomas Economon, and Zico Kolter. Combining differentiable pde solvers and graph neural networks for fluid flow prediction. In international conference on ma- chine learning, pp. 2402-2411, PMLR, 2020.
Steven L Brunton and J Nathan Kutz. Promising directions of machine learning for partial differen- tial equations. Nature Computational Science, 4(7):483-494, 2024.
N Carey, Zanisi, S Pamela, Gopakumar, Omotani, Buchanan, and J Brandstetter. Data efficiency and long term prediction capabilities for neural operator surrogate models of core and edge plasma codes. arXiv preprint arXiv:2402.08561 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through physics-informed informed neural networks: Where we are and what's 's next. Journal of Scientific Computing, 92(3):88, 2022.
Priya L Donti, David Rolnick, and J Zico Kolter. De3: A learning method for optimization with hard constraints. arXiv preprint arXiv:2104. 1225,, 2021.
Karthik Duraisamy, Gianluca laccarino, and Heng Xiao. Turbulence modeling in the age of data. Annual review of fluid mechanics, 51(1):357-377, 2019.
Han Gao, Luning Sun, and Jian-Xun Wang. Phygeonet: Physics-informe geometry-adaptive convo- lutional neural networks for solving parameterized steady-state pdes on irregular domain. Journal of Computational Physics, 428: 10079, 2021.
Jayesh K Gupta and Johannes Brandstetter. Towards muliisppaiotemporrasssaale generalized pde modeling. arXiv preprint arXiv:2209. 15616, 2022.
Philipp Holl and Nils Thuerey. &ow (PhiFlow): Differentiable simulations for pytorch, tensorflow and jax. In International Conference on Machine Learning. PMLR, 2024.
Amanda A Howard, Mauro Perego, George Em Karniadakis, and Panos Stinis. Multifidelity deep operator networks for data-driven and physics-informed problems. Journal of Computational Physics, 493:1 12462, 2023.
Ameya D Jagtap and George Em Karniadakis. Extended chysics-informne neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physics, 28(5), 2020.
Anuj Karpatne, Ramakrishnan Kannan, and Vipin Kumar. Knowledge Guided Machine Learning: Accelerating Discovery Using Scientific Knowledge and Data. CRC Press, 2022.
Georgios Kissas, Yibo Yang, Fileen Hwuang, Walter R Witschey, John A Detre, and Paris Perdikaris. Machine learning in cardiovascular flows modeling: Predicting arterial blood pressure from non- invasive 4d flow mri data using physics-informet neural networks. Computer Methods in Applied Mechanics and Engineering, 358:1 12623, 2020.
Nikola B. Kovachki, Zongyi Li. Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. CoRR, abs/2108. 008000088881.8888 2021.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physics-informe neural networks. Advances in neural infor- mation processing systems, 34:26548-26560, 2021.

---

Published as a conference paper at ICLR 2025
Stefan Kurz, Herbert De Gersem, Armin Galetzka, Andreas Klaedtke, Melvin Liebsch, Dimitrios Loukrezis, Stephan Russenschuck, and Manuel Schmidt. Hybrid modeling: towards the next level of scientific computing in engineering. Journal of Mathematics in Industry, 12(1):8, 2022.
J Nathan Kutz, Steven L Brunton, Bingni W Brunton, and Joshua L Proctor. Dynamic mode decom- position: data data-driven modeling of complex systems. SIAM, 2016.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:
2020.
Phillip Lippe. Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Bjoern List, Li-Wei Chen, Kartik Bali, and Nils Thuerey. How temporal unrolling supports neural physics simulators. arXiv preprint 149777,402.,297,,, 2024.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-22) 2021a.
Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. Deepxde: deep learning library for solving differential equations. SIAM review, 63(1):208-222 2021b.
Lu Lu, Raphaël Pestourie, Steven Johnson, and Giuseppe Romano. Multifidelity deep neural operators for efficient learning of partial differential equations with application to fast inverse design of nanoscale heat transport. Physical Review Research, 4(2):023210, 2022.
Georgios Margazoglou and Luca Magri. Stability analysis of chaotic systems from data. Nonlinear Dynamics, 111(9):8799-881 2023.
Michael McCabe, Bruno Régaldo- Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al. Multiple physics pretraining for physical surrogate models. arXiv preprint arXiv:2310.02994, 2023.
Nikhil Muralidhar, Jie Bu, Ze Cao, Neil Raj. Naren Ramakrishnan, Danesh Tafti, and Anuj Karpatne. Phyflow: Physics-guided deep learning for generating interpretable 3d flow fields. In 202/ IEEE International Conference Data Mining (ICDM), pp. 1246-1251. IEEE, 2021.
Aleksandra Pachalieva, Daniel O' O'Malley, Dylan Robert Harp, and Hari Viswanathan. Physics- informed machine learning with differentiable programming for heterogeneous underground reservoir pressure management. Scientific Reports, 12(1):187344 2022.
Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner. Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, and Alan Edelman. Universal differential equations for scien- tific machine learning. arXiv preprint orXiv:2001.04333 2020.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics- Pnysics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Neil Ashwin Raj. Danesh Tafti, and Nikhil Muralidhar. Comparison of reduced order models based on dynamic mode decomposition and deep learning for predicting chaotic flow in random ar- rangement of cylinders. Physics of Fluids, 35(7), 2023.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 241. Springer, 2015.
12

---

Published as a conference paper at ICLR 2025
Michael Schauperl and Rajiah Aldrin Denny. Ai-based protein structure prediction in drug discov- ery: Impacts and challenges. Journal of Chemical Information and Modeling, 62(13)::142-3155 2022.
Bhargav Siddani, S Balachandar, William C Moore, Yunchao Yang, and Ruogu Fang. Machine
learning for physics-informed generation of dispersed multiphase flow using generative adversar- ial networks. Theoretical and Computational Fluid Dynamics, 35:807-830, 2021.
Justin S Smith, Adrian Roitberg, and Olexandr Isayev. Transforming computational drug discov- ery with machine learning and ai, 2018.
Seung Won Suh, Seung Whan Chung, Peer-Timo Bremer, and Youngsoo Choi. Accelerating flow simulations using online dynamic mode decomposition. arXiv preprint arXiv:2311.18715, 2023.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 15:1596-1611, 2022.
Tapas Tripura, Akshay Thakur, and Souvik Chakraborty. Multi-fidelity wavelet neural operator surrogate model for time-independent and time-dependent reliability analysis. Probabilistic En- gineering Mechanics, 77:103672, 2024.
Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. holver-in-the-low Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems, 33:611 1-22 2020.
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow patholo-
gies in physics-informe neural networks. SIAM Journal on Scientific Computing, 4305)) A3081, 2021.
Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449: 10768, 2022.
Karen Willcox and Jaime Peraire. Balanced model reduction via the proper orthogonal decomposi- tion. AIAA journal, 40(1 1):2233-2330, 2002.
Enrui Zhang, Adar Kahana, Eli Turkel, Rishikesh Ranade, Jay Pathak, and George Em Karniadakis. A hybrid iterative numerical transferable solver (hints) for pdes based on deep operator network and relaxation methods. arXiv preprint 2208::2288.1 21..22222.32333 /3273, pp. 198, 2022.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 HYPER TRAINING PROCEDURE
Below we detail the RI. training procedure of HyPER.
Algorithm HyPER Training Algorithm
Require: Dataset: D, Pretrained Surrogate Model: fa. Simulator: S, Decision Model: de,
Trajectory length: т. Simulator proportion hyperparameter: X, Learning rate: ?)
Error function (MSE in our case): f, Cost function:
1: for u, 2 E do
b For every trajectory in dataset, get field u and conditional features Initialize number of simulator calls Initialize decision model rewards list
2:
k 1<-0 0
3:
Ra + []
4:
Rg []
Initialize baseline rewards list
5:
L ++
List to store log probabilities of each action p Set initial field prediction D For every time-step in trajectory > Get RL model probabilities for next action
6:
ú(x; -1)
+
7:
for 1 E do
8:
p dø(ū(x ( 1), 2,t)
9:
~ Bernoullip)
>> Sample next action b Store log probability of action
10:
-+ (p)) + (1 - a) log(1 - p)
11:
if a 0 then
12:
ú(x, t) + fo (((x, x. t 1),
> Call surrogate for next step prediction
else if a 1 then
14:
ú ar. 1) + S(ū(x, 1),()
Call simulator for next step prediction Track number of times simulator called
k: ++11
16:
end if
Ra += (( 1), u(x t)) c(>, k r)
b Store policy reward based on MSE and cost function
end for
19:
Unffom@Withotr/Rppccemme//dd
Sample l times between [0, r] without replacement, this list will contain the time- steps which the random baseline will call the simulator
20:
for / E [0, T) do do
> Run random baseline for trajectory
21:
if / then
22:
ú(æ, t) + fo(ū(x t- t),t)
> Call surrogate for next step prediction
else ift I then
24:
ú( f) + S(ū(x, - 1),t)
Call simulator for next step prediction
end if
26:
R. + (ū( 1), u(æ, 1)) clA,k,n
D- Store baseline reward using MSE and cost function
27:
end for
28: VoJ - VOL StopGradicnt Ra) Calculate policy gradient, this is done
element-wise and then summed
29:
0 - 0 - VJ
P Update RL decision model parameters
30: end for
A.2 RQ3: HYPER ROLLOUTS WITH NOISY DATA: BIMODAL NOISE EXPERIMENTS
depiets the performance of HYPER in the bimodal noise case. We notice results similar to the unimodal noise case, with a reduction in cumulative MSE error of 41.46% 1...4%%-5.4449 65.24%. We believe this improved performance in more complex noise distribution is result of our RL policy learning the more complex error distribution while the SUG methods accumulate larger error over two different noise windows.
14

---

Published as a conference paper at ICLR 2025
Table Results depicting rollout error for the 2D-Nevier-Stokes experiment with random Gaussian noise added at timesteps 2-4, 15-16. o² is the scale of the noise. The percentage reduction of cumulative MSE rollout error by HyPER, over the best performing SUG model is in parentheses.
Final MSE
Cumulative MSE
Experiment
UNet-P
FNO-P
HyPER
UNet-P
FNO-P
HyPER
Bimodal x² 1.0
0.844
0.097
0.07
6.648
2.914
1.706 (41.46%)
Bimodal o² 0.75
0.25
0.079
0.025
3.436
2.126
0.739 (65.24%)
Bimodal o² 0.5
0.087
0.07
0.027
1.942
1.576
0.683 (56.67%)
Bimodal o² 0.25
0.046
0.068
0.019
1.095
1.266
0.424 (61.28%)
Unet
FNO
HyPER (ours)
Cumulative MSE
0.6
0.4
0.2
10%
30%
40%
50%
60%
70%
80%
90%
100%
Trajectory Progression
Figure 6: Average Cumulative MSE over all test trajectories of HyPER vs. UNet and FNO.
A.3 HYPER AVERAGE PERFORMANCE
Figure 6 shows the average rollout eπrr of HyPER versus UNet and FNO methods for all 200 Havier-Stokes test trajectories. This demonstrates that HyPER is effective reducing rollout error significantly over large set of unseen test trajectories.
A.4 UNET-MULTISTEP PREDICTION QUALITY
In Figure 7 below show five different sample predictions of UNet versus UNet-Multistep. Our investigations revealed that in comparison to UNet, all UNet-Multistep predictions look "smoothed". Essentially found that while UNet-Multistep had reasonably low MSE, it failed capture high frequency detail and suffered from spectral bias. Thus, UNet was chosen as our surrogate model in HyPER across all experiments because it showed comparable cumulative MSE UNet-Multistep (Table I..
A.5 HYPER VERSUS UNET PREDICTION
We contrast the predictions and errors of HyPER rollout and UNet rollout in Figure Notice that at t = [16.5s, 21.0s). UNet suffers from spectral bias and has high error on the plume edges where high frequency detail is required. We that HyPER successfully mitigates this spectral bias error by incorporating the simulator, thereby significantly reducing rollout (cumulative) error over the full trajectory.
15

---

Published as a conference paper at ICLR 2025
Ground Truth
UNet Pred
UNet-Multistep Pred
Sample 1
Sample 2
Sample 3
Sample 4
Sample
Figure 7: Sample predictions of UNet vs UNet-Multistep on five different noise free trajectories. Note that UNlt-Mult predictions suffer from spectral bias and have poor high frequency detail.
A.6 MODEL ARCHITECTURES AND PARAMETERS
A.6.1 MODEL TRAINING DETAILS
We build and train all our neural surrogate models using the PyTorch library on a single Nvidia RTX A6000 GPU. The 4Flow simulator used in the Navier-Stokes experiment runs the same GPU. The DPFEHM Julia simulator runs using multi-threading on an Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz.
The UNet, FNO, UNet-Multistep. MPP, and PDE-Refiner baselines are trained with same number of samples as HyPER for the sake of fair comparison. The UNet-P and FNO-P baseline models are trained with a smaller dataset of 400 trajectories which are separate from the 400 trajectories HyPER's RI. policy is trained with. All baselines trained for 200 epochs.
16

---

Published as a conference paper at ICLR 2025
Ground Truth
UNet Pred
HyPER Pred
UNet Abs Err
HyPER Abs Err
t= 3.0s
t= 7.5s
t= 12.0s
t= 16.55
t= 21.0s
= 25.55
t= 30.0s
0.000.250.500.77
Figure 8: Sample trajectory predictions and absolute errors of HyPER vs UNet. At various timesteps including 21.0s, we see that UNet has high error on the edges the plume where high frequency behavior is present while HyPER successfully reduces this error.
We only train the HyPER RL ResNet model for 30 epochs to demonstrate that HyPER does not require extensive training outperform the baselines.
17

---

Published as a conference paper at ICLR 2025
Table 7: Model training times and datasets for SUG methods and HyPER.
Model
Training Time (Hours)
Epochs
Dataset Size (# Trajectories)
UNet-P
1,46
200
400
FNO-P
0.55
200
400
UNet
3.29
200
800
FNO
2.02
200.
800
UNet-Multistep
4.10
200
800
MPP
9.58
200
800
PDE-Refiner
3.66
200
800
HyPER
5.21
30
400
A.6.2 MODEL PARAMETERS
Our UNet model is built on top of the PDEArena modern UNet architecture with wide residual blocks and training parameters in Table 8
Table 8: UNet parameters.
Parameter
Value
Model Size (# parameters)
12,295,233
Hidden Channels
64
Activation Function
GELU
Channel Multipliers
[1, 2, 2]
Num Residual Blocks Per Channel
Sinusoidal Time Embedding
Yes
Learning Rate
le-4
Optimizer
Adam
Loss Function
MSE
Epochs
200
Our FNO model is built using the neuraloperator library and is constructed to have similar number of parameters as our UNet for fair comparison. See Table 9 for details.
Table 9: FNO parameters.
Parameter
Value
Model Size (# parameters)
12.437,057
Number of Fourier Modes
27
Hidden Channels
Lifting Channels
256
Projection Channels
256
Learning Rate
le-5
Optimizer
Adam
Loss Function
MSE
Epochs
200
The HyPER RL model a lightly modified version of the ResNet34 model from the Torchvision library. See Table 10 for parameters.
The Multiple Physics Pretrained model (MPP) is adapted from MPP github. Note that we train this model for 200 epochs while UNet and FNO only trained for 50. See Table TableTTi for details.
18

---

Published as a conference paper at ICLR 2025
Table 10: RL ResNet parameters.
Parameter
Value
Model Size (# parameters)
11,751,300
Number of Layers
[3, 4, 6, 3]
Hidden Channels
64
Sinusoidal. Time Embedding
Yes ReLU
Activation Function
Learning Rate
le-5
Optimizer
Adam
Reward Function
Equation
Epochs
30
Table 11: Multiple Physics Pretraining parameters.
Parameter
Value
Model Size (# parameters)
28,979,436
Patch Size
16x16
Embedding Dimension
384
Number of Axial Attention Heads
Number of Transformer Blocks Epochs
12
200
The PDE-Refiner model was adapted from PDE-Refiner to work with our UNet model. Note that train this model for 200 epochs while UNet and FNO only trained for 50. See Table 12 for details.
Table 12: PDE-Refiner parameters.
Parameter
Value
Model Size (# parameters)
12,378,241
Number of Refinement/Denoisingg Steps
3
Minimum Noise Scale
4e-7
Hidden Channels
64
Activation Function
GELU
Channel Multipliers
[1, 2, 2] 2
Num Residual Blocks Per Channel
Sinusoidal Time Embedding
Yes
Leaninn Rate
le-4
Optimizer
Adam
Loss Function
MSE
Epochs
200
19


Reviews:
[{"review_id": "t18itYYuS5", "reviewer": "Reviewer_fmMp", "summary": "In this paper the authors tackle model prediction mismatch due to rollout, by proposing a technique that merges hybrid modeling with neural surrogates. The framework differs from previous literature given that it does not use an \"only-surrogate\" approach, but it also uses on-demand data from a rigorous simulator, when it identifies that this is needed. Results were presented using a 2D Navier-Stokes problem, showing that the proposed method provides improved predictions when compared to a random approach, a purely-surrogate approach, and under challenging scenarios of noise and varying physical conditions.", "questions": "- What are some real science or engineering-based problems/case studies that suffer from rollout errors? What is their dimensionality and what time-scales are relevant for these problems with respect to decision-making (e.g., control problems where one needs to perform an action within fractions of second or other?). \n- Based on above answer, how would this method scale to real systems (if they are different to the presented 2D N-S problem)?\n- How is this work relevant or different to earlier hybrid modeling work developed in process systems engineering or for control starting in the 1990's?\n- How would only surrogate techniques predictive errors change if they were re-trained with new data from simulator (if this was not done already)? In other words, is it the hybrid modeling structure or the adaptability or both that are novel and effective in this work?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and organized. \n- The paper merges concepts from different fields (hybrid modeling + neural surrogates) in interesting ways. \n- The results that are presented are favorable for the proposed method.", "weaknesses": "- In the introduction, and in the results presented,  motivation describing real scenarios of rollout is missing. The authors mention the fact that rollout is an issue for simulations, but do not cite or list any real examples. Providing such examples and refs even in the introduction would strengthen the paper. \n- The authors present results only using a 2D Navier-Stokes benchmark problem. This further points to the previous comment on motivation. Moreover, using only this problem does not address the scalability of the proposed approach. How would this approach perform in terms of accuracy and computational cost for larger simulations with many dimensions, parameters, variables affecting predictions? If the authors cannot a larger example in supplementary, they could at least add a discussion on this.\n- There is a large body of literature in hybrid modeling (starting from the 1990's, where different structures of model correction or different fidelity of models are embedded) that is relevant to this work that is not mentioned at all in this paper. The authors should include an earlier reference and clearly describe the novelty of the proposed work compared to earlier work as well.\n- The comparisons presented with only pre-trained surrogates do not seem as fair, unless I have misunderstood the approach. The HyPER approach continuously updates the models by getting new data from high-fidelity simulation. The only-surrogate approaches do not (again, unless I misunderstood). It is thus expected that the HyPER approach would outperform all else. This can be ok, if one considers that the novelty of the HyPER framework is it's adaptive nature. However, given that when new data comes, some re-training happens, would it not be fair to allow for the surrogate-only approaches to also be re-trained with new data? It is likely that they would still perform worse, or require more training time, but such a comparison would help better explain the true novelty of the framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7wrD633Tcs", "reviewer": "Reviewer_QQgX", "summary": "This article represents an effort towards alleviating the rollout errors in neural surrogates for modeling transient dynamics. The authors assume that end-to-end training is not possible because the simulator is not differentiable. They optimize a RL policy that decides to step forward in time either with an accurate non-differentiable solver or with a neural surrogate, the resulting method is called Hybrid PDE Predictor with RL (HyPER). The reward function is a combination of an error term and a computational cost term, which limits the number of calls to the solver. Most of the article contains numerical experiments on 2D Navier Stokes and Subsurface flow applications. The experiments try to assess the benefits of the method 1) against surrogate model only approaches (UNet and FNO), 2) against change of physical conditions, 3) against noisy data, 4) against a random policy, 5) against cost/accuracy trade-offs.", "questions": "Could you please share the mathematical equations for your model (see corresponding weakness above)?\nCould you please share training details of HyPER (see corresponding weakness above)?\nCould you explain how HyPER can theoretically reduce the already accumulate error by using a simulation?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This article focuses on the hard goal of reducing rollout error.\nThe article contains a pragmatic approach to cases where the simulation capabilities may not be differentiable (due, for example, to legacy code).\nThe article lays out clearly the hypothesis that the authors want to test about HyPER.\nThe article has an interesting way to incorporate computational cost into the training objective function.", "weaknesses": "The article lacks details on the model. The fact that the approach is model-agnostic does not mean that the details to make the methodology reproducible should be omitted. \n- The model equations that would predict one rollout at described in Figure 2 are missing. \n- It is not clear why the authors are limiting their approach to one-step auto-regresssive models instead of unrolled networks for the model or the baseline. \n- Details on the computational cost of training the policy as well as its implementation are missing, which makes the results hard to reproduce. RL is known to be unstable when training, the authors should communicate the overall computational cost of training, what hyperparameters they needed to choose, how they choose them.\n- The reported results don't have any error bars.\n- Some terms are not explained, such as y in equation 6.\n- The notations are not consistent across equations (between Eq. 4 and Eq. 6 for example). \n- The illustrative 2D examples are weak baselines (see https://arxiv.org/abs/2407.07218 for more context) that are not representative of the scale of computation where such method would useful. \n- There is no discussion on the convergence with respect to the number training points, and how this would scale with more challenging 3D problems. \n- To finish, the authors' interpretation that the simulation step would correct the accumulated rollout error from the surrogate model is not substantiated. Such a statement is not supported by the equations because u(x,t) is unaltered to compute u(x, t+delta_t), so the error that is already accumulated in u(x,t) can't be reduced.\n\nNote that the number of pages of the manuscript is one page over the strict maximum of ICLR submissions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ue8QARbz18", "reviewer": "Reviewer_t7Zu", "summary": "The authors propose the Hybrid PDE Predictor with RL (HyPER) model,  which utilizes the reinforcement learning that combines a neural surrogate and a physics simulator to reduce surrogate rollout error significantly. This method is knowledge guided and model-agnostic. Here RL is used to decide incorporation of simulators in the loop. HyPER is compared to FNO and U-Net approaches in both accuracy and efficiency.", "questions": "- What is the reason for comparing a knowledge-guided method (with physics or changing BC known) to two data-driven approaches?\n\n- How is HyPER compared to Hybrid approaches and Sim-only approaches?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "The paper is well-written and easy to follow. The aspect of using RL in rollout error reduction is new.", "weaknesses": "- The motivation and necessity of using RL with action space {0 = call surrogate, 1 = call simulator} is questionable. If physics knowledge is known, why not directly use simulators for all steps? According to Fig. 6, the computational cost reduction is not that significant. There is no accuracy comparison between HyPER and Sim-Only, but it can be predicted that Sim-Only can be more accurate. I would suggest the authors include a Sim-Only baseline in their comparisons. In Table 4, when there is no noise, the Random Policy and HyPER have almost the same accuracy, which suggests the RL here  is not meaningful. \n\n- HyPER is compared to two surrogate baselines: UNet-Only and FNO-Only, and improves the performance significantly. However,  HyPER is knowledge-guided with PDE form known and invoked simulator, but the baselines do not require PDE knowledge. HyPER can perform better because of the knowledge imposed. The comparison is not fair. \n\n- When changing physics conditions, the HyPER is trained with a simulator that is “fully aware of the changed boundary condition”, but “both surrogate models (UNet and FNO) are not re-trained with the changed boundary condition”. Again, the comparison is not fair.  The improvement is from the knowledge of PDE conditions.\n\n\n- Could the authors clarify what “Error” and “Cost” specifically represent in formula (3)?\n\n\n- The paper does not sufficiently detail the parameters and training strategies employed. The SUG and S are not specified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zCi6Qywxat", "reviewer": "Reviewer_JFUH", "summary": "The authors propose the Hybrid PDE Predictor (HyPER), which invokes costly computational simulators as knowledge-guided correction to reduce the rollout prediction errors of neural network surrogates whenever required. The proposed model relies on a reinforcement learning policy to invoke the simulator in a cost-aware manner. The resulting framework reduces the rollout error on in-distribution, out-of-distribution, and noisy data and outperforms the compared neural surrogates, at least in the studies presented by the authors.", "questions": "1. l25. Define RL at its first appearance. \n2. Eq. (3). How are 'Error' and 'Cost' defined? The error is estimated with respect to which quantity?\n3. Eq. (4). The total reward R seems to be a function of the true solution field u(x,t). Since the ground truth is not available in the inference period, how the reward will be calculated is unclear to me.\n4. l199. The diffusion coefficient is taken as 0.01. Is it giving rise to laminar flow? What is the Reynolds number? It will be interesting to see the performance in high Reynolds numbers or in small diffusion coefficients. Since high rollout error generally occurs at long prediction horizons for turbulent flows more than laminar flows.\n5. In the 2D Navier Stokes example, the authors consider only 20 timesteps, which is very small when considering long-term predictions. However, in the Subsurface Flow example, the authors seem to consider 100 timesteps, which is considerable.\n6. How many time steps are used for training and how many for testing is not mentioned. If all the time steps are used during training, it defeats the purpose since, in practice, the neural surrogates can not be trained for finitely very long prediction horizons.\n7. Table 1. Since the HyPER is pre-trained on 400 samples and the intelligent RL policy is fine-tuned on another 400 samples, the compared methods should also be trained on 800 samples since 800 datasets are already available. This seems to be acknowledged by the authors in l340.\n8. Why do the fine-tuned models in Fig. 4(b) provide a higher error? Should the fine-tuned models not provide better accuracy than the pre-trained models?\n9. l315. To keep a fair comparison, like the UNet and FNO are not re-trained with the changed boundary condition, the performance of HyPER should also be tested without fine-tuning the intelligent RL policy.\n10. The authors should also mention the number of parameters of the models.\n11. Fig. 5(a). Are the time for all the methods computed on the same type of device, i.e., CPU or GPU? \n12. Alongside the time in Fig. 5(a), it will be interesting to see when the costly computational simulator is activated during inference.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper is written clearly and well-presented, and sufficient illustrations in terms of figures and tables are provided to support the claims of the authors. The hybrid concept of the rollout error correction using a simulator without needing the simulator to be differentiable is original.", "weaknesses": "1. Except for the switching mechanism, the proposed HyPER provides no significant contribution to the existing literature. \n2. While it is true that neural surrogates produce high rollout errors at long prediction horizons, many recent works [1-6] have been carried out to address this issue. However, these are not mentioned by the authors. In order to correctly acknowledge the effectiveness of the proposed framework, a comparison against some of these frameworks is necessary.\n3. The proposed framework uses a hybrid mixture of neural surrogates and costly computational simulators. However, the comparisons are performed against data-driven surrogates. Given the literature on differential physics and other hybrid methods (mentioned by the authors in the paper), it is necessary to compare the HyPER against some of the robust hybrid simulators. \n4. In addition, when the neural surrogates lose temporal correlations with the initial time steps in long prediction horizons, it may be required to perform repeated predictions using the costly computational solvers. In such cases, the cost of simulation is equivalent to directly solving computational solvers like FEM and FDM. \n\n\n[1] Fatone, Federico, Stefania Fresca, and Andrea Manzoni. \"Long-time prediction of nonlinear parametrized dynamical systems by deep learning-based reduced order models.\" arXiv preprint arXiv:2201.10215 (2022).\n[2] Wang, Sifan, and Paris Perdikaris. \"Long-time integration of parametric evolution equations with physics-informed deeponets.\" Journal of Computational Physics 475 (2023): 111855.\n[3] Zeng, Ailing, et al. \"Are transformers effective for time series forecasting?.\" Proceedings of the AAAI conference on artificial intelligence. Vol. 37. No. 9. 2023.\n[4] Navaneeth, N., and Souvik Chakraborty. \"Waveformer for modeling dynamical systems.\" Mechanical Systems and Signal Processing 211 (2024): 111253.\n[5] Lippe, Phillip, et al. \"Pde-refiner: Achieving accurate long rollouts with neural pde solvers.\" Advances in Neural Information Processing Systems 36 (2024).\n[6] Liu, Xin-Yang, et al. \"Multi-resolution partial differential equations preserved learning framework for spatiotemporal dynamics.\" Communications Physics 7.1 (2024): 31.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)