DiffusionRollout: Uncertainty-Aware Rollout Planning in Long-Horizon PDE Solving
Authors: Seungwoo Yoo
Juil Koo
Daehyeon Choi
Minhyuk Sung
Venue: Accepted by TMLR ()

Abstract: We propose DiffusionRollout, a novel selective rollout planning strategy for autoregres-
sive diffusion models, aimed at mitigating error accumulation in long-horizon predictions of
physical systems governed by partial differential equations (PDEs). Building on the recently
validated probabilistic approach to PDE solving, we further explore its ability to quantify
predictive uncertainty and demonstrate a strong correlation between prediction errors and
standard deviations computed over multiple samples—supporting their use as a proxy for
the model’s predictive confidence. Based on this observation, we introduce a mechanism that
adaptively selects step sizes during autoregressive rollouts, improving long-term prediction
reliability by reducing the compounding effect of conditioning on inaccurate prior outputs.
Extensive evaluation on long-trajectory PDE prediction benchmarks validates the effective-
ness of the proposed uncertainty measure and adaptive planning strategy, as evidenced by
lower prediction errors and longer predicted trajectories that retain a high correlation with
their ground truths.

Published in Transactions on Machine Learning Research (02/2026)
RiffusionRollout: Uncertainty-Aware
Rollout Planning in Long-Horizon PDE Solving
Seungwoo Yoo*
dreamy1554 @kaist.oc..k
School of Computing KAIST
Juil Koo* School of Compuling KAIST
63days ac..r
Daehyeon Choi*
dachyeonchhiikaisst.cc..
School of Computing
KAIST
Minhyuk Sung
mhsung@@aist.ooc oc..iis....c..kkk
School of Computing
KAIST
Reviewed on OpenReview: https: // openreview. net/ forum? id= OCzcGOzgzz
Abstract
We propose DiffusionRollont, novel selective rollout planning strategy for autoregres- sive diffusion models, aimed at mitigating error accumulation in long-horizon predictions of physical systems governed by partial differential equations (PDEs). Building on the recently validated probabilistic approach PDE solving, we further explore its ability to quantify predictive uncertainty and demonstrate a strong correlation between prediction errors and standard deviations computed over multiple samples supporting their use as proxy for the model's predictive confidence. Based on this observation, we introduce mechanism that adaptively selects step sizes during autoregressive rollouts, improving long-term prediction reliability by reducing the compounding effect of conditioning on inaccurate prior outputs. Extensive evaluation on long- cong-trajectory PDE prediction benchmarks validates the effective- ness of the proposed uncertainty measure and adaptive planning strategy, as evidenced by lower prediction errors and longer predicted trajectories that retain a high correlation with their ground truths.
1 Introduction
Solving partial differential equations (PDEs) is fundamental to scientific computing, enabling the modeling and prediction of wide range of physical systems from microscopic phenomena like pattern formation driven by chemical reactions and diffusion, to macroscopic dynamics such ocean currents that influence global climate. While traditional approaches have relied numerical solvers due to the lack of closed- form solutions, recent advances such as plysics-informed neural networks (PINNs) (Raissi et al. 2019 Krishnapriyan et al. 2021) and neural operators (Li et al. 2020 2021 have opened new possibilities for reducing computational cost and accelerating solution times. Besides, learning-based solvers enable data- driven simulations of real-world systems with unknown governing equations, as they can be trained directly from observational data without requiring prior knowledge of the underlying dynamics. The computational
Equal contribution.

---

Published in Transactions on Machine Learning Research (02/2026)
efficiency and flexibility of learning-baae PDE solvers have led to their success across diverse domains,
including fluid dynamics (T.i et al. 2023c
2024
Wu et al. 2024 and climate modeling (Kurth et al. 2023
The standard approach to training neural PDE solvers casts the problem as regression, minimizing empirical risk between model predictions and ground truth. Recent work (Huang et al. 2024 Shysheya et al. 2024
however, adopt a generative perspective, framing PDF modeling as conditional generation Seceraaionsamplling plausi- ble future trajectories conditioned on past observations. This shift has enabled neural solvers tackle more complex tasks, such as inverse problems (Huang et al. 2024] and data assimilation (Shysheya et al. 2024 that conventional regression- oegeessoon-based methods struggle to address. Moreover, even in purely regression-basee settings, the training objective, originally designed for modcling high-dimensional data distributions in dif- fusion and score-based generative models (Song & Ermon 2019 Ho et al. 2020; Song et al. 2021). naturally extends to PDE learning, where discretized solution functions are typically high- high-dimensional (Hu et al. 2025 More importantly, it mitigates spectral whhre regression losses overemphasize low-froquency components (Lippe et al. 2023) which adversely affects the long-term stability of solvers, particularly in nonlinear systems.
In this work, we further explore benefits of generative modeling for PDEs, focusing on the challenging task of long-term prediction through autoregressive rollouts, where the model recursively uses its own past predictions to generate future timesteps. Although this approach has been empirically shown to outperform direct prediction methods (Wang & Perdikaris 2023 Li et al.) 2022) which parameterize the solver to estimate the system state at a target time from initial and boundary conditions it also suffers from a key limitation: exposure bias (Arora et al., 2022), where discrepancies between the model's predictions and the training data accumulate over time, leading to compounding errors. This issue becomes particularly pronounced in long-term predictions of physical systerss described by continuous functions over space and time, where even small prediction errors can rapidly accumulate and destabilize the solution.
The core idea behind DiflusionRollout that addresses this challenge is to take advantage of the probabilistic nature of diffusion models and incorporate sampling- based based uncertainty estimates. Our analysis reveals that statistics computed from multiple samples drawn from a conditional diffusion mooeel-traine to approximate the distribution of possible future trajectories are strongly correlated with prediction errors. This makes them reliable metric for estimating the model's confidence in its predictions during rollouts. Based on this, we propose a simple yet effective training-free selective planning stralegy, where unreliable future predictions are discarded and the model advances only to timesteps where it is sufficiently confident. In experiments, we task neural PDE solvers with challenging tongttrayy predictions and demonstrate the effectiveness of the proposed framework, which sets a new taate-offthhe-rr over recent neural operators. Moreover, comparisons with the base model and the original rollout strategy- where the sliding window advances by the largest steps possible clearly show performance gains achieved without any additional training.
2 Related Work
Neural PDE Solvers and Long-Term Rollouts Neural PDE solvers offer several advantages over traditional numerical methods, including modeling of physical systems with unknown governing equations using pifoime-informed neural networks (PINNs) (Raissi et al.
2019 Krishnapriyan et al.) 2021 Cho et al.
2024] or neural operators (Li et al.] 2020 2021), prediction of system dynamics from partial or incomplete
observations (Huan et al. 2024 and solving control tasks (Hu et al.
2025
Together with the development
of network architectures specialized
surrogate learning (Li et al.
2023c
2024 2023b a Hao et al. 2023;
Xiao et al. 2024; Kitaey et al. 2020 Choromanski et al. 2021 Cao
2021
Wu et al.
2024 Wang &
Wang 2024 Alkin et al.) 2024). there is increasing interest in applying such models to achieve accurate
long-ter term rollouts, often spanning hundreds of timesteps (Lippe et al. 2023 Ruhe et al.) 2024 Shysheya
ct al. 2024). Lippe et al. (2023) adapt the iterative denoising scheme of diffusion models (Song & Ermon
2019 Ho et al. 2020 Song el al. 2021) to better capture high-frequency components in predicted solutions,
which is especially critical for nonlinear PDEs. Similarly, [Shysheya et al. (2024) investigate conditioning
strategics for diffusion models in PDE learning and introduce an autoregressive sampling scheme for long- horizon predictions. Likewise, Ruhe al. (2024) propose an alternative diffusion time scheduling method that better captures the varying uncertainty across the prediction horizon. Building on this line of work, we
2

---

Published in Transactions on Machine Learning Research (02/2026)
low predictive uncertainty
P (cccc
high predictive uncertainty
pa(:Iuc)
Rollout with s
Rollout with s =
context window of length C
oulput window of length F
state with & <
state with &
Figure 1: Adaptive rollout step size based on predictive uncertainty. The figure illustrates how predictive uncertainty é, governs the rollout step size. On the left, low sample variance leads all future states to satisfy é < allowing a full-step rollout (s = F). On the right, high variance (i.e., uncertainty) causes only the first state to pass the threshold, leading to = 1. Based on this, our diffusion- basied PDE solver adaptively adjusts the rollout step size during autoregressive prediction.
propose a simple yet effective approach to quantify predictive uncertainty and leverage it to further reduce error during long rollouts.
Exposure Bias in Autoregressive Generation. Exposure bias (Ranzato et al, 2016 Bengio et al, 2015 Arora et al. 2022) refers to the discrepancy between training and inference in autoregressive models: during training, the model is conditioned on ground-truth sequences, whereas al inference time, it must rely on its own previous predictions. This mismatch can lead to error accumulation, particularly in hori-horizon
generation tasks such as language modcling (OpenAI, 2024). visual autoregressive generation (Tian et al.
2024
Han et al. 2024), and PDE solving (Li al. 2023b 2021 Wu al. 2024 Hu et al.
2025 Lippe
et al. 2023; Shysheya et al. 2024p. To address this issue, variety of training st.rateiies have been proposed.
including adversarial training (Xu et al. 2019)) reinforcement learning (Bahdanau et al. 2017
Chen et al.
2020), scheduled sampling or teacher-forcing alternatives (Zhang et al. 2019 Leblond et al
2018 and
unlkeeihhodd-based objectives (Welleck et al. 2020). Despite these ad advances most prior work locuses on modifying the training procedure, often requiring models to be retrained from scratch. In contrast, we propose & simple and practical approach to mitigating exposure bias in long-horizon PDE solving without retraining. Our method adaptively adjusts the rollout step size based on the model's prediction uncertainty, balancing trade-off between network approximation error and inaccurate condition-induced error.
Uncertainty Quantification Using Diffusion Models. Uncertainty quantification in learning- based models has been extensively studied as a means of estimating the confidence model predictions; we re- fer the reader to the comprehensive survey (Gawlikowski et al.) (202)) and provide bricf summary below. Common approaches for estimating uncertainty include aggregating predictions from an ensemble of mod- els (Takshminarayynaan et al., 2017), training Bayesian neural networks (BNNs) (Depeweg et al. 2018 that model posterior distributions over network weights and make predictions by marginalizing over the learned prior. Similar techniques have only recently been applied to diffusion models (Berry et al. 2024 Chan et al. 2024). Berry et al. (2024) estimate epistemic uncertainty by training multiple submodule of pre-trained diffusion model and computing mutual information between outputs and weights, whereas Chan et al. (2024) use Bayesian hyper-networks to generate ensembles of model parameters bypass the need for training multiple networks. However, these approaches inherit the limitations of ensembles and Bayesian neural networks, including high training costs and the need for specialized network designs. In contrast, our method is architecture-agnosti and requires no additional training.
3 Lorg-Horizon Rollouts of PDE Solutions
3.1 Learning to Solve Time- Demee-Dependent PDEs
We focus on time-dependent partial differential equations (PDEs) defined over a temporal domain E [0, T] and d-dimensional spatial domain Ω C Rdd In general, the solution u(x, Ω [0, T) -> Re with temporal
3

---

Published in Transactions on Machine Learning Research (02/2026)
and spatial derivatives satisfies:
ou
Ōu dau
0t = Fu,
0x 0x :.. f(x,t),
(1)
where F' denotes a differential operator and f(x, 8) 2 x [0, + R is a forcing function. Typically, Eqn. is accompanied by the initial condition u(x, 0) (x(x) and the boundary condition Bfu = 0 for
(x, () 0 x [0, T.. The main goal of neural surrogate learning is to approximate the solution operator 9:4 + U for Equ. [ using a neural network Go, parameterized by 0.
Let up:r == [u(x, O)..... u(x, 7)] E ([T++))xxe denote the discretized solution trajectory, obtained by sampling u(x,4) 1) at N spatial points in the domain Q over timestepss each separaled by Al E Rzo- When T is small, the entire trajectory can be modeled from the initial condition ug using a single forward pass of the model. However, this direct entire trajectory modeling approach has been shown to fall short for large T-where T spans tens to hundreds of steps and each temporal solution slice u, is a high-dimensional vector. Instead, autoregressive rollouts have proven more effective than modeling the entire trajectory directly at large T (Lippc et al., 2023 Li ct al. 2022 Wang & Perdikaris 2023).
As such, we choose to model short segments of length T < T, where each segment is divided into a context of length C and a future portion of length F = T C, serving as the input and target output of Go. respectively. At inference time, the model Go predicts long trajoctories in an autoregressive manner, using its previous outpuls inputs for the next segment. For training, an entire trajectory Up.r is now chunked into M pairs of a context and future portion {(uc, up, }} with C, {t.,t + I,..., + - 1} and -1] I} indicating the sets of timestep indices belonging to the context and the future portions of the i-th pair. Note that we use C and F interchangeably to denote both the index sets and their corresponding cardinalities, for notational convenience. In this setting, the parameters O are optimized using these pairs by minimizing the empirical risk:
F. = f4: + C
L = EUD
jug- Go(
(2)
where D is the training set. Once trained, the model Go serves as an approximator of the underlying mapping from context to future, L.e., up R Go(uc).
3.2 Generative Models in PDE Solving
While regression models trained with Eqn. 2 are commonly used in PDE modeling, recent studies have demonstrated the effectiveness of probabilistic approaches to learning Go. These methods not only enable applications where generative capabilities are essential inverse problems) (Huang et al. 2024 Shysheya
et al. 2024), but also improve performance in purely predictive tasks (Lippe et al. 2023 Hu et al. 2025). From the perspective of generative modeling, particularly with diffusion or score- based models (Ho et al. 2020 Song et al.) )021), the mapping uy- = (C) can be viewed as sampling from a conditional distribution. In this framework, a random sample 6 ~ A I) is gradually transformed into the target up the future portion of a segment that is plausible given the context ug:
Up ~ (o(uu(uc)
(3)
where po (uy (uc) is a conditional distribution of possible ug's. Given sufficient past context uc, pe Po(uplue (uc) collapses to a narrow region encompassing set of plausible future states ur with only minor variations. These can be sampled through an iterative denoising or refinement process, where uF is progressively re- constructed by recovering both low- and high-frequency components while removing noises of decreasing magnitudes from the initial sample e. Compared to neural solvers trained with regression-basee losses, such as the one illustrated in Eqn. 2, this approach has been shown to better capture a broader range of fre- quencies in the target signal ur a key factor in stabilizing long rollouts for nonlinear systems (Lippe el al., 2023). In the following section, we investigate how this probabilistic interpretation naturally enables the quantification of the solver's prediction uncertainty or confidence. We then propose a planning strategy that improves prediction accuracy through a minimal modification to the rollout procedure, without requiring specialized training strategies or changes to the model architecture.
4

---

Published in Transactions on Machine Learning Research (02/2026)
Sumpls Fited Line (j-4) 34)
Sanpkis
Sarpees Fited (642:44)
Fibud Life (2-0.13)
(po3)
0-8
0.0
0.4
0.4
04
6.
6.2
0.2
03
0.2
0.5
2.0
0.0
0.2
0.5
1.0
00
6.4
0.8
0.2
644
0.0
98
10
Uncertargy (6)
Uncertaryy (6)
Uncerta
Uncotarrty (k)
Figure 2: Analysis of predictive uncertainty é; as a surrogate for prediction error. The pro- posed quantity cxhibits a strong correlation with the true prediction error across multiple PDE benchmarks, including (from left to right) Gray-Scolt |Ohana et al. 2024), Turbulent Flow (Ohana el al. 2024), Cahn- Hilliard (Soares et al. 2023 and Anisotropic Diffusion (Keehler et al. 2024h. Best viewed when zoomed in.
3.3 Adaptive Rollout Step Size via Uncertainty Estimation
When adopting the C-in, F-out autoregressive scheme discussed in Sec. 3.1 addressing crror accumulation is critical for ensuring long-horizon stability. We identify two primary sources of such accumulation: (I) network approximation error, which arises from the inherent limitations of the learned model Go in approximating the truc solution operator and tends to increase with the number of network function evalua- tions (NFEs); and (2) condition-induced error, which stems from using inaccurate or uncertain predicted states as context in subsequent steps. These two sources are closely interrelated: a prediction error caused by approximation inaccuracies is likely to propagate as a cnndition-inducee error in subsequent steps, com- pounding over time and destabilizing the entire trajectory (Ranzato et al. 2016 Bengio al. 2015 Arora
et al. 2022
Empirically, observe that across the model output uf, the prediction error tends to increase with temporal distance from the context ug. We attribute this trend not only to the inherent difficulty of modcling long- range dependencies, but also to the fact. that lemporally adjacent frames are more frequently included together in training segments, following the training procedure described in Sec. 3.1 This observation implies that F predictions within the same time window covered by Go should not be trusted equally, and that overall model performance can be improved by selectively using only the more reliable predictions as context for subsequent steps.
Specifically, let E (1.2...., denote the step size for selecting predictions from up. A straightforward strategy to mitigate condition-induced error is to follow the most conservative policy by setting s = 1, accepting only single prediction from up, discarding the remaining F-1 predictions, and advancing the time window by just one step. While this reduces reliance on potentially inaccurate predictions, it leads to a higher number of network function evaluations (NFFs), thereby increasing the likelihood of accumulating network approximation error. On the other hand, using the largest step size F' accepts all F predictions and uses the last C of them as context for the next rollout. This reduces the number of NFEs and, consequently, the approximation error. However, it increases reliance on less accurate predictions, which amplifies condition- induced error.
Therefore, an ideal way to suppress error accumulation is to identify which type of error network approx- imation error or iudition-induuced error- is more dominant and adapt the rollout step size accordingly. Inspired by adaptive step size mothods in numerical analysis (Runge) 1895 Kutta, 1901| Hairer et al. 1987) we propose a simple yet effective adaptive rollout step size strategy.
However, the true prediction error required to determine the rollout step size is inaccessible at inference time. While conventional adaptive step size methods address this limitation by estimating error through high-order approximations, such estimates are often computationally expensive and sensitive to numerical instability. To this end, we propose an surrogate measure that estimates model confidence and enables the rejection of unreliable outputs during autoregressive rollouts.
Compared to deterministic regression models, key advantage of modeling a PDE solver as generative model is its ability to estimate uncertainty through sampling. Predictions for uncertain future timesteps
5

---

Published in Transactions on Machine Learning Research (02/2026)
tend to vary across different samples drawn from po(up) uc), providing natural measure of uncertainty. We leverage this by computing statistical measures, specifically, standard deviation, over these samples as a surrogate for the model's predictive uncertainty. Since the underlying PDE dynamics in Eqn. ! considered in our work are deterministic and the spatial discretization is sufficiently fine, the dominant source of predictive uncertainly is epistemic, arising from model approximation and limited context. Our sample-based stan- dard deviation therefore provides a natural mcasure of this epistemic uncertainty in the learned generative surrogate.
Formally, for each E (1,2,.... F), the predicted state u, RN refers to the t-th slice of the future states ur E PxXxe. The uncertainty é, of u, is defined as the norm of the clement-wise standard deviation computed from set of samples { 1k=1 distribution pe(: uc):
where K is the number of samples drawn from the conditional
= Jk=1) 11
(4)
where of ( 1x=1) R. denotes the clement- wise standard deviation of the K samples, and || ll2 the Euclidean norm.
Across all four benchmark PDEs used in our experiments in Sec. the proposed uncertainty measure é, shows strong correlation (p) with the actual prediction error, as illustrated in Fig. 2 logether with the corresponding fitted lines. This supports the use of uncertainty as surrogate for prediction error. Leveraging this, at each prediction step, we compare the uncertainty of each predicted state to a prodefined tolerance т, and accept only the states whose uncertainty falls below т for use as context in the subsequent rollout. The rollout step size is then selected as
= max((t E (1,2,... F)jl < (1)).
(5)
4 Experiments
4.1
Experiment Setup
We evaluate the effectiveness of our method across various PDE types through both qualitative and quanti- tative comparisons against state-of-the-art neural PDE solvers.
Baselines. We compare DiffusionRollout against state-of-the-art neural operators, including regression- based methods-FNO (Li et al., 2021p, FactFormer (Li et al.) 2023b), and Transolver (Wu et al.) 2024
well as the diffusion- based neural operators PDE-Refiner (Lippe et al 2023 and WDNO (Hu et al. 2025 For fair comparison, all models are trained to predict six future states from four past observations, with parameter counts standardized to approximately 3 million, except for WDNO (Hu et al. 2025), which uses 21 million parameters due to convergence issues. Furthermore, for PDE-Refiner (Lippe et al.) 2023), which uses constant rollout step sizes (s = 1,...,00 6) and serves as our baseline, we adopt the same network architecture and weights as ours. This ensures a fair comparison and allows us isolate and evaluate the effectiveness of the proposed adaptive stepping technique described in Sec. 3.3 which selectively retains predicted future states during inference.
Datasets. We benchmark challenging long-trajectory simulations of length T R 100, where autoregres- sive rollout is crucial for inference. Unless otherwise noted, all simulations are conducted on a unit square discretized at 64 x 64 resolution. Our evaluation spans diverse PDE systems- gystens-ranging ranging from reaction- diffusion-driven pattern formations to turbulent flow, as detailed below. Across all PDE systems, we assume the underlying governing equations are unknown and incorporate no prior knowledge during training or inference.
1. Gray-Scott Reaction-Diffusion Equation: Gray-Scott Reaction-Diffusion models the evolution of the concentrations t(x,t) and ((x, of two chemical species, undergoing simultaneous reaction and

---

Published in Transactions on Machine Learning Research (02/2026)
diffusion:
ax δt - Óx - XY + 1(-111))))))
0Y at - br - - (J+
where 6x and dy are the diffuusio coelficients, and ʃ and k denole the feed and kill rales, respectively. The concentrations X(x, t) and Yxx, t), evolving over space and time, produce complex patterns that are highly sensitive to the parameters and k. This sensitivity makes the problem an ideal testbed, as diverse and intricate trajectories can be generated by varying these values. We use 1,200 trajectories from The Well dataset (Ohana et al. 2024), following the provided train-test split and subsampling every 10 frames obtain sequences of length 101. At test time, all models are tasked with predicting 97 future slates based observations from the first four timesteps of each simulation.
2. 2D Turbulent Radiative Flow Equation: In this setup, we consider a turbulent flow simulation as a benchmark problem, which models turbulent interaction between two gases one cold and one hot hot-initially moving relative to each other. As they gradually mix over time, they form a highly reactive internediate--tmmeenuuue gas. This behavior of the system is modeled by the equation:
Op V (pv)= 0t
8pu
=0,
+ V.
(pvr + P) = 0,
Øt
QE
E Looa
P E
+ V ((E P)b)
Ōt
7 7-1"
where y 3. with p denoting the density, the velocity, P the pressure, E the total energy, and tcool the cooling time. Due to the highly stochastic nature of gas mixing, accurately predicting future states during rollouts becomes challenging, making this setup suitable for our experiments. As with the Gray-Scott reaction-diffusion system, we use 90 trajectories of length 101 from the The Well dataset (Ohana et al., 2024 Fielding et al.) 2020), following the same training and test splits. All models use the first seven timesteps context and are tasked with predicting the remaining 94 future states.
3. Cahn-Hilliard Equation: Cahn Hilliard equation models phase separation in binary mixtures, cap- turing how two components as oil and gradualra separate into distinct regions over time. Letting c(x,t) t) denote the concentration of one component, the equation is given by:
dc σt = M (-M---+++''
where denotes the mobility, and f(c) = Wc(1-))² is the bulk free energy density, with representing the height of the thermodynamic barrier. The trajectories exhibit continuously evolving component concentrations, with dynamic patterns that emerge and fade over time. Using numerical solver (Soares) et al. 2023). we generate 960 training trajectories and 120 test trajectories, each consisting of simulation results over 101 timesteps. During inference, models are given the first seven timesteps context and tasked with predicting the remaining sequence.
4. Anisotropic Diffusion: The final problem setup we consider anisotropic diffusion, a generalized form of the heat equation in which the diffusion coefficient varies across the domain:
Ou V AVu,
(6)
Øt
where u denotes the evolving scalar field (c.g., temperature or intensity), and A is spatially varying, typically posilive-definite diffusion tensor that controls the local direction and rate of diffusion. To generate data, we use the anisotropic diffusion simulator from APEBench (Koehler et al. 2024),, producing total of 1080 trajectories, which are split into 960 for training and 120 for testing. All simulations are run with the default parameters provided by the benchmark. At inference time, the models predict the remaining 97 frames using the first four frames context.
7

---

Published in Transactions on Machine Learning Research (02/2026)
Evaluation Metrics. Given a ground-truth solution á and model's corresponding prediction u, we use the relative L2 error to evaluate the accuracy of model predictions:
||u 14--j. áll2
Rel C2(u,
ú)
(7)
|
2023 to evaluate long-term stability, we additionally Furthermore, following PDE-Refiner (Lippe et al.
compute the per-timestep Pearson correlation coefficient r(t) between the predicted and ground truth states at each timestep. Specifically, for each time index t and we compute Pearson correlation r(t) between the predicted state and ground-truth state ûg as:
r(1) = Cov(u,, ú,)
(8)
/Var(u,)arr(i..
where u, ú gN are flattened into vectors prior to computing the correlation. After computing per- timestep Pearson correlations, we quantify long-term stability by measuring the average normalized time up to which the correlation stays above 0.9, denoted T>0.99
T>0.99 T max {t E [1,2] r(t) 2 0.9}.
(9)
This metric reflects the average proportion of the trajectory during which the prediction remains highly corrleated with the ground-truth PDE trajectory.
Implementation Details. We implement our diffusion-based neural solver, including PDE-Refiner (Lippe ct al. 2023), using the DiT (Peebles & Xie 2023) architecture equipped with 3D self-attention layers and rotary embeddings to jointly capture temporal and spatial dependencies. Rather than designing a separate module to process the conditioning input ug, we concatenate the full sequence [uc]|ur]] and feed it directly into the 3D self-attention layers, together with an indicator specifying which frames correspond to condi- tioning inputs. To enable this, we adopt the training objective from Diffusion Foreing (Chen al.j 2024 Song et al. 2025), which assigns an independent noise level to each frame. This simplifies the handling of conditioning frames by assigning them noise level of zero, while applying nonzero noise levels to all remaining frames. At inference time, we employ DDIM (Song et. al.] 2022) sampling with 10 steps. For all regression-based neural operators, including FNO Li et al. (2021). FactFormer Li et al. (2023b). and Transolver Wu et al. 22024) we use the implementations provided by Wu ct al. (2024).
To compule the standard deviation for uncertainly quantification, we draw two samples (K= = 2), which can be obtained efficiently via batched inference, incurring no significant computational overhead. We use uncertainty thresholds of 1, 1, 5, and 0.32 for the Gray-Scott Reaction-Diffusion Equation, the Cahn-Hilliard Equation, the 2D Turbulent Radiative Flow Equation, and Anisotropic Diffusion, respectively.
4.2 Comparisons
We begin by summarizing the performance of our method and the baselines, measured in terms of relative
L2 error and long-term stability T>0.9, in Tab. 1 Across all three long-trajectory benchmarks introduced
in Sec. 4.1
DiffusionRollout outperforms -tate-of-the-art neural operators based on regression losses (Li
let al. 2021 2023b Wu et al. 2024 Hu et al.) 2025). including models based on Fourier layers [Li et al.
2021 and transformer architectures tailored for operator learning (L. et al. 2023b Wu et al. 2024). Tn
addition, DiffusionRollout outperforms recent diffusion-basco ncural operators (Lippe ct al. 2023 Hu ct al.,
2025), including PDE-Refiner (Lippe al. 2023), which uses default autoregressive sampling with fixed step sizes, and WDNO (Hu et al. 2025), which leverages wavelet representations to better capture abrupt changes in solutions. We observe performance drop in the case of WDNO (Hu et al. 2025), which was originally designed to predict entire short-length trajectories at once. This is due to the accumulation of prediction errors, amplified by repeated encoding and decoding through wavelet transforms during long-horize rollouts.
While PDE-Refiner (Lippe et al. 2023), used as our base model and employing naive autoregressive sampling with fixed step sizes, achieves performance comparable or better than other baselines, further improve- ments are possible by adopting adaptive sampling technique described in Sec. 3.3 In particular, we
8

---

Published in Transactions on Machine Learning Research (02/2026)
Gray-Soott Ohana et al. 2024
Turbulent Flow
Cabn-Hillian Soares et al. 2023
Anisotropie Diffusion Koehler et al. 2024
Ohana et al. 2024
Rel Co
7
Rel Co 4 7
Rel Co 4 T
Rel Lo T 7
Regression- Based Neural Solvers
FNO [Li
0.652
0.419
0.450
0.256
0.035
0.998
PactPormer La at.
202/h
0.620
0.187
0.344
0.744 0.773
0.374
0.364
0.250
0.632
Transolver
ot
2024
0.430
0.264
0.336
0.388
0.309
0.405
0.447
Diffusion- Based Neural Solvers
WDNO
Hu et aL
2025
0.792
0.022
0.843
0.000
0.567
0.066
1.152
0.000
PDE-Rehner
Lippe
2023
PDE-Refiner
Lippe o5 al
2124
2
0.406
0.307
0.321
0.800
0.275
0.543
0.040
0.993
PDE-Refluer
Lippe 25 41
2025
3)
0.116
0.303
0.322
0.791
0,267
0.358
0,033
0,991
DDE-Reeiner
2025
0.433
0,293
0.322
0.833
0.275
0.541
0,031
0,991
EDE-Reflnee
Lippe et 41
2023
0.443
0.285
0.322
0.919
0,280
0.531
0,027
0.995
PDE-Reflue Reflue
Lippe
0.462
DiffusionRRolout (Ours)
0.391
0.311
0.307
0.926
0.254
0.626
0.026
0.996
Table 1: Quantitative comparison of long rollouts. For each dataset, we report the relative L2 error and the average duration of high correlation (T>0.9). DiffusionRollout consistently achieves the lowest error and competitive or superior long-term stability, outperforming state-of-the-art neural solvers and naive autoregressive sampling strategy across various rollout step sizes. Bold indicates the best for each column.
SrottScott
Turbulent Flow
Malnn-Hlliaed
OOhana 2024
Ohana 2024
Soares al.
Regression Loss Training Derivative Threshold 4: Threshold
U.479
0.849
0,270
0.395
0.320
0.256
0.391
0.307
0.254
Table 2: Quantitative results from the ablation study. The proposed method achieves the lowest error, outprroominng both its regression-trained variant and the version using an alternative surrogate for predictive uncertainty. Bold indicates the best for each column.
observe notable error reduction on the Gray-Scott dataset (Tab. 1 first column), where our adaptive planning results in 15.37% error reduction compared to the base model's fixed-step sampling strategy, which advances steps per iteration. We highlight that this improvement is achieved purely by modifying the inference procodure based on the analysis in Sec. 3.3 with no changes to the training process.
The trend observed in Tab. T is further supported by the qualitative results in Fig. 3 which illustrate the effectiveness the proposed method capturing the long-term evolution dynamical systems. The first column of the figure shows the field at an early timestep, while the second column displays the target field much later timestep, typically separated by several dozen simulation steps. Despite the apparent differences between the initial and target states, our model generates accurate predictions. In contrast, even state-of- the-art neural solvers such as Transolver (Wu et al.) 2024) struggle in this setting, producing noticeably datasets (fifth and sixth columns) as
different patterns-esseciall especially on the Gray-Scott and reflected by the brighter eтror maps.
4.3 Ablation Study
We validate our design choices by comparing our method against two variants: (1) the same backbone network trained using the regression loss in Eqn. 2 ("Regression Loss Training"), and (2) an alternative adaptive sampling strategy that uses the temporal derivative computed using the predicted future segment as a proxy for predictive uncertainty ("Derivative Threshold"). Qualitative and quantitative comparisons are summarized in Fig. 4 and Tab. respectively. As shown in the first. row of Tab. 2 replacing the diffusion- based training objective with the regression loss in Eqn. results in performance degradation. This is especially pronounced on the Gray-Scott dataset, where the error increases by 8.8%p, highlighting the importance of using a generative training loss. Furthermore, we explored an alternative measure of predictive uncertainty motivated by the observation that systems undergoing abrupt changes inherently

---

Published in Transactions on Machine Learning Research (02/2026)
Inpur
Thrget
PNO
PaotPormer
Trausolver Wi. et al
WDNO
PDE-Relinc
Diffusion Rollout
ul.
2021
Li et ul.
2023b
2024
BE al
2025
Lippe et
2025
Gray-Scott (Ohana et ai
2024
Turbulent Flow fOhans al
2024
Cahn- Hilliard Soares 2023
Anisotropic Diffusion Roehler et al. 2024
Figure Qualitative comparison. Per-pixel error maps (columns 8) are visualized for each model's prediction of the target (column 2), given the initial context (column 1). Darker regions in the error maps indicate lower prediction error.
difficult to forecast. Specifically, consider the temporal derivative of the predicted solution, approximated using the first-order finite difference between consecutive predicted states in up:
[1 11.1
(10)
As shown in Tab. 2 computing uncertainty as the standard deviation over multiple samples proves more Sroctive-particnll in the challenging task of predicting turbulent flows.
4.4 Additional Analyses
In this section, we present additional analyses, including the impact of various hyperparameter choices on model performance, runtime comparison against. PDF-Refiner (Lippe et al. 2023). and investigation
10

---

Published in Transactions on Machine Learning Research (02/2026)
Inpot
Target
Regreasion
Derivative Thres.
6. Thres. (Oure)
Input
Target
Regression Loss
Derivative thres.
é, Thres. (Ous))
Loes
Gray-Seott Ohana et al.
2024
Cahn-Hillladd Soares et al.
2023
Turbulent
(Ohana
Figure Qualitative comparison from the ablation study. Per-pixel error maps (columns 3-5 and 8-10) are shown for each variant of our method, visualizing the prediction error (columns and 7), given the initial context (columns 1 and 6). Darker regions in the error maps indicate lower prediction error.
of the impact of varying rollout step sizes in regression-based baselines (Li et al., 2021 2023b Wu el al. 2024). Unless otherwise noted, the following analyses are conducted on the test set of the 2D Gray Scott dataset (Ohana et al., 2024). We begin by examining the time-accuracy tradc-off induced by the choice of uncertainty threshold We then demonstrate that the proposed uncertainty estimates, based on standard deviations, remain reliable even with small number of samples K, and we analyze how the number of diffusion timesteps S affects model accuracy. Finally, we show that, assuming all samples fit into GPU memory, the proposed method incurs only minimal additional overhead compared to PDE-Refincr (Lippe
et al. 2023
Impact of Varying т. We first explore threshold values T E {0.001, 0.01, 0,6,,,,,, other than the default value т = 1 used for the dataset. We summarize the relative L2 errors corresponding to the step sizes selected by our method for each threshold Fig. 5 For reference, we also plot the step sizes and predictive errors of PDE-Reliner evaluated at different fixed slep sizes. As shown, even with the same model, performance varies with the choice of threshold, highlighting the importance of selecting an appropriate value, Specifically, smaller thresholds make the model more conservative, resulting in smaller steps reflected by the average step size Savg along the trajectories.
Impact of Varying K. Beyond the default number of samples K 2 used throughout the experiments of Sec. 4.2 additionally evaluate K (3,4,5) to verify how many samples are required for reliable uncertainty cstimation. The effect of varying K prodiction accuracy is summarized in Tab. 3 As shown, the relative L2 error remains consistent across all values of K. In particular, K = 2 and K = 5 yield nearly identical results, suggesting that only few samples is sufficient for reliable uncertainty estimates.
Impact of Varying Diffusion Timesteps. In our experiments, we employ the DDIM sampler (Song et al. 2022 with 5 10 timesteps to generate trajectories. To further examine the effect of the number of
11

---

Published in Transactions on Machine Learning Research (02/2026)
0.46
Ours
PDE-Refiner Ours (Best)
0.45
0.44
Relative
0.43
0.42
0.41
0.40
0.6
0.39
1.5
2 Step Size
2.53
4
6
Figure 5: Relative L2 errors of predicted solutions versus the average step sizes Savg under different uncertainty thresh- olds (shown as overlaid text). For comparison, the relative L2 errors of PDE-Refiner with constant step sizes are also visual- ized. As shown, selecting an appropriate threshold crucial for balancing prediction accuracy and inference time, as reflected by both the error and the average step size.
Figure 6: Histograms of step sizes selected along trajectories. Re- sults are shown for the Gray-Scott equation with uncertainly thresholds
t (top) and T (middle), and for the Cahn Hilliard equation (bot- tom).
sampling steps on model performance, we also evaluate the model with larger values $ E (20, 30, 40, 50). As shown in Tab. increasing the number of diffusion timesteps and, consequently, NFEs does not necessarily improve accuracy; instead, it slightly increases prediction error, aligning with our discussion in Sec. 3.3
Runtime Analysis. Since our method builds directly on PDE-Refiner (Tippe et al.; 2023). its forward- pass computational cost is identical to that of PDE-Refiner, provided that all particles used for uncortainty estimation into GPU memory and can be processed within single batch. Accordingly, we measure both the network forward-pass time for PDE-Refiner (Lippe et al. (2023) and our method, as well as the additional time required to adaptively determine the step size. We report the network forward-pass time and the adaptive step-size selection time based our uncertainty estimates in Tab. Compared to PDE- Refiner (Lippe et al., 2023) with 1, whose step size is most similar to that selected by our method, the overall runtime remains comparable, while our approach achieves higher predictive accuracy as shown in Tab. [ As indicated in the "Adaptive Step Selection (s)" column, computing standard deviations across samples does not introduce a significant. bottleneck.
Selected Step Size Analysis. To better understand the behavior of our framework, we track the step size selected at each timestep for every test-set trajectory of Gray-Scott and Cahn-Hillian equations, and visualize the distribution of these values as a series of histograms. Specifically, we collect the step sizes selected in the Gray-Scott experiments with 7 (best) and - examine how the selected step sizes vary with the threshold and across datasets. As shown in Fig, [6] the selected step sizes vary noticeably with the choice of and across datasets. On the Gray Scott dataset with T= (Fig. 0 top), consistently taking smaller steps leads to improved performance, as also reflected in the PDE-Refimer Lippe et al., 2023) baselines, and our method automatically discovers this strategy. With the increased uncertainty threshold of - (Fig. 6 middle), our method initially behaves conservatively and then gradually increases the step size as the dynamics stabilize. On the other hand, for the Cahn-Hilliard equation, where larger rollout steps are beneficial (as demonstrated by PDE-Refine (Lippe
well as in the Cahn-Hilliard experiment, to
2023) with
5), the proposed framework successfully adapts and selects larger steps accordingly (Fig. bottom).
12

---

Published in Transactions on Machine Learning Research (02/2026)
K
2
3
10
20
30
40
50
Rel Co
0.391
0.396
0.395
0.394
Rel
0.391
0.399
0.401
0.401
0.402
Table 3: Relative L2 errors predicted solu- tions for varying sample counts K. Reliable uncertainty estimates can be obtained with as few two samples.
Table 4: Relative L2 errors of predicted solu- tions for different diffusion timesteps. Increas- ing the diffusion timesteps during generation, and thus the NFEs, results in higher errors.
Method
Network Forward (s)
Adaptive Step Selection
Total (s)
PDE-Refiner
Lippe
2023
554.623
554.623
PDE-Refiner
Lippe cb al
2123
2)
289.360
289.360
PDE-Refiner
Lippe 20 a
2025
202-421
202.421
DDF-Refiner
Lappe
154.558
154.558
PDE-Refiner
Lippe et a1
20/3
126.524
126.524
PDE-Refiner
Lippe
2026
105.005
DiflusionHollout (Uuis)
503.251
4.250
507.501
Table 5: Runtime breakdown of PDE-Refiner and DiffusionRollout. The proposed mothod incurs
only minimal computational overhead for adaptive step selection. In the experiment on Gray-Scott dataset, where method uses average step size of Savg 1.12, the overall runtime remains comparable to that of PDE-Refiner with similar fixed step size (s 1).
Gray-Scott
Turbulent Flow
Cahn-Hilliard
1.0
1.2
X
1.0
Dil- Regresinnn
0.9
0.8
1.0
INO
0.8
FuciFormer
0.6
0.7
Translver
0.6
0.6
0.4
0.4
0.5
0.4
0.2
0.2
5
Figure 7: Relative L2 errors of regression-based models with varying autorcgressive step sizes. The predictive error generally increases as the step size decreases, indicating that network approximation error is primary source of error accumulation.
Varying Step Sizes with Regression-Based Neural Operators. In addition to our main experiments, which use fixed step sizes for regression-basse based neural operators (Li et al. 2021 2023b Wu et al. 2024 we further examine the effect of varying step sizes in these methods. As shown in the trends in Fig. 7 the relative error generally increases as the step size decreascs. We attribute this behavior to regression-basct models being more susceptible to network approximation error, where a larger number of NFEs leads to increased error accumulation over long-horizs rollouts. Notably, the trajectories predicted by FNO (Li et al. 2021) diverge and produce invalid numerical values when smaller step sizes are usod, as indicated by x markers the plots. Nevertheless, as discussed in Sec. 4.2, regression- based methods already underperform relative diffusion-based approaches, and our adaptive rollout strategy is applicable only to the latter.
5 Conclusion
We proposed DiffusionRollout, diffusion-based neural PDE solver designed to predict long rollouts of physical systems governed by PDEs. Building on the benefits of generative modeling for PDE modeling, we introduce a method that leverages the norm of standard deviations computed from multiple samples drawn from pre-trained diffusion model-as surrogate mcasure of predictive uncertainty. This uncer- tainty estimate is then used during the autoregressive sampling process to reject erroncous predictions that could otherwise propagate and significantly degrade rollout. accuracy. Tn experiments on three challenging
13

---

Published in Transactions on Machine Learning Research (02/2026)
benchmark problems, the proposed method outperforms state-of-the-art baselines and demonstrates clear performance gains over conventional autoregressive sampling, which, by contrast, does not account for po- tential prediction errors during inference.
Limitations and Future Work. While DiffusionRollout improves long-term prediction accuracy and inherits key benefits of generative modeling, it relies on an iterative sampling process, which slower than regression- bassed neural operators. A promising direction for future work integrate recent one-step few-step diffusion models into adative rollout framework for improved inference efficiency. In addition, although our work primarily focuses addressing exposure bias along the temporal dimension, similar issue also arises during the diffusion sampling process due to its iterative nature. We believe this represents another important axis that merits further investigation. Finally, an interesting future direction is to explore learning a dedicated planner module that adapts the uncertainty threshold based on the current context.
Acknowledgments
This work was supported by the NRF of Korea grant (RS-2023-00209923)) IITP grants (RS-2024-00399817, 2025-2544,331, RSS-225--54443188, RS- 2SS22025-02631113) and the Industrial Technology Innovation Pro- gram (RS- (5--22222-2313326) all funded by the Korean government (MSIT and MOTIE), well grants from the DRB-KAIST SketchThhFutte Research Center.
References
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brand- stetter. Universal physics transformers. In NeurIPS, 2024.
Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Cheung. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. In Findings of the Association for Compulational Linguislies, 2022.
Dzmitry Bahdanau, Philemon Brakel, Kelvin Xu, Anirudh Goyal, Ryan Lowe, Joelle Pineau, Aaron Courville, and Yoshua Bengio. An actor-critic algorithm for sequence prediction. In ICL.R, 2017.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In NearIPS, 2015.
Lucas Berry, Axel Brando, and David Meger. Shedding light on large generative networks: Estimating epistemic uncertainty in diffusion models. In Uncertainty in Artificial Intelligence, 2024.
Shuhao Cao. Choose a transformer: fourier galerkin. In NearIPS, 2021.
Matthew A Chan, Maria J Molina, and Christopher Metzler. Estimating epistemic and aleatoric uncer- tainty with single model. In NewrIPS, 2024.
Boyuan Chen, Diego Martf Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. In NewrIPS, 2024.
Yu Chen, Lingfei Wu, and Mohammed J. Zaki. Reinforcement learning based graph-to-sequence model for natural question generation. In ICLR, 2020.
Woojin Cho, Minju Jo, Haksoo Lim, Kookjin Lee, Dongeun Lee, Sanghyun Hong, and Noseong Park. Pa- rameterized physics-informed neural networks for parameterized pdes. In ICML, 2024.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In ICLR, 2021.
Stefan Depeweg, José Miguel Hernández-Lobate Finale Doshi- Velez, and Steffen Udluft. Decomposition of uncertainty in bayesian deep learning for efficient and risk-sensitive learning. In JCML, 2018.
14

---

Published in Transactions on Machine Learning Research (02/2026)
Drummond B Fielding, Eve C Ostriker, Greg L Bryan, and Adam S Jermyn. Multiphase gas and the fractal nature of radiative turbulent mixing layers. The Astrophysicol Journal Letters, 894(2):1.24, 2020-
J. Gawlikowski, C. R. N. Tassi, M. Ali, et al. A survey of uncertainty in deep neural networks. Artificial Intelligence Review, 56(Suppl 1):1513 1589, 2023. doi: 10.1007/1114422-22--05556, URL https://doi. org/10 10111.0.0001/// 100.1007/10462--000 -26-------922------9-
Ernst Hairer, Gerhard Wanner, and Syvert P. Nørsett. Solving Ordinary Differential Equations I. Springer, 1987.
Jian Han, Jinlai Liu, Yi Jiang, Bin Yan, Yuqi Zhang, Zehuan Yuan, Bingyue Peng, and Xiaobing Liu. Infinity: Scaling bitwise autoregressive modeling for high-resolution image synthesis. In CVPR, 2024.
Zhongkai Hao, Chengyang Ying, Zhengyi Wang, Hang Su, Yinpeng Dong, Songming Liu, Ze Cheng, Jun Zhu, and Jian Song. Gnot: A general neural operator transformer for operator learning. In JCML, 2023.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. In NewrIPS, 2020.
Peiyan Hu, Rui Wang, Xiang Zheng, Tao Zhang, Haodong Feng, Ruiqi Feng, Long Wei, Yue Wang, Zhi-Ming Ma, and Tailin Wu. Wavelet diffusion neural operator. In ICLR, 2025.
Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park. Diffusionpde: Generative pde-solving under partial observation. In NewrIPS, 2024.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR, 2020.
Felix Koehler, Simon Niedermayr, Nils Thuerey, et al. Apebench: A benchmark for autoregressive neural emulators of pdes. Advonces in Neural Information Processing Systems, 111120252-120333 2024.
Aditi S. Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Character- izing possible failure modes in physics-informed neural networks. In NewrIPS, 2021.
Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkunar. Fourcastnet: Accelerating global high- resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, 2023.
Wilhelm Kutta. Beitrog zur näherungsweisen Integration totaler Differentialglecchungn. Teubner, 1901.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncer- tainty estimation using deep ensembles. In NewrIPS, 2017.
Rémi Leblond, Jean-Baptiste Alayrac, Anton Osokin, and Simon Lacoste-Julien. SEARNN: Training RNNs with global-local losses. In ICLR, 2018.
Zijie Li, Kazem Meidani, and Amir B. Farimani. Transformer for partial differential equations' operator learning. 2023a.
Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. In NeurIPS, 2023b.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadeneshelli Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural Operator: Graph Kernel Network for Partial Differential Equations. arXiv, 2020.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential Equa- tions. In ICLR, 2021.
15

---

Published in Transactions on Machine Learning Research (02/2026)
Zongyi Li, Miguel Liu-Schiaffni, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Learning chaotic dynamics in dissipative systems. In NewrlPS, 2022.
Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Moham- mad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Arizzadenesheli, and Anima Anandku- mar. Geometry-Informed Neural Operator for Large-Scale 3D PDEs. In NewrIPS, 2023c.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. JMLR, 2024.
Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E Turner, and Johannes Brandstetter. PDE-
Refiner: Achieving Accurate Long Rollouts with Temporal Neural PDE Solvers. In NewrlPS, 2023.
Ruben Ohana, Michael McCabe, Lucas Thibaut Meyer, Rudy Morel, Fruzsina Julia Agocs, Miguel Beneitez, Marsha Berger, Blakosley Burkhart, Stuart B. Dalziel, Drummond Buschman Fielding, Daniel Fortunato, Jared A. Goldberg, Keiya Hirashima, Yan-Fei Jiang, Rich Kerswell, Suryanarayana Maddu, Jonah M. Miller, Payel Mukhopadhyay, Stefan S. Nixon, Jeff Shen, Romain Watteaux, Bruno Régaldo-Saint Blan- card, François Rozet, Liam Holden Parker, Miles Cranmer, and Shirley Ho. The well: a large-scale collection of diverse physics simulations for machine learning. In NeurIPS, 2024.
OpenAL. Gpt-4 technical report. arXiv, 2024.
William Peobles and Saining Xic. Scalable diffusion models with transformers. In ICCV, 2023.
Maziar Raissi, Paris Perdikaris, and Goorge E Karniadakis. Physics-informed neural networks: A deep learn- ing framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
Marc' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojeiech Zaremba. Sequence level training with recurrent neural networks. In ICLR, 2016.
David Ruhe, Jonathan Heek, Tim Salimans, and Emiel Hoogeboom. Rolling diffusion models. In JCML, 2024.
Carl Runge. Über die numerische auflösung von differentialglichnngen. Mathematische Anmalen, 1895.
Aliaksandra Shysheya, Cristiana Diaconu, Federico Bergamin, Paris Perdikaris, José Miguel Hernández- Lobato, Richard E. Turner, and Emile Mathieu. On conditional diffusion models for pde simulations. NewrIPS, 2024.
Elvis do A. Soares, Amaro G. Barreto, and Frederico W Tavares. Exponential Integrators for Phase- Field Equations using Pseudo-spectral Methods: A Python Implementation. arXiv, 2023.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2022.
Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du, Russ Tedrake, and Vincent Sitzmann. History- guided video diffusion. arXio, 2025.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Tn NewrIPS, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Scoro-basod generative modcling through stochastic differential equations. In ICLR, 2021.
Keyu Tian, Yi Jiang, Zehuan Yuan, BINGYUE PENG, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. In NewrIPS, 2024.
Sifan Wang and Paris Perdikaris. Long- integration of parametric evolution equations with physics- informed deeponets. J. Comput. Phys., 2023.
16

---

Published in Transactions on Machine Learning Research (02/2026)
Tian Wang and Chuang Wang. Latent neural operator for solving forward and inverse pde problems. In NewrIPS, 2024.
Sean Welleck, Ili Kulikov, Stephen Roller, Emily Dinan, Kyunghyun Cho, and Jason Weston. Neural text generation with unlikelihood training. In ICLR, 2020.
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A fast transformer solver for pdes on general geometries. In ICML, 2024.
Zipeng Xiao, Zhongkai Hao, Bokai Lin, Zhijie Deng, and Hang Su. Improved operator learning by orthogonal attention. In NeurIPS, 2024.
Yifan Xu, Kening Zhang, Haoyu Dong, Yuezhou Sun, Wenlong Zhao, and Zhuowen Tu. Rethinking exposure bias in language modeling. arXiv preprint arXi::1900.1 11111990..111.,... 2019.
Wen Zhang, Yang Feng, Fandong Meng, Di You, and Qun Liu. Bridging the gap between training and inference for neural machine translation. In Proceedings of the 57th Annuaal Meeling of the Association for Computational Linguistics, 2019.
17

---

Published in Transactions on Machine Learning Research (02/2026)
Appendix
A.1 Additional Qualitative Results
We present additional qualitative results in Fig. A8 where RiflusionRollout produces the most accurate predictions, as visually indicated by the darker error maps.
Input
Target
FaetFormce
Transovver
PDE-Refiner
Diffusion
Li
2021
Li es al. 2023b
Wi ol.
2024
Hi et
2025
Lippe et al 2023
Rollout
Gray-Scott
Ohana et al.
2024
Turbuleut Flow OOhana al.
2024
Figure A8: Additional Qualitative Results.
18

---

Published in Transactions on Machine Learning Research (02/2026)
Input
Target
FNO
FacePormer
Tranzolee
WDNO
PDE-Refaecr
Diffusion
Cahn-Hilliare (Surrss al. 2023
Anisotropic Diffusion Koebler et al. 2024)
B
Figure A8: Additional Qualitative Results.
19

---

Published in Transactions on Machine Learning Research (02/2026)
Ioput
Target
FNO
FacePormer
Tranzover
WDNO
PDE-Refaeer
Diffusion
Figure A8: Additional Qualitative Results.
20


Reviews:
[{"review_id": "4l404RjRy5", "reviewer": "Reviewer_U8mi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors tackle the problem of PDE solving, framing it as conditional generation, and sampled via autoregressive rollouts.\nIn particular, they stress that although successful, this approach's performance degrades for long horizon since error tends to accumulate leading to what is referred as 'exposure bias' (i.e. distribution shift when chaining the conditional generation).\nThey also show that the model's variance is predictive of the prediction error, and use this to automatically adjust the rollout window.", "claims_and_evidence_yn": "Yes", "claims_explanation": "## Strengths\nThe manuscript is overall pretty easy and pleasant to read.\nThe introduced sampling procedure is quite simple yet neat.\nThe suggested method is evaluated on a wide range of datasets.\n\n## Weaknesses\nI feel like the motivation for using the conditional's variance as an estimate of prediction error could be strengthen.\nAlbeit quite impressive, I find the empirical result section a bit sparse in details, and some of the result not obvious to interpret. As such I have written quite a couple of remarks/questions in the hope to get answers and better understand these results.", "audience_yn": "Yes", "audience_explanation": "This manuscript is tackling a relevant task since solving PDE is ubiquitous and extremely costly as the resolution or horizon increases.\nWhat's more this manuscript shows that one can adapt the 'time step' depending on the 'model's confidence' to reach a better accuracy-compute tradeoff.", "requested_changes": "- Sec 3.2: I feel like it would be worth saying a bit more about uncertainty. Here there is no aleatoric uncertainty right since we are interested in PDEs there is no 'real' noise. The only uncertainty is thus epistemic. Given $u_s$ a solution of the PDE at time $s$, any future state $u_s$ at time $t>s$ can be _deterministically_ predicted via Eq (1). What breaks this is the space discretisation, and thus $p(u_t|u_s)$ is not fully 'determined' as a such this conditional has non zero variance. That is why conditioning on longer past trajectory narrows the conditional's variance.\n- Sec 3.3:\n    - Assuming that the space discretisation is not too coarse, the most of the uncertainty is _epistemic_ which is exactly the uncertainty that is aimed to be measured. I would suggest making this connection explicit as it may help justify why this is a 'natural' measure of uncertainty'.\n    - Also worth stressing that using the conditionally generated sample's _standard deviation_ as the measure of uncertainty is reasonable because these conditionals are very likely _unimodal_ as they are only conditioned and predicting over a small period of time.\n    - I am not totally sure to know/understand in Eq (4) what is the l2 norm for there? What is $\\sigma$? Worth clarifying these notations. I would use $\\hat{\\mathrm{Std}}$. Estimating this Std though necessitate to get several samples though right? Which would not be necessary otherwise, thus somewhat counterbalancing the original motivation of getting a better accuracy/compute tradeoff. Or is the Std estimated across the 6 prediction window states?\n- Sec 4.2:\n  - I am confused about the difference with PDERefiner, as I perhaps wrongly assumed that this work was building on PDERefiner, but changing the sampling procedure to choose on the fly the prediction window step $s$. If so it feels that the performance of 'DiffusionRollout' should be upper bounded by 'PDERefiner' $(s = 1)$ and lower bounded by 'PDERefiner' $(s = 6)$. I thought that the hope was that this adpaptive step procedure would perform competitively than always choosing $s=1$ but for a fraction of the compute. If 'DiffusionRollout' is _not_ based on 'PDERefiner', then Tab 1 should have results for 'DiffusionRollout' with $s=1,...,6$.\n  - Since efficiency is at the core of the motivation for this sampling method it it quite necessary to have some metric measuring this across different models. E.g. is one NFE of PDERefiner the same as one NFE for 'DiffusionRollout'? Are both using the same number of sampling step (per conditional step)?\n  - Are the sampling procedure normalised for compute? I assume not since PDERefiner with $s=6$ should be roughly 6 times cheaper than with $s=1$. \n- Sec 4.3: 'temporal derivative computed using the predicted future segment' -> Can you say more about this, e.g. write down the formula? Is this looking at the normal of the finite difference estimate of the temporal derivative? Still this method seems quite good too!\n- Sec 4.24:\n  - Worth making Tab 3 into a figure! \n  - Since $s_avg=1.12$ for $\\tau=1$ this means that actually it's really not so different then simply choosing $s=1$ right.\n  - Thus what's interesting is the hopefully better trade off for higher value of $s_avg$. It would be worth having on that Tab (or Figure) the comparison with the naive approach of always choosing $s=1,..,6$.\n  - Also it would be really useful for higher $s_avg$ to see the average value of $s$ for each time along the trajectory. I assume that it starts large ish and then decreases pretty quickly till reaching 1?\n  - It is not obvious to me that the optimal performance is for $s_avg>1$ and not $s_avg=s=1$. As in there is a range of $tau$ for which there is no trade off between accuracy and speed, it's a win-win. What is your explanation for this?\n  - Tab 4 & 5: Similarly, why would a worse estimate lead to a better choice for $s$?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PHriRLbFiS", "reviewer": "Reviewer_rtig", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes the DiffusionRollout method, whose core idea lies in adaptive step size selection. The authors construct a proxy task based on the prediction error of diffusion models to evaluate the confidence of predictions at different steps of the Solver model. Furthermore, extensive experimental evaluations validate the effectiveness of the method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The author effectively validated the core motivation of this paper in Figure 2, demonstrating a positive correlation between the prediction error of the diffusion model and the uncertainty of the Solver model, consistent with the design of the proposed method.", "audience_yn": "Yes", "audience_explanation": "Investigating PDE solvers based on generative models is a meaningful task, and the adaptive step selection explored in this paper is a key factor in the solution process. I believe readers of TMLR will find this topic of interest.", "requested_changes": "To my knowledge, the exposure bias issue is not confined to autoregressive generation but is also prevalent in diffusion models. It manifests specifically in samples generated with fewer inference steps (e.g., 300 NFE), outperforming those from full inference (1000 NFE). For instance, DDPM-IP [1] employs input perturbation to mitigate this problem. I believe this is relevant to the methods explored in this paper. The authors should analyze and discuss this issue to provide a more in-depth exploration of exposure bias.\n\n[1] Input Perturbation Reduces Exposure Bias in Diffusion Models. ICML-2023", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Y09NdhCuGV", "reviewer": "Reviewer_9oSw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The contributions are:\n- The authors propose the standard deviation of a conditional diffusion model to be the surrogate for model's predictive uncertainty and demonstrates a linear relationship between model predictive error and uncertainty on Gray-Scott ReactionDiffusion Equation dataset.\n- With the above observation, the authors propose DiffusionRollout, a strategy that adaptively select rollout step size for an autoregressive conditional diffusion model.\n- With extensive experiments, the authors demonstrate their strategy helps to mitigate the error accumulation effect over time.\n\nKey strengths:\n- The framework is easy to follow.\n- Most of the claims are backed by extensive experiments.\n\nWeaknesses:\n- While I find the method intuitively correct and backed with numerous experiments, I find it lack some theoretical grounds, without which it is more of an inference heuristic.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The claims are supported by evidence. For example:\n\n**Claim**: DiffusionRollout improves long-term prediction reliability.\n\n**Evidence**: DiffusionRollout does best with respect to the $T^{>0.9}$ metric on almost all datasets, indicating that it models long range predictions better than baselines.\n\nHowever, there is one thing I want to bring up. For the predictive error vs uncertainty relationship (Figure 2), only the Gray-Scott ReactionDiffusion Equation dataset is shown. I wonder if this relationship appears in all datasets.", "audience_yn": "Yes", "audience_explanation": "This paper provides a simple training-free inference strategy to adaptively select step size for modeling a PDE. This would be interesting to the operator learning audience.", "requested_changes": "1. I would like to see figures like Figure 2 for all other datasets to confirm if predictive error vs uncertainty relationship, a key claim by the authors, is actually universal.\n2. The ablation study looks good but to fully evaluate whether the strategy is helpful, I think the following may be useful:\n    1. Compare the performance of the model: do naively s=1, s=F and adaptive step size. This is because I find in Table 3, $s_{\\text avg}=1.120$ yields the best result, which means the model essentially only do 1 step forward at a time. A histogram of the selected step sizes may also be helpful.\n    2. For the above comparison, it is also helpful to show the computational cost for each case, for example NFE, number of FLOPs, etc so we can get a sense of the performance gain and the extra computational overhead.", "broader_impact_concerns": "This is no ethical concern.", "additional_comments": "As described above, I would hope more theoretical analysis is available for this framework to be sound and theoretically grounded rather than a heuristic.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)