On conditional diffusion models for PDE simulations
Authors: Aliaksandra Shysheya
Cristiana Diaconu
Federico Bergamin
Paris Perdikaris
José Miguel Hernández-Lobato
Richard E. Turner
Emile Mathieu
Venue: NeurIPS 2024 poster (2024)
Keywords: neural PDE solver
PDE
partial differential equation
forecasting
data-assimilation
diffusion
denoising
autoregressive
neural surrogate
reconstruction guidance
conditional score
Primary area: diffusion_based_models
TL;DR: This paper provides a comprehensive analysis and extension of the current state of score-based diffusion models trained on short segments from PDE trajectories, and evaluated on forecasting and data assimilation tasks.

Abstract: Modelling partial differential equations (PDEs) is of crucial importance in science and engineering, and it includes tasks ranging from forecasting to inverse problems, such as data assimilation. However, most previous numerical and machine learning approaches that target forecasting cannot be applied out-of-the-box for data assimilation. Recently, diffusion models have emerged as a powerful tool for conditional generation, being able to flexibly incorporate observations without retraining. In this work, we perform a comparative study of score-based diffusion models for forecasting and assimilation of sparse observations.  In particular, we focus on diffusion models that are either trained in a conditional manner, or conditioned after unconditional training. We address the shortcomings of existing models by proposing 1) an autoregressive sampling approach, that significantly improves performance in forecasting, 2) a new training strategy for conditional score-based models that achieves stable performance over a range of history lengths, and 3) a hybrid model which employs flexible pre-training conditioning on initial conditions and flexible post-training conditioning to handle data assimilation. We empirically show that these modifications are crucial for successfully tackling the combination of forecasting and data assimilation, a task commonly encountered in real-world scenarios.

On conditional diffusion models for PDE simulations
Aliaksandra Shysheya' University of Cambridge as2975@cam. ac. uk
Cristiana Diaconu" University Cambridge cdd43@cam. ac uk
Federico Bergamin"
Technical University of Denmark fedbe@dtu.dk
Paris Perdikaris Microsoft Research Al4Science
José Miguel Hernández-Lobato
University Cambridge
paperdikaristhicrooss
jmh2330cam uk
Richard E. Turner University of Cambridge Microsoft Research AI4Science The Alan Turing Institute ret230cam ac uk
Emile Mathieu University of Cambridge
ebm320cam ac.uk uk
Abstract
Modelling partial differential equations (PDEs) is of crucial importance in science and engineering. and it includes tasks ranging from forecasting to inverse problems, such as data assimilation. However, most previous numerical and machine learning approaches that target forecasting cannot applied out-of-the-box for data assim- ilation. Recently, diffusion models have emerged as powerful tool for conditional generation, being able to flexibly incorporate observations without retraining. this work, we perform comparative study of score-basee based diffusion models for fore- casting and assimilation of sparse observations. In particular, we focus on diffusion models that either trained in conditional manner, conditioned after uncondi- tional training. We address the shortcomings existing models by proposing 1) an autoregressive sampling approach, that significantly improves performance in fore- casting, 2) a new training strategy for conditional score-based models that achieves stable performance over range of history lengths, and 3) a hybrid model which employs flexible pre-training conditioning on initial conditions and flexible post- training conditioning handle data assimilation. We empirically show that these modifications are crucial for successfully tackling the combination of forecasting and data assimilation, task commonly encountered real-world scenarios.
Introduction
Partial differential equations (PDEs) are ubiquitous they are powerful mathematical framework for modelling physical phenomena, ranging from the motion of fluids to acoustics and thermody- namies. Applications
varied weather forecasting [4], train aerodynamics profiling [70], and concert hall design [76]. PDEs can either solved through numerical methods like finite element methods [43], or, alternatively, there has been rising interest in leveraging deep learning to learn neural approximations of PDE solutions [56]. These approaches have shown promise for generating accurate solutions in fields such fluid dynamics and weather prediction [37, 52, 54).
Some of the most common tasks within PDE modelling are 1) forecasting, where the goal is to generate accurate and long rollouts based some initial observations, and 2) inverse problems,
"Equal contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Table 1: score-based methods considered in this work. Each method can be classified in terms of the score network (joint or conditional), the rollout strategy at sampling time, and the conditioning mechanism at inference time.
MODEL
SCORE
ROLLOUT
CONDITIONING
Joint AAO [60]
st (1,L
AAO
Guidance
Joint AR (ours)
So F1.1
AR
Guidance
Universal amortised (ours)
Sy (t, (1.
AR
Architecture/Guidance
which aim to reconstruct a certain aspect of the PDE (i.e. coefficient, initial condition, full trajectory. etc.) given some partial observations of the solution to the PDE. Data assimilation (DA) refers to a particular class of inverse problems [62]. with goal to predict or refine trajectory given some partial and potentially noisy observations. While in general these tasks are addressed separately, there are cases in which a model that could jointly tackle both tasks would be beneficial. For example, in weather prediction, common task is forecast future states, as well as update them after more observations from weather stations and satellites are available [29]. Currently, traditional numerical weather prediction systems tackle this in a two-stage process (a forecast model followed by data assimilation system), both with a similarly high computational cost. A simpler option would be to have a model that is flexible enough to both produce accurate forecasts as well as condition on any incoming noisy observations.
In this work, we investigate how a probabilistic treatment of PDE dynamics [9, 81] can achieve this goal, with a focus on score-based diffusion models [26, 64, 68] as they have shown remarkable performance in other challenging applications, ranging from image and video [27, 58] to protein generation [72. 78, 791. We perform a comparative study between the diffusion model approaches in Tab. 1, and evaluate their performance on forecasting and DA in PDE modelling. The unconditionally- trained model that is conditioned post-traing pest-training through reconstruction guidance [65] is referred to as
the joint model, while the amortised model directly fits conditional score during training.
The main contribution of this work is that we propose extensions to the joint and amortised models to make them successfully applicable tasks that combine forecasting and DA:
Joint model. We propose an autoregressive (AR) sampling strategy to overcome the limitations of the sampling approach proposed in [60]. which samples full PDE trajectories all-at-once (AAO). We provide theoretical justification for why AR outperforms AAO, and empirically show, through extensive experiments, that AR outperforms performs on par with AAO DA, and drastically improves upon AAO in forecasting, where the latter fails.
Amortised model. We introduce novel training procedure for amortised models for PDEs, that, unlike previous work [42], allows their performance to remain stable over a variety history lengths. Moreover, we propose hybrid model that combines amortisation with reconstruction guidance that is especially well-suited to mixed tasks, that involve both a forecasting and DA component.
Finally, this work provides, to the best of our knowledge, the first quantitative comparison between joint and amortised models for PDE modelling, with all other factors such as training dataset and model architecture carefully controlled.
2 Background
Continuous-tinn diffusion models. In this section we recall the main concepts underlying continu- ous diffusion models, and refer Song et al. [68] for a thorough presentation. We consider a forward noising process (z(()))20 associated with the following stochastic differential equation (SDE)
dx(()) (t)dtt dw((), x00) po.
with (t) an isotropic Wiener process, and Po the data distribution. The process defined by (1) is the well-known
process, which geometrically converges to the standard Gaussian distribution (0, 1). Additionally, for any 2 0, the noising kernel of (1) admits the following form (00))) N'' ((=((())eaz()) I)), with μe et2 and of 1 et Importantly, under mild assumptions on Po, the time-reversal process (E (t))c20 also satisfies SDE [10, 24] which is

---

given by
d'F(t) (t) (1-)( v log pr ( { dt + dwo(t),
where pr denotes the density of ar()). In practice, the score log Pt unavailable, and approximated by a neural network So (4, V log Pt. referred to as the score network. The parameters 0 are learnt by minimising the denoising score matching (DSM) loss [30, 68, 73]
(3) where the expectation is taken over the joint distribution of t ~ u(O0, 1]). =00) ~ Po and ær(t) ~ Pe)((-(z(()) from the noising kernel. Sampling involves discretising and solving (2). Optionally, to avoid errors accumulating along the denoising discretisation, each denoising step (predictor) can be followed by small number of Langevin Monte Carlo (corrector) steps [67].
( (t]
V (zgg)(
Conditioning with diffusion models. The methodology introduced above allows for (approxi- mately) sampling from s=()) ~ Po- However, in many settings, one is actually interested in sam- pling from the conditional p(ar(D)lv) given some observations y E ) with likelihood A(r))(a)) and A X' + y being measurement operator. The conditional is given by Bayes' rule p(z(0)ll) y) pp)) a (0)) /p((), yet it is not available in closed form. An alternative is to simulate the conditional denoising process [13, 15, 68]
d - =
log pi (T (01b)} di + dwo().
In the following, we describe the two main ways to sample from this process.
Amortising over observations. The simplest approach directly learns the conditional score v log Pe To (1)v) [26, 68). This can be achieved by additionally feeding y to the score network 8g, i.e. So (1, x(), y). and optimising parameters by minimising a loss akin to (3).
Reconstruction guidance. Alternatively, one can leverage Bayes' rule, which allows to express the conditional score w.r.t. ar(t) as the sum of the following two gradients
V llgg((()))l)
V Vlogp(z(t)) +
V log logl(t)))))
(5)
prior
guidunce
where V log((((()) can be approximated by score network (t v V log logp(r(()) trained on the prior data-with no information about the observation y. The conditioning comes into play via the V log p(v)x((t) term which is referred to as reconstruction guidance [11, 13, 27, 48, 65, 66, 80]. Assuming Gaussian likelihood p(y]=(0)) = N(y)Ar((), o with linear measurement A(x)
Ax, we have:
p(yl=(t)) p(ox=(0)))p((())((()) dr(0)= NNWAA(0), dz(0)
N (s)A (()
(6)
where (())) is not available in closed form and is approximated by Gaussian ((x()))((t) & ((=()))
with mean given by Tweedie's formula [18, 57] a(x(t)) é ([)(())) = (a(t) + (o and variance Var(z(0) [x([)] r?. where is the guidance schedule and is chosen monotonically increasing function [19, 53, 60, 65].
3 PDE surrogates for forecasting and data assimilation
In this section, we discuss the design space of diffusion models, outlined Tab. 1, for tackling forecasting and DA in the context of PDE modelling.
Problem setting. We denote by P :4 + li a differential operator taking functions u R4 x Z + X E lt as input. Along with some initial u(0, -) uo and boundary u(T, 8z) = f conditions, these define partial differential equation (PDE) 0.u P(u;a) = = (0.u, u. L.u... ... ax) with coefficients and where θ is a shorthand for the partial derivative θu/Or, /T, and θ u. 0 denote the partial
3

---

derivatives 0u/02, /2 respectively. We assume access to some accurate numerical solutions from conventional solvers, which are discretised both in space and time. Let's denote such trajectories by 21:1. = (d1....;,75)) ....) E Dx ~ Po with D and T. being the size of the discretised space and time domains, respectively.
We are generally interested in the problem of generating realistic PDE trajectories given observations, i.c. sampling from conditionals p(zi. Specifically, we consider tasks that involve conditioning on initial states, as well as on sparse space-time observations, reflective of some real-life scenarios such as weather prediction. Thus, we focus on two types of problems, and the combination thereof: (a) forecasting, which involves predicting future states HHLL given a length- H past history 21:H-11 ie. sampling from ((
(b) data assimilation (DA), which concerned with inferring the ground state of a dynamical system given some sparse observed states 2,, i.c. sampling from p(1111 (xo). Then, the goal is to learn a flexible PDE surrogate, i.e. machine learning model, that, unlike most numerical solvers [43], is able to solve the underlying PDE while accounting for various initial conditions and sparse observations.
3.1 Learning the score
Learning a diffusion model over the full joint distribution p(æ1:z (t)) can become prohibitively expensive for long trajectories, as it requires score network taking the full sequence input- sy (t, (t)), with memory footprint scaling linearly with the length of the sequence.
One approach to alleviate this, as suggested by Rozet and Louppe [60], is to assume Markov structure of order & such that the joint over data trajectories can be factorised into a series of conditionals p(Z111) = p(=1))((z2)e)
... p(xx+1 ( I 0(2,|2x--ks-1)) Here we are omitting the time dependency Z1.L = Z1.z (0) for clarity's 's sake. The score w.r.t. 24 can be written as
i-k
Vx logp(2x1..)) = V. logp(sr(ar>n: kn- 1) +
Σ F=1+1
Va,, log (1j-)))
Vx, logp(æ,, Fet1s -(- = Vz, 2: log p( ((g((())++t))
(7)
with k
1, whilst formulas for i S h and i L - i can be found in App. B.1. In the case of modelling the conditional p(z ly). we use the same Markov assumption, such that learning V., log p(xi:z/y) reduces learning local score Vz, log The rest of Sec. 3.1 describes learning the local score of the joint distribution, but the same reasoning can easily be applied to fitting the local score of the conditional distribution.
When noising the sequence with the SDE (1), there is no guarantee that (1) will still be k-order Markov. However, we still assume that (t) is Markov but with larger order k >> k, as motivated in Rozet and Louppe [60, See 3.1]. To simplify notations we use k = k in the rest of the paper. Consequently, instead of learning the entire joint score once V (1)2(1) log 1og.(211. (t)), we only need
to learn the local scores (4, k (t)) & V (1) log (Tiikkii (()) with window size of 2k + 1. The main benefit is that the network only takes as input sequences of size 2k 1 instead of L, with k << L. The local score network sy is trained by minimising the following DSM loss L(0):
V log P:jo (t) where the expectation is taken over the joint i ~ u(k + 11, k)), t ~ 4(10, 1)), x4 -o ~ Pu and (t) (=) given by the noising process.
E |s bx:t&
3.2 Conditioning
Reconstruction guidance. We now describe how tackle forecasting and DA by sampling from conditionals (r.))_l of the prior p(Z1.L)~win the trained joint local score st (tt,:;; described in Sec. 3.1. For both tasks, the conditioning information y To E R is a measurement of subset of variables in the space-time domain, i.c. IG A(x) = A vec(x) +n with vec(z) E N the space-time vectorised trajectory, 7) (O, oı and masking matrix {0, A= OxNN where rows indicate observed variables with 1. Plugging this in (6) and computing the score we get the following reconstruction guidance term
(y - At(x(t)) TA @+(x(t)) dx(()
V (logp(zo/((t))
(8)
r +

---

Summing up this guidance with the trained score over the prior as in (5), we get an approximation of the conditional score V log (t) and can thus generate conditional trajectories by simulating the conditional denoising process.
Conditioning via architecture. Alternatvely, instead of learning a score for the joint distribution and conditioning posterioni, one can directly learn an approximation a score network Se (t, Zes+ (t) x x{-Cx) is trained to fit the conditional score given the states 24-0: by minimising DSM loss akin (3). Here, we denote the number of observed states by C, while P stands for the length of the predictive horizon. Most of the existing works [35, 75] propose models, which we refer to as amortised models, where the number of conditioning frames C is fixed and set before training. Thus, a separate model for each combination of P and C' needs to be trained. To overcome this issue, similarly to [28]. we propose a universal amortised model, which allows to use any P and C. such that P + C = 2k + 1, during sampling. In particular, during training, instead of determining C' beforehand, we sample C¹ Uniform( ((0,..., 2k}} ) for each batch and train the model to fit the conditional score So zr; Cx:+ P ((()ati Ca)) thus amortising over C' as well.
the conditional score. Such
The conditioning on the previous states is achieved by feeding them to the score network as separate input channels. The conditioning dimension is fixed to the chosen window size of 2k + 1. To indicate whether a particular channel is present we add binary masks of size 2k + Even though we condition on Ci-C5i, the score for the whole window size Ti-C:i+ f is predicted. In our experiments, we found that this training scheme works well for our tasks, although another parameterization as [28], where the score is only predicted for the masked variables, could be used.
Architecture conditioning and reconstruction guidance. In the universal amortised model, conditioning on previous states happens naturally via the model architecture. However, this strategy is not practical for conditioning on partial observations (e.g. sparse space-time observations as in DA), as new model needs to be trained every time new conditioning variables are introduced. To avoid this, we propose to simultaneously use reconstruction guidance (as estimated in (8)) to condition on other variables, as well as keep the conditioning on previous states through model architecture.
3.3 Rollout strategies
In this section, we describe sampling approaches that can be used to generate PDE rollouts, depending on the choice of local score and conditioning. We provide in App. K pseudocode for each rollout strategy.
All-at-once (AAO). With trained joint score (t (x)):++((t) can generate entire trajectories Ti.z in one go, as in Rozet and Louppe [60]. The full score So (1, (()) s V log Pr ()))..(()))) is reconstructed from the local scores, as shown in App. B.1 and summarised in Rozet and Louppe [60, + 1] which is the central output
Alg 2. The ith component is given by s; = -2- of the score network, except for the start and the end of the sequence for which the score is given ( + 1 :]. respectively. If we observe some values 2o --(AAAT 1)), the conditional sampling is done by summing the full score 21. with the guidance term V log p(z, (1.t.))))))))))))) estimated with (8).
by $1.b = sa 11:22+1 (t))l: k + 1] and S
Autoregressive (AR). Let us assume that want to sample #C+1:L given Fı.C initial states. An alternative approach to AAO is to factor the forecasting prediction problem as p (1rc)) = Il: ((=:|2;--:+-1) with C < k. As suggested in Ho et al. [27], it then follows that we can sample the full trajectory via ancestor sampling by iteratively sampling each conditional diffusion process. More generally, instead of sampling one state at time,
can autoregressively generate P states, conditioning on the previous C ones, such that P C = 2k + 1. This procedure can be used both by the joint and the amortised conditional score. DA can similarly be tackled by further conditioning on additional observations appearing within the predictive horizon P at each AR step of the rollout.
Scalability. Both sampling approaches involve a different number of neural function evaluations (NFEs). Assuming p predictor steps and e corrector steps simulate (4), AAO sampling requires (1 + c)p NFEs. In contrast, the AR scheme is more computationally intensive as each autoregressive
step also costs (1 + e)p NFEs but with (1 (1+ (22(2++1) AR steps. We stress, though, that this
P
5

---

assuming the full score network evaluation $1:L = memory constraints, the full score in AAO would have be sequentially computed. Assuming only one local score evaluation fits in memory, it would require as many NFEs as AR steps. What's s more, as discussed in the next paragraph, AÃO in practice typically requires more corrector steps. See App. F.1 for further discussion on scalability.
is parallelised. In practice, due to
Modelling capacity of the joint score. As highlighted in Tab. I, when using a joint score, both AAO and AR rollout strategies can be used. We hypothesise there two main reasons why AR outperforms AAO. First, for the same score network with an input window size of 2k + 1, the AAO model has a Markov order & whilst the AR model has an effective order of 2k. Indeed, shown in (7), process of Markov order k yields a score for which each component depends on 2k: + 1 inputs.
Yet parameterising a score network taking 2k + 1 inputs, means that the AR model would condition on the previous 2k states at each step. This twice greater Markov order implies that the AR approach is able to model a strictly larger class of data processes than the AÃO sampling. See App. B.1 for further discussion on this.
Second, with AAO sampling the denoising process must ensure the start of the generated sequence agrees with the end of the sequence. This requires information to be propagated between the two ends of the sequence, but at each step of denoising, the score network has limited receptive field of 2k + 1, limiting information propagation. Langevin corrector steps allow extra information flow, but come with additional computational cost. In contrast, in AR sampling, there is no limit in information flowing forward due to the nature of the autoregressive scheme.
4 Related work
In the following, we present works on diffusion models for PDE modelling since these are the models we focus on in this paper. We present extended related work section in App. C, where we broadly discuss ML-based and classical solver-based techniques for tackling PDE modelling.
Recently, several works have leveraged diffusion models for solving PDEs, with a particular focus on fluid dynamics. Amortising the score network on the initial state, Kohl et al. [35] introduced an autoregressive scheme, whilst Yang and Sommer [81] suggested to directly predict future states. Recently, Cachay et al. [9] proposed to unify the denoising time and the physical time to improve scalability. Lippe et al. [42] built upon neural operators with iterative denoising refinement, particularly effective to better capture higher frequencies and enabling long rollouts. In contrast the above-mentioned work, others have tackled DA and super-resolution tasks. Shu et al. [63] and Jacobsen et al. [32) suggested using the underlying PDE to enforce adherence to physical laws. Huang al. [29] used an amortised model and inpainting techniques [46] for conditioning at inference time to perform DA on weather data. Rozet and Louppe [60] decomposed the score of long trajectory into a series of local scores over short segments to be able to work with flexible lengths and to dramatically improve scalability w.r.t. memory use. Concurrently to our work, Qu al. [55] extended this approach to latent diffusion models to perform DA on ERA5 weather data [25]. In this work, we show that such a trained score network can alternatively be sampled autoregressively, which guaranteed to lead to a higher Markov order, and empirically produce accurate long-range forecasts. Concurrently, Ruhe et al. [61] proposed to use a local score, not justified by any Markov assumption, together with frame-dependent noising process for video and fluid dynamics generation. Diffusion models for PDE modelling are indeed closely related to other sequence modelling tasks, such as video generation and indeed our universal amortised approach is influenced by Höppe et al. [28]. Voleti et al. [75].
Diffusion models for PDEs.
5 Experimental results
We study the performance of the diffusiin-based models proposed in Tab. I on three different described in more detail below. The code to reproduce the experiments is publicly available at https//ghbbb.omm/aabrgggggggg///
Data. In this work, we consider the 1D Prramotoo-Sivahinnse (KS) and the 2D Kolmogorov flow equations, with the latter being variant of the incompressible Navier-Stokes flow. KS is

---

fourth-order nonlinear ID PDE describing flame fronts and solidification dynamics. The position of the flame u evolves as 0-u + ud, u + 02u + v0ğu is solved on periodic domain with 256 points for the space discretisation and timestep Ar = 0.2. Training trajectories are of length 140AT, while the length of validation and testing ones is set to 640Ar. The Kolmogorov flow is 2D PDE that describes the dynamics of an incompressible fluid. The evolution of the PDE is given by 8 + u Vu - " + v - f = 0 and V u = O, where u represents the velocity field, y/ the viscosity, p the fluid density, p the pressure, and f is the external forcing. The considered trajectories have 64 states with 64 x 64 resolution and Â 0.2. The evaluation metrics are measured with respect to the scalar vorticity field = 02 -, We focus on these since their dynamics are challenging. and have heen investigated in prior work [17, 42]. We refer to App. D for more details on the data generation. In addition, in App. F.1 we present results for the simpler ID Burgers' equation.
where // is the viscosity. The equation
Models. For forecasting and DA, we train a single local score network on contiguous segments randomly sampled from training trajectories. We use window 9 model for KS and window 5 model for Kolmogorov, as these settings give good trade-off between performance and memory requirements. We show in F.3 that increasing the window size in KS does lead to improved perfor- mance in forecasting, but the gains are relatively small. We parameterise the score network so with a modern U-Net architecture [22, 59] with residual connections and layer normalization. For sampling, we use the DPM solver [44] with 128 evenly spaced discretisation steps. As a final step, we return the posterior mean over the noise free data via Tweedie's formula. For the guidance schedule we set ri = Yoo / u² and tune via grid search. We refer to App. D.4 for more details.
Evaluation metries. We compute two per time step metrics-mean squared error MSE Cov(ELL.rIz)
E[( - )² and Pearson correlation PI:L Var(as:.)Var(51: between model samples
d ~ pa and ground truth trajectories Z1:L ~ Po- We measure sample accuracy through RMSD =
Z 1=1 MSE, and high correlation time max = Imax x At with Imax = arg maxtgl:L {pr > 0.8).
5.1 Forecasting
As mentioned in Sec. 3, forecasting is a crucial component of the combined task we consider, so we investigate the ability of trained diffusion models sample long rollouts. The models evaluated on 128 and 50 test trajectories for KS and Kolmogorov, respectively. To understand the current state of score-based models for PDE forecasting we benchmark them against other ML-based models, including the state-of-the-art PDE-Refince [42] and a MSE-trained U-Net and Fourier Neural Operator (FNO) [39]. We follow the setup in [42] for the architectures and training settings of the baselines. For the diffusion models and the MSE-traine trained U-Net baseline, we experiment with both an architecture inspired from [60] and one inspired from [42] and report the best results. The performance both neural network architectures considered can be found in Tab. 5. More details about the different architectures can be found in App. D.4.
High corr. time (s)
KS
Kolmogorov
10.0
Joint AR Universul amortised Amortised
75
50-
7.5
T
25
5.0
1/8
2/7
3/6
4
5/4
6/3
7/2
811
1/4
2/3
3/2
4
P PIC
P
PIC
Figure 1: High correlation time (f) on the KS (left) and Kolmogorov (right) datasets for different P C' conditioning scenarios of the joint, amortised and universal amortised models, where P indicates the number of generated states and C the number of states conditioned upon. We show mean 3 standard errors, as computed on the test trajectories.
Fig. shows the performance of the joint, amortised and universal amortised models for different conditioning scenarios P C, with P the number of states generated, and C' the number of states conditioned upon at each iteration. For the KS dataset, the universal amortised model outperforms the plain amortised one across all P C; for Kolmogorov their performance is comparable (i.c. within 3
7

---

standard errors). For both datasets, the best (plain and universal) amortised models outperform the best joint AR model, with the exact best high correlation times shown in Fig. 2.
Stability over history length. The most interesting comparison between the models lies in how they leverage past history in the context of forecasting. Previous work notes that longer history generally degrades the performance of amortised models [42]. We find that this behaviour is dataset- dependent- the observation holds for KS, whereas for Kolmogorov the plain amortised models with longer history (i.e. 4 and 2| 3) tend to outperform the ones with shorter history (i.e. 2 and 4 1). Due to our novel training procedure for the universal amortised model (i.e. training over a variety of tasks), we obtain stable performance over all conditioning scenarios for the KS dataset. For Kolmogorov, the universal amortised model benefits from longer history, similarly to the plain amortised one. However, note that, unlike the plain amortised model, where a different model needs to be trained for each conditioning scenario, the universal amortised can tackle any scenario, allowing for trade-off between accuracy and computational speed at sampling time. The joint model generally benefits from longer conditioning trajectories, but is also fairly stable for the majority of P |C settings. We show more results App. F.3.
For complete comparison between the models in Tab. 1, we also perform AAO sampling for the joint model. However, as shown in App. F.I, the predicted samples rapidly diverge from the ground truth. Thus, the AR method of querying the joint model is crucial for forecasting.
High corr. time (s)
KS
Kolmogorov
Joint AR
100
73.2
88.7
7
78.6
76.5
X
91.4
112.2
10-
8.2
I
10
T
10.6
I
7.2 2
10.4
MSE-trained FNO
Amortised
Universal amortised
50-
5-
MSE-trained U-Net
PDE-Refiner
0-
Figure 2: The best joint and amortised models compared against standard ML-based benchmarks for forecasting on the KS (left) and Kolmogorov (right) datasets. High correlation time (f) is reported, showing the mean + 3 standard errors, as computed on the test trajectories. We show Tab. 5 the configurations for the best models.
Benchmarks. In Fig. 2, we report the comparison against popular methods for forecasting. The universal amortised model performs similarly to the MSE-trained U-Net on both datasets, while the joint model comparable to the MSE-trainedd FNO, performing on par with it for KS and slightly outperforming it for Kolmogorov. However, we stress that, unlike these benchmarks, which can only be straightfforwadde applied to forecasting, our models are significantly more flexible-they manage to achieve competitive performance with some of the MSE-trained benchmarks, yet can be applied to a wider variety of tasks (e.g. forecasting combined with DA).
Additional results. For a more thorough understanding of the models' behaviour, we perform further analysis on the frequency spectra of the generated trajectories in App. I. We show that, as noted in previous work [42], the generated KS samples tend approximate the low-frequency components well, but fail to capture the higher-frequency components, while the Kolmogorov samples generally show good agreement with the ground truth spectrum. We also investigate the long-term behaviour of our proposed models in App. J. Even when generating trajectories that are significantly longer than those used during training, the models are capable of generating physically-plausee states.
5.2 Offline data assimilation
In offline DA, we assume we have access some sparse observations. The experimental setting goes as follows. We first choose a number of observed variables, and then uniformly sample the associated indices. We assume to always observe some full initial states. More details can be found in App. G. We repeat this for different numbers of values observed, varying the sparsity of observations, from approximately third of all data to almost exclusively conditioning on the initial states. We compute the RMSD of the generated trajectories for six values of the proportion of observed data, evenly spaced on a log scale from 10-3 to 10 10-0.5 We show results for the joint AR model, joint AAO with 0 and corrections (indicated between brackets), and the universal amortised model. To provide
8

---

consistent comparison with the AAO approach from [60], all models use the U-Net architecture inspired by [60]. We tune the guidance strength y, and for AR we only report the P C setting that gives the best trade-off between performance and computation time (see Fig. 28 for comprehensive summary of performance depending on 7 and P C). We show the results on 50 test trajectories for both KS and Kolmogorov. More results can be found in App. G.
KS
Kolmogarov
RMSD
20
10-
0-
0
Comp. cost (min)
Joint AR
100-
Joint AAO (0)
100
Joint AAO (1)
50
50-
Universal
amortised
10
10
10
10 3
10
Proportion observed
Figure 3: RMSD (mean +3 3 standard errors) for KS (left) and Kolmogorov (right) the offline DA setting for varying sparsity levels (top), and the computational cost associated with each setting (bottom). The latter the same for all sparsity settings for AAO, but differs for AR since it depends on the P C setting that was used. The e in AAO (c) refers to the number of corrector steps used.
More efficient sampling. We observe in Fig. 3 that AR sampling outperforms AAO in the sparse ob- servation regimes, and performs on par with it in the dense observation regimes, provided corrections are used. However, Fig. shows that corrections significantly increase the associated computational cost. We believe the difference in performance is mostly due to the limiting information propagation in the AAO scheme, combined with the lower effective Markov order of AAO AR (k vs 2k for local score network with a window 24 - 1), as discussed Sec. 3.3. To further investigate this hypothesis, in App. G.2.1 we compare the results for window 5 (k = 2) AR model with window (k = 4) AAO model for KS These should, in theory, have the same effective Markov order. We observe that for the sparsest observation regime, there is still gap between AR (k 2) and AAO (k: = 4). potentially because the high sparsity of observations make the information flow in the AAO scheme inefficient. In fact, this is consistent with the findings from forecasting, where that the AAO strategy is unable to produce coherent trajectories if it does not have access to enough future observations. In contrast, for the other sparsity scenarios the performance of the two schemes becomes similar provided corrections are used. However, this comes with a significant increase in computational cost.
observe
Reconstruction guidance for universal amortised model. As outlined in Sec. 3.2, we can combine the universal amortised model with reconstruction guidance for conditioning on sparse observations. Amortised models do not grraihhtfoowardl handle DA - this can be seen in Fig. 29a, where we show that conditioning just through the architecture gives poor results in DA. In KS, the universal amortised model achieves lower RMSD for the very sparse observation scenarios (10-3 proportion observed), and for the very dense ones (100.5 proportion observed). However, the joint model outperforms it in the middle ranges. In Kolmogorov, the universal amortised model slightly outperforms the joint for all sparsity settings, with the results being within error bars. These findings indicate that they both viable choices for offline DA, with the best performing choice depending on the application.
Interpolation baseline. In App. G.3 we also provide quantitative and qualitative comparison to a simple interpolation baseline, proving that the studied models manage to capture the underlying physical behaviour better than such simple baselines.
5.3 Online DA: combining forecasting and data assimilation
Inspired by real-life applications such as weather prediction, we compare the models in the online DA setting, which represents the combination of forecasting and DA. We assume that some sparse

---

and noisy variables are observed once every states (in blocks of length s), and whenever new observations arrive (i.e. each DA step), we forecast the next states. At the first step, we perform forecasting assuming we have access to some randomly sampled indices, accounting for 10% of the first s-length trajectory. At the second step (after s time-steps), we refine this forecast with the newly observed set of variables. This repeats until we have the forecast for the entire length of the test trajectories. The performance is measured in terms of the mean RMSD between the forecasts and the ground truth (averaged over all the DA steps, each outputting trajectory of length f)) We set 8 = 20 and f = 400 for KS, and 5 = and f = 40 for Kolmogorov, resulting in 13, and 6 DA steps, respectively. The setup is visualised in App. H. and the key findings from the results in Tab. 2 are
AAO is unsuitable for tasks involving forecasting - As opposed to the offline setting, AAO fails at online DA, leading to significantly higher RMSD on both datasets.
The universal amortised model shows the best performance out of all the models, which is in line with the findings from the previous sections it outperforms joint AR in forecasting and performs on par with it in DA.
Table 2: RMSD (4) (mean + 3 standard errors over test trajectories) for online DA for the KS and Kolmogorov datasets. The e in AAO (c) refers to the number of corrector steps used.
DATA
Joint AR
Joint AAO (0)
Joint AAO (1) Universal amortised
KS Kolmogorov
11.6 + 0.6
25.3 + 0.1
26.3 + 0.1
7.7 + 0.6
7.8 + 0.5
27.4 + 0.2
30.1 + 0.2
4.9 + 0.4
Discussion
In this work, we explore different ways condition and sample diffusion models for forecasting and DA. In particular, we empirically demonstrate that diffusion models achieve comparable performance with MSE-trained models in forecasting, yet significantly more flexible as they can effectively also tackle DA. We theoretically justify and empirically demonstrate the effectiveness of AR sampling for joint models, which is crucial for any task that contains a forecasting component. Moreover, we enhance the flexibility and robustness of amortised models by proposing a new training strategy which also amortises over the history length, and by combining them with reconstruction guidance.
Limitations and future work. Although achieving competitive results with MSE-trained baselines in forecasting, the models still lack state-of-the-art performance. Thus, they are most practical in settings where flexible conditioning is needed, rather than a standard forecasting task. Although we have explored different solvers and number of discretisation steps, the autoregressive sampling strategy requires simulating a denoising process at each step, which is computationally intensive. We believe that this can be further improved, perhaps by reusing previous sampled states. Additionally, the guided sampling strategy requires tuning the guidance strength, yet we believe there exist some simple heuristics that are able to decrease the sensitivity of this hyperparameter, Finally, we are interested in investigating how the PDE characteristics (i.e. spatial / time discretisation, data volume, frequency spectrum, etc.) influence the behaviour of the diffusion- -ased models.
Acknowledgements
AS, RET and EM are supported by an EPSRC Prosperity Partnership EP/T005386/1 between the EPSRC, Microsoft Research and the University of Cambridge. CD is supported by the Cambridge Trust Scholarship. FB is supported by the Innovation Fund Denmark (0175-000 (0075-00 and the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science
(NNF200C0662066 and NNF200C0065611. 1. JMHL acknowledges support from a Turing AI Fellowship under grant EP/V023356/1. RET is supported by the EPSRC Probabilistic AI Hub (EP/Y028783/1) and gifts from Google, Amazon, ARM, and Improbable. We thank the anonymous reviewers for their key suggestions and insightful questions that helped improved the quality of the paper.
10

---

References
[1] William F. Ames. Numerical Methods for Partial Differential Equations. Computer Science and Applied Mathematics. Academic Press, New York, 1977. ISBN 0120567601. URL
[2] Anima Anandkumar, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Nikola Kovachki, Zongyi Li, Burigede Liu, and Andrew Stuart. Neural operator: Graph kernel network for partial differential equations. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2019.
[3] Simon Arridge, Peter Maass, Ozan Öktem, and Carola-Bibiane Schönlieb. Solving in- verse problems using data-driven models. Acta Numerica, 28:1-174, 2019. doi: 10.1017/ $0962492919000059,
[4] Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature, 525(7567), 2015.
[5] Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Prediction of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64(2), 2019.
[6] Kaushik Bhattacharya, Bamdad Hosseini, Nikola B. Kovachki, and Andrew M. Stuart. Model Reduction And Neural Networks For Parametric PDEs. The SMAI Journal of computational mathematics, 7:121-157, 2021.
[7] Johannes Brandstetter, Max Welling, and Daniel E. Worrall. Lie Point Symmetry Data Aug- mentation for Neural PDE Solvers. In International Conference on Machine Learning. PMLR, 2022.
[8] Johannes Martinus Burgers. A mathematical model illustrating the theory of turbulence. Advances in applied mechanics, 1:171-199, 1948.
[9] Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. DYffusion: a dynaamics-iffomeed diffusion model for spatiotemporal forecasting. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[10] Patrick Cattiaux, Giovanni Conforti, Ivan Gentil, and Christian Léonard. Time reversal of diffusion processes under finite entropy condition, September 2022.
[11] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving Diffusion Models for Inverse Problems using Manifold Constraints. Advances in Neural Information Processing Systems, 35, 2022.
[12] Hyungjin Chung. Byeongsu Sim, and Jong Chul Ye. Saine-Closc--Diffusse ating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Acceler-
[13] Hyungjin Chung, Jeongsol Kim, Michael Mccann, Mare L. Klasky, and Jong Chul Ye. Diffusion Posterior Sampling for General Noisy Inverse Problems. In The Eleventh International Conference on Learning Representations, ICLR, 2023.
[14] Steven M Cox and Paul C Matthews. Exponential time differencing for stiff systems. Journal of Computational Physics, 136(2):430-455 2002.
[15] Kieran Didi, Francisco Vargas, Simon V. Mathis, Vincent Dutordoir, Emile Mathieu, Urszula Ju- lia Komorowska, and Pietro Lio. A framework for conditional diffusion modelling with applications in motif scaffolding for protein design. 2023.
[16] Tobin A Driscoll, Nicholas Hale, and Lloyd N Trefethen. Chebfun guide, 2014.
[17] Yifan Du and Tamer Zaki. Evolutional Deep Neural Network. Physical Review E, 104(4), 2021.

---

[18] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106, 2011.
[19] Marc Anton Finzi, Anudhyan Boral, Andrew Gordon Wilson, Fei Sha, and Leonardo Zepeda- Nunez. User-defined Event Sampling and Uncertainty Quantification in Diffusion Models for Physical Dynamical Systems. In International Conference on Machine Learning. PMLR, 2023.
[20] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for -o-mage synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[21] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers. In International Conference on Learning Representations, 2021.
[22] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporal scale generalized PDE modeling. 2023.
[23] Ernst Hairer and Gerhard Wanner. Solving Ordinary Differential Equations II. Stiff and Differential-Algerruie Problems, volume 14. 01 1996. doi: 10.1007/978-33-62--09477
[24] Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, 1986.
[25] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz- Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. The era5 global reanalysis. Quarterly Journal of the Royal Meteorological Sociery, 146(730):1999-225) 2049, 2020.
[26] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[27] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. In Deep Generative Models for Highly Structured Data Workshop, ICLR, volume 10, 2022.
[28] Tobias Höppe, Arash Mehrjou, Stefan Bauer, Didrik Nielsen, and Andrea Dittadi. Diffusion models for video prediction and infilling. Transactions on Machine Learning Research, 2022.
ISSN 2835-8856.
[29] Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter Dueben, and Torsten Hoefler. Diffda: a diffusion model for weather-scale data assimilation. arXiv preprint 12Xiv:2401.0593: 2024.
[30] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[31] Victor Isakov. Inverse Problems for Purtial Differential Equations. Springer Publishing Company, Incorporated, 3rd edition, 2018. ISBN 3319847104.
[32] Christian Jacobsen. Yilin Zhuang, and Karthik Duraisamy. CoCoGen: Physically-Consistent and Conditioned Score-based Generative Models for Forward and Inverse Problems, December 2023.
[33] Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine acceleraaeeeeerrteed computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21), 2021. ISSN 0027-8424. doi: ...073/paas22101788111
URL https: pnas ottp////ooee////////////0002222
[34] Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning-aceo accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21), 2021.
[35] Georg Kohl, Li- Chen, and Nils Thuerey. Turbulent flow simulation using autoregressive conditional diffusion models. arXiv preprint arXiv:2309.01745 2023.
12

---

[36] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli. Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function Spaces. Journal of Machine Learning Research, 24(89), 2023.
[37] Lizao Li, Robert Carver, Ignacio Lopez-Gomez, Fei Sha, and John Anderson. Generative emulation of weather forecast ensembles with diffusion models. Science Advances, 10(13): eadk4489, 2024. doi: 20.1126sciadv.adkkel
[38] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential Equations. In 9th International Conference on Learning Representations, ICLR, 2021.
[39] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
[40] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03799 2021.
[41] Wenyuan Liao, Mehdi Dehghan, and Akbar Mohebbi. Direct numerical method for an in- verse problem of parabolic partial differential equation. Journal of Computational and Applied Mathematics, 232(():351366) 2009. ISSN 0377-0427. doi: https://do. 10116f.cam.2209.06.0011 URL https: //t::///www... http://wwww...iceeco www.scedieecdrcttetccc cttccccm///c//ccceccc htt///cienccerarriicccee p11//00377042099003622
[42] Phillip Lippe. Bastiaan S. Veeling, Paris Perdikaris, Richard E. Turner, and Johannes Brandstet- ter. PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[43] Wing Kam Liu, Shaofan Li, and Harold S. Park. Eighty Years of the Finite Element Method: Birth, Evolution, and Future. Archives of Computational Methods in Engineering, 29(6), 2022.
[44] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, 2023.
[45] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3), 2021.
[46] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[47] Emile Mathieu, Vincent Dutordoir, Michael Hutchinson, Valentin De Bortoli, Yee Whye Teh, and Richard E. Turner. Geometric Neural Diffusion Processes. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[48] Xiangming Meng and Yoshiyuki Kabashima. Diffusion Model Based Posterior Sampling for Noisy Linear Inverse Problems, May 2023.
[49] Nicholas H Nelsen and Andrew M Stuart. The Random Feature Model for Cpputt-utputtta Maps between Banach Spaces. SIAM Journal on Scientific Computing, 43(5), 2021.
[50] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, pirf-peeformance deep learning library. In Advances Neural Information Processing Systems
32. 2019.
13

---

[51] Ravi G. Patel, Nathaniel A. Trask, Mitchell A. Wood, and Eric C. Cyr. A physics-informed operator regression framework for extracting data-driven continuum models. Computer Methods in Applied Mechanics and Engineering, 373, January 2021.
[52] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadeneshelii Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. FourCastNet: A Global Data- driven High-rcsolution Weather Model using Adaptive Fourier Neural Operators, February 2022.
[53] Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, and Brian Karrer. Training-free Linear Image Inversion via Flows, September 2023.
[54] Ilan Price, Alvaro
Ferran Alet, Tom R. Andersson, Andrew El-Kadi, Do- minic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson. Gencast: Diffusion-based ensemble forecasting for medium-range weather, 2024.
[55] Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine. Deep generative data assimilation in multimodal setting. arXiv preprint arXiv:2404.06660 2024.
[56] Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics Informed Deep Learning
(Part I): Data-driven Solutions of Nonlinear Partial Differential Equations. 2017.
[57] Herbert E Robbins. An empirical bayes approach to statistics. In Breakthroughs in Statistics: Foundations and basic theory, pages 388-394. Springer. 1992.
[58] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[59] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assisted Inter- vention, 18th International Conference, Springer, 2015.
[60] François Rozet and Gilles Louppe. Score-based Data Assimilation. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[61] David Ruhe, Jonathan Heek, Tim Salimans, and Emiel Hoogeboom. Rolling diffusion models. arXiv preprint 00X0v:240..094400, 2024.
[62] Daniel tion. arXiv preprint arXiv:1810.06199 2018
Andrew M Stuart, and Armeen Taeb, Inverse problems and data assimila-
[63] Dule Shu, Zijie Li, and Amir Barati Farimani. A Physics-informed Diffusion Model for High-fidelity Flow Field Reconstruction. Journal of Computational Physics, 478, 2023.
[64] Jascha Dick-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning. 2015.
[65] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Gseudoinverse-Guided Diffusion Models for Inverse Problems. In International Conference Learning Representations, 2022.
[66] Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-Guided Diffusion Models for Plug-and-Play Control- lable Generation. 2023.
[67] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[68] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 1021.
14

---

[69] Antonio Stanziola, Simon Arridge, Ben T. Cox, and Bradley E. Treeby. A research framework for writing differentiable pde discretizations jax. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
[70] Maciej Szudarek, Adam Piechna, Piotr Prusiñski, and Leszek Rudniak. Cfd study of high-speed train in crosswinds for large yaw angles with rans- based turbulence models including geko tuning approach. Energies, 1581)) 2022.
[71] Jos Torge, Charles Harris, Simon V. Mathis, and Pietro Lio. Difthopp: A graph diffusion model for novel drug design via scaffold hopping. arXiv preprint urXiv:2308.074j1 2023.
[72] Brian L Trippe, Jason Yim. Doug Tischer. Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. International Conference on Learning Representations (ICLR), 2023.
[73] Pascal Vincent. connection between score matching and denoising autoencoders. Neural 2011.
computation,
[74] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapcau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Erie Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SeiPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020.
[75] Vikram Voleti, Alexia Jolicoeur- Martinnea, and Christopher Pal. Mevd: Masked conditional video diffusion for prediction, generation, and interpolation. In (NeurIPS) Advances Neural Information Processing Systems, 2022.
[76] Michael Vorlaender, Dirk Schröder, Sönke Pelzer, and Frank Wefers. Virtual reality for architectural acoustics. Journal of Building Performance Simulation, 8, 01 2015.
[77] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed deeponets. Science advances, 7(40), 2021.
[78] Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel. Samuel Pellock, Alexis Courbet, William Sheffler, Jue Wang. Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Back, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[79] Joseph I.. Watson, David Juergens, Nathaniel R. Bennett, Brian Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel. Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay, Tommi Jaakkola, Frank DiMaio, Minkyung Back, and David Baker. De novo design of protein structure and function with RFdiffusion. Nature, 2023.
[80] Kevin E Wu, Kevin K Yang, Rianne van den Berg, James Zou, Alex X Lu, and Ava Amini. Protein structure generation via folding diffusion. arXiv preprint (tti::22099766111 2022.
[81] Gefan Yang and Stefan Sommer, A denoising diffusion model for fluid field prediction. arXiv e-prints, pages arXiv-2301, 2023.
[82] Qinsheng Zhang and Yongxin Chen. Fast Sampling of Diffusion Models with Exponential Integrator. The Eleventh International Conference on Learning Representations, ICLR. OpenReview.net, 2023.
15

---

[83] Yinhao Zhu and Nicholas Zabaras. Bayesian deep convolutional encoder encoder-decoder decoder networks for surrogate modeling and uncertainty quantification. Journal of Computational Physics, 366, August 2018 8.
16

---

On conditional diffusion models for PDE simulations (Appendix)
A Organisation of the supplementary
In this supplementary, we first motivate in App. B the choice made regarding the (local) score network given the Markov assumption the sequences. We then provide App. C, followed by description of the data generating process in App. D) for the datasets used the experiments in Sec. 5. We also provide further implementation details. After, in App. we assess the ability of the all-at-once sampling strategy-whine relies on reconstructing the full score from the local ones-to perform accurate forecast, and in particular show the crucial role of correction steps in the discretisation. We then provide more results on forecasting in App. F, including study of the influence the guidance strength hyperparameter y and of the conditioning scenario, a comparison with widely used benchmarks in forecasting, such as climatology and persistence, and results for different U-Net architectures. In App. G we show more experimental evidence on the offline DA task, including the influence of the guidance strength and of the conditioning scenarios, well as perform a comparison between an AR model and an AAO model with twice higher Markov order. We also include comparison with an interpolation baseline. In App. H we provide further experimental evidence for the online DA task for a range of conditioning scenarios for the joint and amortised models. In App. we provide frequency spectra of the trajectories generated by our methods and the forecasting baselines, and App. J we analyse the long-term behaviour of the diffusion models. Finally, provide pseudo-code in App. K.
reference related work in
B
Markov score
In this section, look at the structure of the score induced by Markov time series. In particular, we lay out different ways to parameterise and train this score with a neural network. In what follows assume AR-1 Markov model for the sake of clarity, but the reasoning developed trivially generalises to AR-k Markov models. Assuming an AR-1 Markov model, we have that the joint density factorise
p(11.)) p(x1) II p(as|1r-1).
(9)
1=2
B.1 Window 2k + 1:
Looking the score of (9) for intermediary states have
logp(x1.L)
= Va, logp(1(41(xs) + V., log p(1,)(=11))
(10)
Vz. log p(ge:41,,
Va, log 11/2:11) V log p(pp)()1)
V., log 1i, 2(-1)
= V.. log log 1+41)
(11)
17

---

Similarly, for the first state, we get
Vay log p(11:L)
= Vay log p(1)( + 2z, log pog (x1)
Vaj ;
Van p(1,, 22) + V log p(1g)(2))
Vx, p(z1, x2))
Ve, log plaz|22 22 x1) since ا(اد
Va, 22))
(12)
and for the last state
Vay logp(xn).)
= V Fn (ggp(rr/21111)
(13)
Vx, log (
since =p(2 (21-)
Vx, 1
V xx logp(oL-1,
1,
Vx, 10g
VEL log ( 2,
(14)
Noised sequence. Denoising diffusion requires learning the score of noised states. The sequence Ti.z is noised with the process defined in (1) for an amount of time t: (1) (1: ~ Pejo- With T1:L an AR-k Markov model there is no guarantee that T1.L (t) is still k-Markov. Yet, we can assume that (1) is AR-k' Markov model where K' >> k, in Rozet and Louppe [60]. Although may seem in contradiction with the previous statement, in what follows we focus on the AR-1 assumption since derivations and the general reasoning is much easier to follow, yet reiterate that this generalises to the order k.
Local score. Eqs. (11),
(12)
and (14) suggest training score network fit
Sa (t, (1) x, (t), x:+1((t)) (t)) 2 V (logm(a--1((,,((()).))))
by
sampling tuples
Tz-1, 2441 po and noising them.
Full score. The first, intermediary and last elements of the score are then respectively given by
Van (4) log (x1. s=(t,
Va,(t) log pi (#1.1 so(t, st(t
Vay(t) log Pr (x 2 ~ sa st
Joint sampling. The full score over the sequence æ1:L can be reconstructed via the trained local score network taking segments of length 3 as input. This allows for sampling from the joint. All of the above generalise to AR-k models, where the local score network takes segment of length 24 + as input.
Guided AR. Alternatively for forecasting. having learnt local score network (t), (t), (t)), iee modelling joints p(r;-1, 2; Ii+1), we can condition on T;-1, 2i to sample I;+11 leveraging reconstruction guidance-as as described Sec. 3.3. Effectively, this induces an
of AR-1- Markov model, or in general, an AR-2k-instead of k-Markov model. As such, this score parameterisation takes more input than strictly necessary to satisfy the le Markov assumption.
Amortised AR. Obviously we note that, looking at (10) and (13), one could learn a con- ditional score network amortised the previous value such that so (t))=;;-)0)) s Va,(t) log (2g)(z(()))
This unsurprisingly allows forecasting with autoregressive rollout as in Sec. 3.3. Additionally learning a model over p(æ1) and summing pairwise the scores as (10) would also allow for joint sampling.
18

---

B.2 Window + 1:
In App. B.I, we started with an AR-k model, but ended up with learning a local score taking 2k: + inputs. Below we derive alternative score parameterisation. Still assuming AR-1 model (9), we have that the score for intermediary states can be expressed as
V xe log p(11.))
V log (( + v., log (( (1i-1)
= Vx, logp(z;( Vx, log p(x;) + +V x; logp(r;(2;i))) + Va, log p(a;-1
via Bayes' rule
Va, 2, (ggp(z//<411) Vx, log logp(1;+1)
logp(z;) + Vx, log ((t,,
vx, 204, v., logp(z;) + Vx, logp(r),
(15)
Local score. From (15), we notice that we could learn a pairwise score network such that (e (1),,, 22,(1) log ()) (4)), and x1)) R Va.(9 logp(z;(()) (Q), trained by sampling tuples I;-1, 24, I++1 Po and noising them. Note that this requires 3 calls to the local score network instead of as in App. B.I.
For k >
Vz, logp(x).L) v., log
, 3++1: ................................ Fitk)
it
Vx, 1og
Vx, logp(0)loj
/=++1
i
Vz, log
Σ
Va, log log
Σ
V
logp(om-
i+
i
V.. logp - Σ
Va, log
(16)
j=
From (16), we obtained 2k: + window size model, while having the scores with window size of and k + 1. This would require roughly twice more forward passes during AAO sampling, but will double the effective window size.
Full score.
The first, intermediary and last elements of the score are then respectively given by
V (1( (og)(x):2()) Vni (t)
V (n log + Vayte) (og
(t)
Vaat)) log so 8; (1), n (1), 1+)
So (t), (((t so(t,,,(t))
Fal() logp(si.1
V Van( log 1o
V Vaz(1) (og
5o (1),, (1)
Joint sampling. joint sampling.
The full score over 21.((t) could therefore be reconstructed from these, enabling
AR. Additionally, autoregressive forecasting could be achieved with guidance by iterating over pairs ( J. thus using local score of window size k: + contrast App. B.I which requires an inflated window of size 2k 1.
19

---

C More related work
In Sec. 4, we presented relevant works that focus machine learning models and more specifically diffusion models for PDE modelling. Our work builds upon different conditional generation protocols used for conditional sampling in diffusion models. Sec. 2 presents the techniques considered for our models. In the following, we summarise all the available methods for building a conditional diffusion model.
Conditional diffusion models. The development of conditional generation methods for diffusion models is an active area of research, whereby a principled approach is to frame this as the problem of simulating a conditional denoising process which converges to the conditional of interest when fully denoised [68, 65, 12]. One approach is to leverage Bayes's rule, and combine pretrained score network, with an estimate of the conditional over observation given noised data reftrreeferre to as a guidance term [13, 48, 66, 65, 27]. Several refinements over the estimation of this guidance have been suggested [60, 19]. Another line of work has focused on 'inpainting' tasks paokless where the observations lie in subspace of the diffused space- and suggested 'freezing' the conditioning data when sampling [68, 72, 46, 47]. Lastly, another popular approach, sometimes referred to as elassifier-free guidance, is to directly learn the conditional score amortised over the observation of interest [26, 79, 71]. This approach requires having access at training time to pairs of data and observation.
Machine learning models for PDEs. One class of ML-based methods that can learn to approxi- mately solve PDEs given some initial state assumes fixed mimte-dimensiooal spatial discretisation (i.e. mesh), typically a regular grid, and, leveraging convolutional neural networks, learns to map initial values of the grid to future states [20, 83, 5J. These are data-driven methods, which are trained on trajectories generated by conventional solvers, typically at fine resolution but then downscaled.
operators via neural networks, coined neural operators [2, 36, 45, 51, 49, 6]. These methods are able to map parameters along with an initial state to set of neural network parameters that can then be used for simulation at different discretisation. Neural operators typically achieve this by leveraging the Fourier transform and working in the spectral domain [38, 21, 52]. Fourier neural operators (FNOs) can be used as the backbone of the diffusion model, but U-Net-based [59] models have been shown to outperform them [22].
Alternatively, mesh-free approaches learn
Classical solvers for PDEs. Conventional numerical methods, e.g. finite element methods [43], are commonly used solve dependent PDEs in practice [1]. In contrast to ML-based methods, numerical methods have the advantage of providing theoretical guarantees of convergence and bounds on error, which depend on time and space discretisation. However, many real-world problems, e.g. fluid and weather modelling, are described by complex time-dependent PDEs, and thus to obtain accurate and reliable solutions, the space and time discretisations have to very small [23]. As such, in this case, the conventional methods require significant computational and time resources. To tackle this issue, several recent works [33, 69] attempted to speed up numerical solvers using automatic differentiation and hardware accelerators (GPU/TPU).
While conventional numerical methods have been intensively studied and applied in forecasting, they are not suitable for solving inverse problems out-oo-the-box. box- To overcome this, several works [31, 41] propose numerical approaches specifically developed for certain PDEs, while Arridge et al. [3] provides an overview on data-driven, regularisation- based solutions to inverse problems. Moreover, unlike diffusion models, numerical solvers not offer probabilistic treatment, which crucial in certain applications of PDE modelling. For example, in weather and climate modelling, being able to predict range of probable weather scenarios key in decision-making, such as issuing public hazards warnings and planning renewable energy use [54].
20

---

128
128
128
0-
ÅT 50 ÂT Training example
100 AT
AT
50 AT
100 ÂT
50 ÄT Test example
100 ÀT
Validation example
Figure 4: True test trajectories from solving the Burgers' equation following the setup of [77]. In this case, both training, validation, and test trajectories have the same length. We used Ar = 0.01s, so the trajectory contains states from time Os 1s. The spatial dimension is discretized in 128 evenly distributed states
D Data generation and implementation details
D.1 Burgers' equation
Ground truth trajectories for the Burgers' equation [8]
du Ou 82u
+ u Oz 0.
obtained from the open-source dataset made available by Li et al. [40], where they consider (I, E (0, 1) x (0, 1j. The dataset consists of 1200 trajectories of 101 timesteps, i.e. Ar = 0.01, and where the spatial discretisation used 128. 800 trajectories were used as training set, 200 as validation set and the remaining 200 as test set. These trajectories were generated following the setup described in [77] using the code available sampled from Gaussian random field defined as (0, (25 + 521) (although be aware that in the paper they state to using (0, (- + 5 () which different from what the code using viscosity of 0.01. Specifically, they solve the equations using a spectral Fourier discretisation and fourth-order stiff time-stepping scheme (ETDRK4) [14] with step-size 10-4 using the MATLAB Chebfun package [16]. Examples of training, validation, and test trajectories are shown Fig. 4.
their public repository. Initial conditions are
suggests) and then they integrate the Burgers' equation using spectral methods up
D.2 1D Kuramoto-Sivashinsky
We generate the ground truth trajectories for the ID Kuramoto-Sivashinst equation
du Ou 8.u
8'u
+ u 0=
O,
θr
022
0z4
following the setup of Lippe et al. [42]. which is based the setup defined by Brandstetter et al. [7]. In contrast to Lippe et al. [42], we generate the data by keeping Ax and AT fixed, i.c. using a constant time and spatial discretisation'. The public repository of Brandstetter et al. [7] can be used to generate our dataset by using the same arguments as the one defined in Lippe et al. [42, App D1.].
Brandstetter et al. [7] solves the KS-equation with periodic boundaries using the method of lines, with spatial derivatives computed using the pseudo-spectral method. Initial conditions are sampled from a distribution defined over truncated Fourier series, i.e. un(x) >k= A₁ sin I. + Φx)). where Ag. lk, Φb random coefficients representing the amplitude of the different sin waves, the phase shift, and the space-dependent frequency. The viscosity parameter is set to v = 1. Data is initially generated with float64 precision and then converted to float32 for training the different models. We generated 1024 trajectories of 140Ar, where AT 0.2s, as training set, while the validation and test set both contain 128 trajectories of length 640Ar. The spatial domain discretised into 256 evenly spaced points. Trajectories used to train and evaluate the models can be seen in Fig. 5.
2This just requires changing lines 166-171 of the generate_date data. file in the public repository of Brandstetter al. [7].
21

---

256
256
256
128-
128
128
0 Ar
320 AT
640 AT 0 Ar
320 AT Test example
640 ÄT
Validation example
0
0 AT 70 AT 140 AT Training example
Figure 5: Examples from the than those used to evaluate the models. Training trajectories have 140 time steps generated every Ar 0.2s, i.e. training trajectories contain examples from Os to 28s. Validation and test trajectories, instead, show the evolution of the equation for 640 steps, i.e. from Os to 128s. The spatial dimension discretized in 256 evenly distributed states.
dataset. Training trajectories shorter
D.3 2D Kolmogorov-Flov equation
We use the setup described by [60] and their public repository to generate ground truth trajectories for the Kolmogorov- equation. Their implementation relies on the jax-cfd library [34] for solving the Navier-Stokes equation. As have seen in Section 5, the Kolmogorov equation is defined as
follows:
du Ōt
vA Vp f-0
+ u Vu-
(17)
V 1=
where y > is is the viscosity also defined as Rs with Re being the Reynolds number, p is the fluid density, p the pressure field, and is external forcing term. For the data generation they follow Kochkov et al. (34] setup, which used Re = 1000, or v = 0.001 respectively, constant density sin(4y) 0. lu where the first term be seen forcing term and the second one as a dragging term. They solve the equation on a 256 x 256 domain grid defined on [0, 2nr² with periodic boundaries conditions but then they coarsen the generated states into 64 64 resolution. While other works in the literature directly model the vorticity, i.c. the curl of the velocity field, they directly model the velocity field They simulate Eq. (17) starting from different initial states for 128 steps and they keep only the last 64 states. The training set we used contains 819 trajectories, the validation contains 102 trajectories, and we test all the different models on 50 test trajectories. We show in Fig. 6 some example states from the training, validation, and test sets, corresponding to time step E {Ar, 8AT, 16AT, 24A7, 32AT, 40Ar, 48Ar, 56AT, 64AT].
p= and
D.4 Implementation details
Our implementation is built on PyTorch [50] and the codebase will soon be made publicly available. For generating samples simulate the probabilistic ordinary differential equation (ODE) associated with the backward or denoising process which admits the same marginals as E (t), referred to as the probability flow
d=(t)= = (- - loogp, dt.
using the DPM solver [44] in 128 evenly spaced discretisation steps. As a final step, return the posterior mean over the noise-free data via Tweedie's formula. Unless specified otherwise, we set the guidance schedule to r? yo? /ue [19, 60], tune via grid-search and set o, 0.01.
In the following, we present all the hyperparameters used for the U-Net that parameterises the score function. For the forecasting task, we studied two different architectures for KS and Kolmogorov- one considered [60] (SDA) and another one in [42] (PDE-Refiner). Tab. 3 contains the architecture and training parameters for the SDA architecture, while Tab. 4 contains the parameters used for the PDE-Refiner architecture. For KS, the SDA and PDE--Refine architectures have very similar capacity in terms of the number of trainable parameters. The only difference is in the way they encode
22

---

t= AT
t= 162T
== 24AT
== 32AT
== 40AT
== 48AT
1 56AT
t == 64Ar
Training example
ÀT
8AT
16AT
24AT
32Ar
40Ar
48AT
56AT
644T
Validation example
t= AT
== 8AT
== 16AT
t= 24Ar
== 32Ar
== 40Ar
t 484r t 564T t 64AT
Test example
Figure Examples from the Kolmogorov dataset. Training, validation, and test trajectories have 64 time steps. For each example, we only show subset of the states, corresponding to (Ar, SAT, 16AT, 24Ar, 32AT, 40Ar, 48Ar, 36AT,64A7]
the diffusion time SDA encodes PDE-Refiner encodes into the parameters of the LayerNorm. For Kolmogorov, the architectures considered were different in terms of the number residual blocks per level and channels per level, resulting in the PDE-Reee Refiner architecture having more capacity than the SDA architecture.
separate channel that concatenated the input, while
The results in Figs. 1 and 2 correspond the best-performing architecture for each model, but show the full set of results Tab. 5. The results for the data assimilation experiments are all obtained using the SDA architecture.
Baselines implementation details. In the forecasting experiment of Sec. 5, we tested the different score-based models against other MI.-based models, including the sfate-of-the-art PDF-Refiner [42] and an MSE-trained U-Net and Fourier Neural Operator (FNO) [39]. For the implementation of the baselines, we used the following open-source code, all using the PDE-Refiner architecture from Tab. In addition, for the MSE-trained U-Net, we also implemented model using the SDA architecture, with results shown Tab. 5.
Hardware. All models and the experiments we present in this paper were used a GPU cluster with a mix of RTX 2080 Ti, RTX A5000, and Titan X GPUs. The majority of the GPUs used have 1 IBB of memory. A direct map between experiments and specific GPU was not tracked, therefore we cannot report these details. The results in Fig. are computed on a RTX 2080 Ti GPU with 1IGB of memory.
single GPU. We
E All-at-once (AA0) joint sampling and study of solvers in forecasting
One approach for generating long rollouts conditioned on some observations using score-based generative modelling is directly train joint posterior score st 21. (t)ly) approximate Vain(t) log (()):1(t)lv). The posterior can be composed from prior score and likelihood term, just as the main paper. However, this approach cannot be used for generating long rollouts (large L). as the score network becomes impractically large. Moreover, depending on the dimensionality of each state, with longer rollouts the probability of running into memory issues increases signifi- cantly. To deal with this issue, in this section we investigate the approach proposed by Rozet and Louppe [60], whereby the prior score (sp (1, (()) approximated with series of local scores (so (t, - (t))), which casier learn. The size of these local scores is in general significantly
23

---

Table Hyperparameters used for the SDA U-Net architecture inspired by [60] in the three datasets considered in the paper. Blue refers to specific information about the joint model and red refers to the universal amortised model.
BURGERS
KS
KOLMOGOROV
Residual blocks per level
(3,3,3)
(3,3.3.3)
(3.3.3) (96, 192, 384)
Channels per level
(64,128,256)
(64,128,256,102)
Kernel size
Padding
circular
circular
circular
Activation
SiLU
SiLU
SiLU
Normalization
LayerNorm
LayerNorm
LayerNorm
Optimizer
AdamW
AdamW
AdamW
Weight decay
le-3
le-3
1e-3
Learning rate
2e-4
2e-4
2e-4
Scheduler
linear
linear
linear
Epochs
102411024
4kl15k
10kl10k
Batch size
32
32
32
Table 4: Hyperparameters used for the PDE-Refiner U-Net architecture inspired by [42] in the forecasting experiment for the KS and Kolmogorov datasets. Blue refers to specific information about the joint model and red refers the universal amortised model.
KS
KOLMOGOROV
Residual blocks per level
(3,3,3,3)
(3,3,3,3)
Channels per level
(64.128,256.1024)
(128, 128, 256, 1024)
Kernel size Padding
3
circular
circular
Activation
GELU
GELU GroupNorm
Normalization
GroupNorm
Optimizer
AdamW
Adam W
Weight decay
1c-3
1e-3
Learning rate
2e-4
2e-4
Scheduler
linear
linear
Epochs
2k|156 5k
700110k
Batch size
128
128
smaller than the length of the generated trajectory and will be denoted as the window of the model W.
At sampling time, the entire trajectory is generated all-at-once, with the ability to condition on initial and/or intermediary states. We study two sampling strategies. The first one is based on the exponential integrator (HI) discretization scheme [82], while the second one uses the DPM++ solver [44]. For each strategy, the predictor steps can be followed by few corrector steps (Langevin Monte Carlo) [68, 47]. For the results below, we consider the forecasting scenario for the Burgers' dataset. We forecast trajectories of 101 states, conditioned on noisy version of the initial states (with standard deviation о, = 0.01). The results are computed based on 30 test trajectories, each with different initial conditions. The sampling procedure uses 128 diffusion steps and varying number of corrector steps.
Influence of solver and time scheduling. We show the results using two solvers: El [82] and DPM++ [44]. For each, we study two time scheduling settings, which determine the spacing between the time steps during the sampling process
Linear/uniform spacing: l, =1,N (IN (0) v
Quadratie:
(N
NN 1/k
N
24

---

where N indicates the number of diffusion steps, tn and to are chosen to be 1 and 10 10-3 and is chosen be
Fig. shows the RMSD between the generated samples and the true trajectories, averaged across the trajectory length. It indicates that for AAO sampling with 0 corrections, there is little difference between the two time schedules, and that DPM++ slightly outperforms EI in terms of RMSD. However, the differences are not significant when taking into account the error bars. As the number of corrections is increased, the difference in performance between the two solvers gets negligible. However, it does seem that with large number of corrections (25), the quadratic time schedule leads to lower RMSDD. Nevertheless, given that, in general, only 3 5 corrector steps are used, we argue that the choice of solver and time schedule has a small impact on overall performance.
Conrections:
Corrections: 5
Conrections: 25
DPM
EI
DPM
Hl
DPM
3-
RMSD
2-
quadiniic
Tinear
quadiaiic
linear
quadmatc
Tincor
Figure 7: RMSD for Burgers' for the solvers (DPM+ and El), and the time schedules (quadratic and linear) used. The results are shown for (left), 5 (middle), and 25 (right) corrector steps. When using no/only few corrector steps, the RMSD highly impacted by the choice of solver and time schedule. However, at 25 corrector steps, the quadratic time schedule seems to lead lower RMSD.
Influence of corrector steps and window size. Figs. and 9 illustrate the effect of the window size and number of corrector steps the mean squared error (MSE) and the correlation coefficient (as compared to the ground truth samples). Based on the findings from above, here we used the DPM++ solver with quadratic time discretisation.
The corrector steps seem to crucial to prevent the error accumulation, with performance generally improving with increasing corrections for both models (window 5 and 9). However, this comes at a significantly higher computational cost, as each correction requires one extra function evaluation. Unsurprisingly, the bigger model (window 9) is better at capturing time dependencies between states, it comes with increased memory requirements.
10-
MSG
00
60
100
20
40
100
Trajectory longth
Trajectory length
Figure 8: Evolution of the MSE for Burgers' along the trajectory length for window of 5 (left) and (right). The error bars indicate + standard errors. Each line corresponds to a different number of corrector steps. Unsurprisingly, the window 9 model performs better than the window 5 model (when using the same number corrector steps). Performing corrector steps, alongside predictor steps, is crucial for preventing eπrоrs from accumulating.
Finally, the findings from above are also confirmed Fig. 10, which illustrates that the predictive performance increases with the window size and the number of corrector steps. Both of these come at an increased computational cost; increasing the window size too much might lead to memory issues, while increasing the number of corrections has negative impact the sampling time.
25

---

1.0-
Corceution
0..0
-0.5
25
10
40
100
0
100
Trjectory length
Trajectory length
Figure 9: Evolution of the Pearson correlation coefficient for Burgers' along the trajectory length for a window of 5 (left) and 9 (right). The error bars indicate standard errors. The findings are entirely consistent with the ones suggested by the MSE results.
3.0
Window
2.5
RMSD
1.34
1.0-
0.0
Conrections
Figure 10: RMSD for Burgers' as function of corrector steps for the two studied models (window 5 and window 9). Increasing both the window size and the number of corrections positively impacts the predictive performance.
Additional experimental results on forecasting
This section shows further experimental results in the context of PDE forecasting. We show results on the Burgers', KS, and Kolmogorov datasets. We study the difference between all-at-once (AAO) and autoregressive (AR) sampling in forecasting, the influence of the guidance strength, and of the conditioning scenario in AR sampling for the joint and amortised models. Moreover, we also provide comparison with some widely used benchmarks in forecasting: climatology and persistence. Finally, we provide results for two U-Net architecture choices for the diffusion-basee models and the MSE-trained U-Net baseline.
F.1
AAO vs AR
Once the joint model se (t)) se Varn(1)) log (ogp(x1.1(tt) is trained, there are multiple ways in which it can be queried
All-at-once (AAO): sampling the entire trajectory in one go (potentially conditioned on some initial and/or intermediary states);
Autoregressively (AR): sampling a few states at a time, conditioned on a few previously generated states (as well potentially on some initial and/or intermediary states).
Let us denote the length of the generated trajectory with L. the number of predictor steps used during the sampling process with p. and the number of corrector steps with c. Moreover, in the case of the AR procedure, we have two extra hyperparameters: the predictive horizon P (the number of states generated at each iteration), and the number of conditioned states C. Finally, let the window of the model be W. Each sampling procedure has its advantages and disadvantages.
Sampling time: the AAO procedure (with no corrections) faster than the AR one due to its non-recursive nature. We express the computational cost associated with each procedure in terms of the number of function evaluations (NFEs) (i.e. number of calls of the neural network) needed generate the samples.
AAO: NFE = (1 + +)
AR: NFE- + c) x x L-W L-W
2. Memory cost: the AAO procedure more memory-intesiive than the AR one, as the entire generated trajectory length needs to fit in the memory when batching is not performed. If
26

---

batching is employed, then the NFEs (and hence, sampling time) increases. In the limit that only one local score evaluation fits in memory, the AAO and AR procedures require the same NFEs.
3. Redundancy of generated states: in the AR procedure, at each iteration the model re- generates the states conditioned upon, leading to computational redundancy.
We show that the AAO procedure is not well-suited for forecasting tasks, by performing experiments on the Burgers' and KS datasets.
Burgers'. As mentioned above, the AR method comes at an increased computational cost at sampling time, but generates samples of higher quality and manages maintain temporal coherency even during long rollouts. However, as illustrated in App. E, the quality of the AAO procedure can be improved by using multiple corrector steps. In particular, to obtain the same NFE for AAO and AR sampling (assuming the same number of predictor steps and no batching for AAO),, we can use corrector steps.
Fig. 11 shows the results for trajectory with L = 101, using model with W = = The results are expressed in terms of the RMSD, calculated based on 30 test trajectories. For AAO sampling = 48, while for AR sampling we set c = 0. In the AR sampling, we generated 2 states at a time conditioned on the previously generated states. In both cases we conditioned on noisy version of the initial 3 true states (with standard deviation dy 0.01). For all experiments, we used the El discretization scheme, with quadratic time schedule. The AR procedure is clearly superior the
set
RMSD
AALMS
AKN
Rollout/curretioos
Figure 11: RMSD for Burgers' for AAO and AR generation with window of 5. The number of corrections for AAO generation was chosen to match the computational budget the AR sampling technique, yet the AR technique is clearly superior.
AAO, even with 48 corrector steps. This can be easily observed by comparing the generated samples the ground truth shown in Fig. 12.
True
AAO
AR
Figure 12: Comparison between the true samples (top), AAO samples (middle), and AR samples (bottom) for Burgers'. The AAO samples are only able to capture the initial states correctly, while the AR AR-genertted samples maintain very good correlation with the ground truth even at long rollouts (101 time steps).
For window of 9, used the following number of corrector steps for each conditioning scenario:
27

---

2, 7: c 101-9 46
3=3,C= c: - =31
3
PP=40 c= 101-9 23
P 5, C- 4 - 19
Moreover, the number of states we initially condition upon is equal to C' (i.e., when comparing AAO with 46 corrections to AR with 2, = 7 (2) 7), we condition on the first initial noised up states).
RMSD
1.0-
0.0
3
23.4 19.0
Figure 13: RMSD for Burgers' for AAO and AR generation with a window of and with a varying computational budget. In the AAO case the computational budget is dictated by the number of corrector steps (indicated on the x-axis through the number before the comma). In the AR case, the budget is determined by the number of states generated at each iteration (P). This is also indicated on the x-axis in the format P C (e.g.. implies we generate two states at time and condition on 7). AR is clearly superior AAO for forecasting.
Once again, the AR- AR-generated trajectories achieve significantly better metrics, which also reflects in the quality of the generated samples. Fig. 14 shows examples for AAO: 46/AR: 2 7.
Irue
AR AAO
Figure 14: Comparison for Burgers' between the samples (top), AAO samples with 46 corrections (middle), and AR 2 samples (bottom). With 46 corrector steps, some AAO trajectories manage to maintain good correlation with some of the easier test trajectories. However, the quality of the AR samples is clearly superior to the AAO samples, managing to model all test trajectories well.
KS. We show similar findings on the KS dataset, where we present the performance of AAO and AR sampling 128 test trajectories of length 100At (corresponding to 20 seconds) in terms of MSE along the trajectory length. For AAO we employ a varying number corrections (0, 1, 3, and 5) and for each setting show the results for the best 7 E {0.03, 0.05, 0. 0.5). Setting to lower values led to unstable results. For AR we set y = 0.1 and show results for 1. We condition on the first true states and show the results for the remaining 92 states.
Fig. 15a clearly shows that AAO sampling has poor forecasting performance as compared to AR sampling. Employing corrections has positive impact close to the initial states, but they no longer seem to help for longer rollouts. Even we employ Langevin corrector steps, the AAO metrics are significantly worse than the AR metries.
28

---

Value
Error
800
600
ANO (1
MSE
AAO (0)
400
AAO (1)
AA)
200
(5)
A4
20
40
60
80
Trajectory length
(a) Evolution of the MSE for KS along the trajectory length for a window of 9. The bars indicate 3 standard errors. For AAO we show the results for 1,3, and corrector steps, indicated between brackets. Averaged over the entire trajectory, employing 1 correction gives best performance.
Time (in seconds)
(b) Example trajectories
from the window 9 model.
Figure 15: AAO vs AR comparison for the KS dataset.
Fig. 15b shows an example trajectory of length 100, alongside the predictions from AAO with 1 correction and AR sampling (left column). The right column shows the difference between the predictions and the ground truth. The AR procedure leads to results that are very similar to the ground truth for the first 100 states, whereas the predictions from AAO sampling quickly diverge from the ground truth.
Kolmogorov. We also confirm the previous findings in the Kolmogorov dataset, with experimental evidence shown in Fig. 16. We compare the performance between AAO and AR sampling for a model with window size of in terms of MSE along the trajectory length. We use 50 test trajectories of length 64. For AAO we show the results for 0 and 1 corrections, and show the results for the best setting for the guidance strength between E f0.03, 0.05, 0.1, 0.3], corresponding to y = 0.05 in both cases. For AR show results for and y 0.1. We condition on the first states and show the results for the remaining 60 states.
We plot in Fig. 16b the vorticity associated with the true, AAO-predicted, and AR- AR-predicted states eight different time points. The AAO-predicted states diverge from the ground truth quickly, while the AR-predict states stay close to the ground truth for more than half of the trajectory length.
F22 Influence of guidance strength in forecasting
KS. We studied three different settings for the parameter that controls guidance strength Y {0.01, 0.03, 0.1} } for the KS dataset for joint AR sampling. We found 7 0.01 to be unstable in the case of forecasting, so in Fig. 17 we show the results for E {0.03, 0.1). We consider a variety of window sizes W and conditioning scenarios P C. with P C W, where P represents the number of states generated at each rollout iteration, and C' the number of states we condition upon. As shown in Fig. 17, 0.1 is the better choice for low / middle-range values of P, and decreasing it to = 0.03 can lead to unstable results. In contrast, ? = 0.03 gives enhanced performance for longer predictive horizons P, compared
0.1.
Kolmogorov. We studied five different settings for the guidance strength 7 in the Kolmogorov dataset 7 {0.01, 0.03, 0.05, 0.1, 0.2, 0.3). The lowest value 0.01 gave unstable results, not showing Fig. 18. The findings are similar to those for the KS dataset, indicating that the guidance strength needs be tuned depending on the conditioning scenario. For P and P the best setting of y is 0.03. However, for smaller P, 7 0.03 gives unstable results, and the best settings are 7 0.05 for 2, and 0.1 for 1.
E33 Forecasting performance of differrntt conditioning scenarios P C
The autoregressive sampling method offers some flexibility in terms of the number of states generated at each iteration (P). From computational budget perspective, the higher P. the fewer iterations
29

---

AÃO (0)
MSE
1000
AAO (1)
AR
0
60
Trajectory length
(a) Evolution of the MSE for Kolmogorov along the trajectory length for model with window 5. The bars indicate + standard errors. For AAO we the results O, and corrector steps, indicated between
brackets. Averaged over the entire trajectory, employing corrections does not improve performance.
Ar == 9åT
== 17AT
25A7
t 33
== 41Ar
49Ar == 57
True
AAO (0)
AR
(b) The vorticity of the predicted states from the window 5 model. The top row shows the true vorticity, the middle one shows the vorticity of the AAO-predicted samples, and the bottom row, the vorticity of the AR- AR-predicted states. We indicate by the time index of the state (i.e. L= Ar means first forecast state).
Figure 16: AAO vs AR comparison for the Kolmogorov dataset.
20-
0.03
0.1
RMSD
25
3/7
67
10
1112
2111
3/10
5/8
PIC
100-
0.03
0.1
time (S)
80-
High
118
1/12
2111
3/10
PIC
Figure 17: Forecasting performance of the joint model the KS dataset for y E (0.03, 0.1} for varying window sizes and P C' conditioning scenarios. The top row shows the RMSD, while the bottom one shows the high correlation time. Where we do not show results for 7 0.03 it because they were unstable.
need to be performed. However, this might come the cost decrease in the quality of the generated samples. In this section, we investigate this trade-off for the Burgers' and KS datasets. The results for the Kolmogorov dataset are shown in the main paper in Fig.
We present results for a window of 9 in the Burgers' dataset and study several window sizes W for the KS dataset. We also vary the predictive horizons P. and the number of conditioning frames C W P. The trajectories are conditioned upon the first W P true states.
The results are presented in terms of RMSD, which is computed based on 200 test trajectories for the Burgers dataset, and on 128 test trajectories for the KS dataset. The number diffusion steps used 128. For all experiments, used the DPM++ solver, with a linear time discretisation scheme. At the
30

---

0.01
0.
0.3
RMSD
2/3
3/2
411
PIC
High corr. time (5)
10-
0.03
01
0.3
0.05
1/4
23
3/2
P PIC
Figure 18: Forecasting performance of the joint model in the Kolmogorov dataset for y (0.03, 0.05, 0.1, 0.2, 0.33 for varying window sizes and P C conditioning scenarios. The top row shows the RMSD, while the bottom one shows the high correlation time.
end, we also performed an additional denoising step, and adopted the dynamic thresholding method for the Burgers' dataset. The guidance strength is fixed at y 0.1 for the Burgers' dataset, while for KS we show the best setting out of y E {0.03, 0.13.
Burgers'. Fig. 19 shows that the performance is not very sensitive to P for low to medium values of P. although the model seems to perform better with lower P (and consequently, higher number of conditioning frames C). However, the performance quickly deteriorates for very high values of P (i.e., 8).
x10
-10-2
RMSD
1.0-
RMSD
0.54
2
3
5
0
7/2
217
3
4
514
PIC
Figure 19: Burgers' dataset. RMSD for different values P. The right plot a zoomed in version of the left one, where we excluded the 8 setting. Outliers have been removed.
KS We repeat the same experiment for the KS dataset, but study a wider variety of window sizes W and conditioning scenarios P C, where C W. Moreover, as mentioned G.I, for each setting we test two values of the guidance parameter yy In Fig. 20 only show the results corresponding the best guidance strength. We also show results for the amortised model introduced in Sec. 3.2 in Fig. 21, where different model was trained for each P C scenario. Unlike the joint model, which was trained for 4000 epochs, the amortised model was only trained for 1024 epochs, since after that it started overfitting.
RMSD
x PIC
time (3)
High
PIC
Figure 20: RMSD (top) and high correlation time (bottom) for the joint model on the KS dataset for varying window sizes W and P C' conditioning scenarios (where P + == WW. The model performs best for low values of P (P = 1), and high values of the number of conditioning frames
31

---

Joint model. As observed in Fig. 20, the best conditioning scenario for the joint model for each window corresponds to 1 (i.e. low number of predicted frames). We observe an increase in performance by increasing the window size, with the best performing setting corresponding to model with window 13 12). In this case the model achieves 77.3s high correlation time, as opposed to 73.2s for the window 9 model presented in the main paper. We believe further gains in performance could be achieved by increasing the window size further, but this would also come with higher memory requirements.
30+
RMSD
PEE
1.1
di o1 68
80 s4 si8
IhE 1214 1218
PIC
tiinl (5)
High
25-
214 2/8
611
s12
sia
12
1212
124 12
P
Figure 21: RMSD (top) and high correlation time (bottom) for the amortised model in the KS dataset with varying P and C. The model performs best when generating relatively large blocks (P 2 4) and conditioning lower number of frames (C < 4).
Amortised model. In contrast to the joint model, Fig. 21 shows that the amortised model under- performs for small P 1, with high correlation time of less than 50s. We hypothesise that the reason for this might be the combination of the following: i) generating larger number of frames at the same time requires less steps for the same trajectory length, so there is less error accumulation; and ii) for larger P the model is trained with a better learning task compared to P = 1, as larger P provides information to the model. For P 4, the model performs similarly for any P. being slightly more sensitive to larger C' when P is lower. Regardless of P, the performance of the model deteriorates significantly with larger C 8, similar to the behaviour shown Lippe et [42].
Universal amortised model. As indicated in Sec. 3.2, more flexible alternative to the amortised model is the universal amortised model. Due to our novel training procedure, latter does not need re-training for each P C conditioning scenario. Moreover, shown in Fig. I and Fig. 22, the performance of the universal amortised model stable for all C scenarios, and, unlike other models from the literature [42], is not negatively affected by longer previous history.
30-
30-
RMSD
2014
1/3
20
3/6
5/4
1/3
3/0
45
3/c
P
Figure 22: RMSD for universal amortised model (left) and amortised model (right) in the KS dataset with varying and C. Universal amortised model more stable as well as outperforms the amortised model across all combinations of P and C.
F.4 Comparison with climatology and persistence
We also compare the performance of our models (joint AR and universal amortised) to two common forecasting benchnakss-ciimtoolg.. determined by computing the average value over the entire training set, and persistence, obtained by repeating the last observed state for the entire trajectory. Rather than looking at RMSD or high correlation time, we look at a per time step triterion-MSS
32

---

This allows us to find out at which point predicting just the mean value is preferred to the predictions of the models (i.e. when climatology gives lower MSE than the models).
Fig. 23 shows that the joint and universal amortised models outperform persistence throughout the trajectory length. For KS, they significantly outperform climatology until approximately 400 time steps, corresponding to about 80s. As observed in the main paper through the high correlation time, the amortised model seems be better long-term prediction than the joint model. For Kolmogorov, they outperform climatology for more than 40 time steps, corresponding to 8s.
X
X10
1.0
Guided 1|8
Guided 1|4
Universal
Universal
MSE
amortised 8
MSE
amortised 1
Persistence
0.5
Persistence
Climatology
Climatology
0.0-
0.0
400
600
Trajectory length
Trajectory length
(a) KS.
(b) Kolmogorov.
Figure 23: MSE along the trajectory length for KS (left) and Kolmogorov (right) for 1) the joint AR model (blue), 2) the universal amortised model (orange), and the two forecasting benchmarks considered- -3) persistence (black) and 4) climatology (grey). The error bars indicate + 3 standard error. Our models outperform climatology for about two thirds of the trajectory length.
F.5 Example samples
We show in Fig. 24 some example predicted trajectories using the joint 1 8 model and the universal amortised 8 model for the KS dataset, alongside the ground truth trajectories. We also show the error between the predictions and the trajectories.
Value
Error
Value
Error
Truc
True
Joint
Joint
Amortised
Amortised
120
120
Time (in seconds)
Time (in seconds)
Value
Error
Value
Error
Truc
True
Joint
Joint
Amortised
Amortised
80
120
40
120
120
40
80
120
Time (in seconds)
Time (in seconds)
Figure 24: Example true trajectories, and model predictions for the joint 1 8 (middle) and universal amortised 8 (bottom) model on the KS dataset. The top left/right example corresponds to one of the trajectories with the highest/lowest RMSD.
33

---

For Kolmogorov, we show in Fig. 25 examples of the vorticity at different time steps t corresponding the states (top), the states predicted by the joint 4 model (middle), and to the states predicted by the universal amortised model (bottom).
Ar
== BAT
t= 17A7
2527
334
41Ar
== 49A
57Ar
True
Joics
Amcribsed
A-Ar
9AT
1727
25AT
33AF
41A7
49A
57A
Joine
Amonised
T-Ar
9AT
1727
25AT
33AT
41AT
49A
57AT
Joint
Amortised
A=
1747
33Ar
True
Joint
Amortited
Figure 25: Example vorticity fields at different time steps and for four different initial states. Each plot consists of three different rows corresponding to the true trajectories (top), to the trajectories predicted by the joint 4 model (middle), and to those predicted by the universal amortised model (bottom).
34

---

F.6 Results for different architectures
For the diffusion models and the MSE-trained U-Net baseline outlined in Sec. 5, we study two different U-Net architectures and report the best results in Fig. 2. One follows [60] (SDA), while the other one follows [42] (PDE-Refiner). More details about the exact hyperparameters and training details can be found in App. D.4.
As illustrated in Tab. 5, obtain better results with the SDA architecture for the joint AR model, while the DDE--eeinnrr architecture yields better performance for the amortised models. For the MSE-trained U-Net baseline, the SDA architecture gives better results for the KS dataset, but worse for Kolmogorov.
Table 5: The best high correlation time (mean + 3 standard errors over test trajectories) achieved for two different architectures-SDDe inspired by [60], and PDE-Refiner inspired by [42]. For the diffusion-basse based models, the conditioning scenario that corresponds to the best performance indicated between brackets.
KS
KOLMOGOROV
SDA [60]
PDE- Ref-nef [42]
SDA [60]
PDE RefiRReiner [42]
JOINT AR UNIV. AMORTISED
73.2 + 4.3s (1
8)
70.9 + 4.2x (1
8)
6.2+0.66 (1 4)
6.8 + 0.5s (1 [4)
82.0 22.0 4.4s (8
1)
88.7 15.7 4.4s (7
2)
8.9 0.6s (1 4)
10.0 0.5s (1 4)
AMORTISED
63.1 3.8s (7
4.4s
L)
10.0 0.5s (1
10.6 0.5s (2 3)
MSE U-NET
91.9 4.9s
+
7.6s
10.0 + 0.3s
G Additional experimental results on offline data assimilation
In this section, we provide more details about the offline DA setup. For KS, for each sparsity level we uniformly sample the indices associated with the observed variables in space-time. For Kolmogorov, we uniformly sample indices in space each time step. For both datasets, we assume to always observe some full initial states. The number of such states for AR sampling depends on the chosen conditioning scenario and is equal to C. For AAO we condition on W initial states, where W stands for the window size.
Furthermore, this section provides further experimental results for the offline data assimilation task for the KS and Kolmogorov datasets. We did not explore this for the Burgers' dataset since we decided focus PDEs with more challenging dynamics. We investigate the influence of the guidance strength, and the performance of the AR models with varying conditioning scenarios. We also perform a comparison between joint model with Markov order of that is queried autoregressively and a 2k-Markov order model queried all-at-once. This allows us to gain further insight into the differences between the two sampling schemes. Finally, we provide quantitative and qualitative comparison to simple interpolation baseline.
G.1 Influence the guidance strength for the joint model
We study several settings of the guidance parameter for 1) the joint model sampled AAO, and 2) the joint model sampled autoregressively. We show the results in terms of RMSD evaluated on 50 test trajectories for varying levels of sparsity, in the main paper.
Joint AAO. For the KS dataset, we consider (0.01,0.00 0.03, 0.05, 0.1, 0.5). For Kolmogorov, we show the results for a window 5 model and Y {0.03, 0.05, 0.13. We query them AAO, and employ 0 and 1 corrections (AAO (0) and AAO (1)).
window 9 model, with ?
For KS, using 7 = 0.01 results in unstable rollouts, therefore we only report the results for the other four values of 7. For the sparse observation regimes, we find that the best results are achieved by employing the lowest value that still leads to stable results (n 0.03). For the denser regimes (i.e. proportion observed 2 10-1 for 0 corrections and > 10 1.5 for 1 correction), the best setting corresponds to 7 0.05.
Fig. 27 shows the results on the Kolmogorov dataset. In contrast to the KS dataset, we found that the best 7 value is consistent among all observation regimes, and corresponds to 7 0.05 for both AAO
35

---

Window 9, AAO (0)
Window 9, AAO (1)
0.03
20
7 0.05
RMSD
0.03
7 0.1
10-
0.05
y 0.5
10
10
10
10-3
10
Proportion observed
Figure 26: RMSD on the KS dataset for the joint model queried AAO with 0 (AAO (0) left) and 1 (AAO (1) right) corrections. We show results for a window of size 9. Confidence intervals indicate 3 standard errors.
(0) and AAO (1). this case, decreasing the y value (i.e. increasing the guidance strength) too much does not lead to improved performance and leads much higher standard error, compared the other studied settings.
Window 5, AAO (0)
Window 5, AAO (1)
YY = 0.03
Y 0.03
30-
0.05
0.05
RMSD
20
0.1
0.1
10
10 3
10 2
10
10
111
10
2
10
Proportion observed
Figure 27: RMSD on the Kolmogorov dataset for the joint model queried AAO with 0 (AAO (0) left) and 1 (AAO (1) right) corrections. We show results for window of size 5. Confidence intervals indicate standard errors.
For the joint AR scheme not only can we vary the guidance strength, but we can also change the conditioning scenario P C. In Fig. 28 show the RMSD for the KS dataset for two values of y E {0.03, 0.1} for a window model. Each plot shows the results for a different sparsity level.
Joint AR.
Proportion observed: 10 3.0
Proportion observed: 10 2.5
Proportion observed: 10 20
Proportion observed: 10 1.5
RMSD
Proportion observed: 10-10
Proportion observed: 10 10-0.
12-
1/8
2/7
P C
Figure 28: RMSD for the KS dataset on the offline DA task for varying conditioning scenarios and y 0.03 (light blue) and 0.1 (dark blue). For each sparsity level (proportion observed), we consider all the P C scenarios for model with window 9. When results are not presented for 7 0.03, it is because they were unstable. Error bars indicate + 3 standard errors. Note the logarithmic y-scale for the bottom two rows of plots.
36

---

Fig. 28 shows that for the lower values of P. i.e. settings where we predict only few frames at time, the results become unstable for low values y. However, especially the sparser regimes. decreasing to 0.03 can lead to significant improvement in performance for the larger values of P, i.e. settings with a higher predictive horizon. We also tried decreasing Y further 0.01 but the results were unstable for the majority of P C' settings. The only scenario where 0.01 achieves a better performance than the settings presented in Fig. 28 for and proportion observed 10-1.5, and this is the scenario considered in Fig. 3 for the above-mentioned proportion of observed data.
For Kolmogorov, we studied y E {0.03, 0.05, 0.1, 0.2, 0.3} for each conditioning scenario of window 5 model (P E {1,2, 3, 4)). }) We show the best setting for each conditioning setting in Tab. 6. Similarly to the KS dataset, we find that lower y values give better results for larger P. while they are unstable for lower P.
Table 6: Best y corresponding to each conditioning scenario P for different proportion of data observed on the Kolmogorov dataset for the joint model with window size 5.
Proportion observed best
best y
0.32
0.03
0.1
0.03
0,032
4
0.03
0.01
0.03
0.0032
0.1
0.001
0.2
G.1.1 Conditioning for DA in the universal amortised model
In our experiments, we found that to get better performance for the universal amortised model with reconstruction guidance, both y and By need to be tuned given a particular proportion of data observed. What is more, the universal amortised model is also more sensitive to the choice of the reconstruction guidance parameters than the joint one. Due to computation constraints and sensitivity of the universal amortised model to the choice of hyperparameters in reconstruction guidance, limit our experiments to a model with window 9 for the KS dataset.
For each proportion of observed data we perform a hyperparameter sweep varying (0.1, 0.01} and By E {0.01, 0.005, 0.001}. Tab. shows the value for the best predictive horizon P, y and Fy for each proportion of data observed for model with window size of 9 on the KS dataset, and Tab. 8 shows the results for the Kolmogorov dataset for a model with window size Empirically, we found that the more data are observed, the lower values of and oy required for optimal performance. Lower values of y and oy increase the strength of the guidance term, which indeed is expected to guide the process better when more data are available.
As described in the Sec. 3.2, for the universal amortised model we propose to use conditioning via architecture for past states and reconstruction guidance for sparse observations DA. However, there few other options that could be used for incorporating sparse observations in amortised models: i) resampling [46], where at each reverse diffusion time step more sampling iterations are introduced, which mix observed and unobserved variables, and ii) conditioning via architecture, which requires separate model, taking not only previous states, but also sparse observations as input, to be trained.
Table Best y oy and P for different proportions of data observed on the KS dataset for universal amortised model with window size 9.
Proportion observed
best
best y
best Ty
0.32
0.05
0.001
0.1
0.05
0.001
0.032
0.1
0.001
0.01
0.005
0.0032
0.1
0.01
0.001
0.1
0.01
37

---

Table Best 7, By and P for different proportions of data observed on the Kolmogorov dataset for universal amortised model with window size 5.
Proportion observed
best
best y
best By
0.32
0.05
0.001
0.1
0.05
0.001
0.032
2
0,05
0.001
0.01
0.05
0.001
0.0032
0.1
0.01
0.001
0.1
0.01
We have not investigated the resampling strategy, as in their work [15] showed that it has worse performance compared to reconstruction guidance in image and protein design applications. Moreover, in practice resampling requires around 10 mixing sampling steps, with one step accounting for single call to the score network, for each reverse diffusion iteration. Therefore, this makes the conditional sampling with resampling approach much more compute-intense and less effective in terms of sampling time.
We have studied the second approach for conditioning, where all the conditioning is happening via the architecture, For this, we trained separate model, which takes in fully observed previous states and some random data observations with masks as an input to the score network. During training we sample the number of fully observed states C' as well as random masks for sparse observations within the rest of the P frames for each batch. As a result, the model is trained to fit the following Fig. 29a shows the performance of this model (green) for different proportions of observed data. We considered a larger proportion of observed data, as for the lower ones considered in the main text of the paper, the model was significantly underperforming. From the Fig. 29a, it can be clearly seen that the model with all the conditioning done through the architecture performs significantly worse compared to our proposed version of the universal amortised model with reconstruction guidance. More specifically, the level of the universal amortised model's s performance with 10% of data observed is reached by the other model only when more than 70% of data is observed. We hypothesise that such poor model performance can be explained by the fact that at some point during the AR rollout, the input condition (previously generated C' states and partially observed sparse data within window of size P) that is provided to the model becomes duitr--diittiuutiin due to the mismatch between the previously generated states and the partially observed variable from the true trajectory. Fig. 29b shows inputs to the score network depending on the AR step on KS data. The model was trained with window 9 and uses 3 previously generated states for conditioning. The mismatch between previously generated states and observed trajectory data is visually noticeable at AR step 60, while becoming even more drastic with more AR steps. Therefore, there a mismatch between input distribution used during training and sampling.
conditional score: sp(t, -2-C++++ (t)---sst(()tC::::: Tt-Cx;,
Univonal asautize
staps
Coalitining orchiettere
50
100
150
200
250
RMSD
10
0
AR siep: 60
10-1
50
200
250
step: 100
0.1
0.2
0.3
Proportion observed
0.4
0.6
0.8
0.7
0.8
0.9
50 100
150
200
250
(a) Conditioning via architecture
(b) Input to the model with conditioning via architecture
Figure 29: Left: RMSD for the KS dataset on offline DA for the universal amortised model with all conditioning done via architecture. The model is compared against an identical universal amortised model that, in contrast, uses reconstruction guidance for conditioning on sparse observations. Right: The input that is used for model conditioning depending on the AR step (the longer the trajectory, the more AR steps are performed). The red line divides the previously generated C' states used for conditioning and the observations of the true trajectory within window P. The observations of the true trajectory are not fully observed by the model, but are illustrated without masks for more clarity. The vertical axis stands for the state within the window size used in the score network (9 in this case), while the horizontal one illustrates space discretisation of the PDE.
38

---

G.2 Performance for varying window sizes
Joint model. We compare the performance of a joint model with window 5 (k 2) with that of a model with window (k = 4) on the KS dataset. For both models, we show results for AR sampling, as well as AAO sampling with 0, and 1 corrections (AAO (0) and AAO (1)) in Fig. 30. The window model is represented with dashed lines, while the window model with solid lines. For each proportion of data observed, we show the results corresponding to a value of y and P C' for AR sampling that gave the best performance.
10
AR
RMSD
10-1_
AAD (0)
AAO (1)
10
10
10
Proportion observed
Figure 30: RMSD for the KS dataset on the offline DA task for two models, one with a window of size 5 (dashed lines), and one with window size of (solid lines). The left plot shows the results for AR sampling, while the right plot shows the results for AAO sampling with 0 and 1 corrections. The error bars indicate + standard errors. Note the logarithmic scale.
Fig. 30 shows that for both AAO and AR sampling the performance of the window 9 model is generally higher than that of the window 5 model, although for sparsity settings the difference not significant. The biggest improvement is seen for the middle range of sparsity (proportion of observed data between 10-2 and 10-1). We hypothesise that the reason why do not a significant improvement in performance between the two models in some settings is because Markov order of k: 2 (corresponding to window 5 model) is enough deal with the task hand, and increasing the Markov order offers limited additional useful information, leading to only marginal improvements.
G.2.1 Comparison between AR (k)) and AAO (2k)
As mentioned in Sec. 3.3, we hypothesise there are two reasons why the AR approach outperforms the AAO approach. The first reason has to do with modelling capacity- window of size 2 1, the effective Markov order of the AR model is 2k, while the AAO model has an effective Markov order of k. Secondly, in AAO sampling information needs to be propagated from the start of the sequence to its end, in order for the trajectory to be coherent. However, the score network only has limited receptive field 2k + 1, limiting information propagation. The way to tackle this through Langevin corrector steps, but they come an additional computational cost. To empirically assess these claims the context of DA, we compare window 5 model (k: 2) queried autoregressively with window 9 model (k 4) sampled AAO. For the latter, apply 0, 1. and corrector steps (AAO (0), AAO (1), and AAO (2)).
10a
AR, k=2
20-
AAO (0), k=4
RMSD
10
AAO (1), ked
10-
AAO ke4
10-1
0-
10
10
10
10
10
Proportion observed
Figure 31: RMSD for the KS dataset on the offline DA task for two models, one with a window of size 5 queried autoregressively (dashed blue), and one with window size of and sampled AAO (solid purple). We show results for 0, and corrector steps, indicated with the in AAO (c). The error bars indicate + standard errors. The right plot shows the same results as the left one, but logarithmic y scale.
39

---

The findings shown in Fig. 31 are consistent with our theoretical justification. In general, increasing the number of corrections (and hence, enhancing information propagation) increases the performance of AAO. Indeed, for the majority of sparsity levels (proportion observed > 10 10-2.5), the performance of the AR l: 2 and AAO
with least correction become similar.
The Langevin corrector steps do not seem to make a significant difference for the two extremes considered-prooorrtio observed of 10-3 and 10-0.5 In the first case, there are so few data observed, that the task becomes similar to forecasting. We have already shown that AAO is inefficient in that setting, regardless of whether corrections are applied or not (see Fig. 15). For a proportion observed of 10 0.5 the conditioning information is dense enough that corrections are not needed.
We can also corroborate the findings from App. F.I regarding the performance of AAO vs. AR on forecasting with the findings from this section on DA to better understand the difference between AAO and AR sampling. The results both tasks clearly show that window 5 AR models (k 2) does not lead to the same results as a window 9 AAO (0) model (I: 4), although they have the same effective Markov order. In the context of DA, the performance of the AR model can be recovered by applying at least one corrector step, which allows for more efficient information propagation. In contrast, in forecasting, applying corrector steps helps in the vicinity of the observed states, but makes little difference for long rollouts. The lack of observations (after the initial states) leads to inefficient information flow, which results in the AAO procedure failing to produce coherent trajectories. For AR, this is naturally tackled through the autoregressive fashion of the rollout.
G.2.2 Performance for varying conditioning scenarios P
We study varying conditioning scenarios for the offline DA task. For the KS dataset consider two joint models of different window sizes (W and = 9) and show in Fig. 32 the performance for each P C' scenario (indicated in the legend), where P C = W. For each proportion of observed data and for each conditioning scenario P C, we show the results corresponding to the best guidance strength 7
Window
Window 9
10a-
5
RMSD
111 4
3
3
7 2
10-1_
10
10
10
10
10
Proportion observed
Figure 32: RMSD for the KS dataset the offline DA task for two models, with a window of size 5 (left), and one with window size of 9 (right) for varying P C scenarios (indicated in the legend). Confidence intervals indicate 3 standard errors. Note the logarithmic y scale.
We observe in Fig. 32 that for both window sizes the best performance is generally achieved for the larger values of P. This is probably due the increasing amount of observations within the generated frames. However, for a proportion of observed data of 10-3, the observations are so sparse, that the task becomes closer to forecasting, where, as previously noted, lower value of P is preferred. In that setting, the best performance is achieved by the 4 and 1 scenarios, respectively, although there are only small differences between the different scenarios.
In the main paper, in Fig. 3, we show the figures corresponding to the best P C' scenario for each proportion of observed data, with the exception of 10-3. There, found little difference between the performance of / {1, 2, 3, 4). and hence chose to show the scenario that offers the best trade-off between performance and computational complexity (i.e. P 4) increase in RMSD of 0.25(= 1.4%) for an almost four-fold decrease in sampling time.
In Fig. 33a we show the performance of the amortised model with window 9 for different P and C' for KS. The best performance is achieved by the 3 model across most the proportions of data observed. The universal amortised model is again more stable compared to the joint one when varying and C, however struggles to get the same level of performance for intermediate values of observed data considered. Fig. 33b shows the same for the Kolmogorov dataset. In contrast to KS,
40

---

the best result is achieved by the model that has the longest history (1 4). The performance of the model again more stable for different scenarios compared to the Joint AR.
2
10'
RMSD
10"
RMSD
3
4
100
10
10
Proportion observed
10"2
10
Proportion observed
10
(a) KS.
(b) Kolmogorov.
Figure 33: RMSD for the KS (left) and Kolmogorov (right) datasets on the offline DA task for the universal amortised model for varying P C scenarios (indicated in the legend). Confidence intervals indicate = standard errors. Note the logarithmic scale.
We performed the same analysis for the Kolmogorov dataset and show the results in Fig. 34, In this case we only study model of window for both the joint and universal amortised models. We observe similar findings as in the KS dataset for the joint model, with lower P values being preferred for lower proportions of observed data, and higher when the amount of observations increases.
Window
10
RMSD
2
10
10
10
10
Proportion observed
Figure 34: RMSD for the Kolmogorov dataset on the offline DA task for the joint model with window size 5 for varying P C scenarios (indicated in the legend). Confidence intervals indicate = 3 standard errors. Note the logarithmic y scale.
G.3 Comparison to interpolation
In this section we perform a quantitative and qualitative comparison to a simple interpolation baseline. This implemented using the interpolation functionality within the Scipy library [74]. For KS, we investigated three interpolation methods: linear, cubic, and nearest, and report the results for the method that gave the best results for each percentage of observed data. For Kolmogorov, investigated linear and nearest interpolation, since the cubic method is not applicable to 3D data (2 spatial dimensions and time dimension).
We present the results Fig. 35. For the 1D KS dataset, all the methods studied in the main text (joint AR, AAO, and universal amortised) generally outperform the interpolation baseline for the sparse regimes (up to 10 proportion observed). The only exception for universal amortised with 10 data observed, where the performance is par. As more data is observed, the performance of all five methods becomes similar, indicating that the field can be easily reconstructed based on the observations alone. For the sparsest regimes, linear interpolation gave the best results, while for the last four, we used cubic interpolation.
For Kolmogorov, there is significant gap between the AR methods and the interpolation baseline, potentially indicating that interpolation is not as efficient when the number of data dimensions grows. The AAO methods also significantly outperform the interpolation baseline for all but the sparsest regimes. All results correspond to linear interpolation, since this gave the best performance.
For both datasets, we also provide a qualitative comparison between the joint AR predictions and the interpolation predictions, alongside the ground truth and the observed data. This is illustrated for two levels of sparsity for each dataset (10-2 and 10-1). Fig. 36 shows the results for the KS dataset. For the smaller proportion of data observed, the interpolation baseline produces unphysical behaviour, possibly because it does not contain any prior knowledge about the dynamics of the dataset and
41

---

KS
Kolmogorov
Joint AR
Joint AAO (0)
RMSD
20
Joint AAO (1)
10-
Universal amortised
Interpolation
10
10
10
10
10
10
Proportion observed
Figure 35: RMSD (mean =3 standard errors) for the KS (left) and Kolmogorov (right) datasets on offline DA task for varying sparsity levels. We show the results in Fig. 3, alongside interpolation baseline (red). Note the logarithmic x scale.
the observations alone do not offer enough information for good trajectory reconstruction. In the case of the joint AR model, the learned dynamics from the prior score make the predictions more consistent with the dynamics describing the KS dataset. When the data is dense enough, even simple interpolation baseline gives reasonable predictions (Fig. 36b).
Tue
Observed
Thue
Observed
Predicted
Interpolated
Predicted
Interpolated
Proportion observed 10
(b) Proportion observed 10
Figure 36: Qualitative comparison between the true, predicted (with the joint AR model), and interpolated trajectories, as well as the observed values for the KS dataset. The left plot corresponds to proportion of observed data of 10-2, while the right corresponds to 10-1
In contrast, for the Kolmogorov dataset, the interpolation results still show unphysical behaviour even for the higher percentage of data observed (10-1). This is probably because of the higher dimensionality of the data, as well as because we are plotting the vorticity (curl velocity field), rather than the predicted velocity fields, which involves differentiating the predictions. Overall, the visual quality the predictions outputted by the methods studied in the main text are clearly superior to the interpolation baseline.
H Additional experimental results online data assimilation
In this section we provide a more detailed explanation of the setup from the online DA task and provide results for several conditioning scenarios for the joint and amortised models.
The online DA scenario goes as described in Sec. 5.3. In Fig. 38 we also present a diagram, showing example of which observations would be available at each DA step in the KS dataset. The darker dots indicate observations, while the lighter background indicates that no observations are available in that region that DA step. The example corresponds
proportion of data observed of 0.1.
We study varying conditioning scenarios for the online DA task for the joint and universal amortised models. We do not show results for P = 1 due to the excessive computational time required for sampling, and only study the results for one window size for each dataset: window 9 for KS and window 5 for Kolmogorov.
Joint model. For the window joint model, we show the results on the KS dataset corresponding to the best guidance strength out of y E f0.03, 0.05, 0.1} for AR sampling. For AAO sampling we studied E (0.01, 0.03, 0.05, 0.1, 0.53 with and corrections. We found that the results did not significantly vary with with = 0.01 giving marginally better performance, and lower values giving unstable results. Tab. 2 we show the results corresponding y = 0.01 for AAO with (AA0 (0)) and 1 (AAO (1)) corrections, and for 0.05 and 4 5 for AR.
42

---

t= At
t= 9At
t= 174t
25
== 334
t= 41
494
t 57
True
Observed
Predicted
Interpolated
(a) Proportion observed= 10-2
9At
174t
25
33At
t
41
49
57
True
Observed
Interpolated Predicted
(b) Proportion observed = 10-1
Figure 37: Qualitative comparison between the true, predicted (with the joint AR model), and interpolated trajectories, as well as the observed values for the Kolmogorov dataset. The top plot corresponds to proportion of observed data of 10-2, while the bottom corresponds to 10-1
Fig. 39a shows that the joint model has similar performance for P E (4,5, 6). with the lowest RMSD given by 3. Lower values P E (2,3} give poorer performance, just as the offline DA task. Higher values of P (7,8) also perform worse, just in the forecasting task. Thus, the fact that the middle values of P perform best probably due to the fact that the online DA task is a mix between DA and forecasting.
For the Kolmogorov dataset we studied 7 E (0.03, 0.05, 0.1, 0.3) for model with window queried AAO. We employed 0 (AA0 (0)) and 1 (AAO (1)) corrections. Similarly to the KS dataset, we found little difference between the 7 settings, and found that 7 = 0.03 gave marginally better performance. For AR, we studied y E (0.03, 0.05, 0.1, 0.2) for each possible value of P (P e 2, 3, 4)). We found 7 = 0.03 to be unstable for low values P E {1, 2). The best setting was achieved for and 7 0.05. We show the results corresponding the best y setting in Fig. 39b.
Universal amortised model. For the universal amortised model, we use the best 7 = 0.05 and Ty 0.001 parameters found for the offline DA for the same proportion of observed data, due to
43

---

DA Step
DA Step
DA Step
Figure 38: Setup for the DA online task for the KS dataset for a proportion of observed data of 0.1. Every time steps, we get access to sparse observations from blocks of 20 states. At the first DA step predict the first 400 states based the sparse observations from the first 20 states. At the next DA step, we use the previously -ennuuyygeeee forecast, as well as the new sparse observations coming from the next block (time step 20-40) to predict the next trajectory of length 400 (from time step 20 to 420). This repeats until the end of the sequence (640 time steps).
RMSD
20
RMSD
20-
6/3
7/2
8/1
2/3
3/2
P PIC
PIC
(a) KS.
(b) Kolmogorov.
Figure 39: RMSD for the joint model on the KS (left) and Kolmogorov (right) datasets for varying P C' scenarios. For KS we use window 9 model, whereas for Kolmogorov we use a window model.
computational constraints. Fig. 40a shows the performance on the KS dataset for each scenario, where P + C = W. Fig. 40a indicates that the universal amortised model shows stable performance different and C, with the setting of 5 being slightly better than the other ones. Fig. 40b shows the performance on the Kolmogorov dataset for each P C' scenario, where P + C = W. Fig. 40b suggests that conditioning the model on larger history results in better performance. Even though 1 |4 shows the best result, all other conditioning scenarios also have competitive performance compared to both AAO(I) and Joint AR models.
RMSD
20
I
RMSD
10-
2/7
3/6
4/5
5/4
6/3
7/2
8
1/4
2/3
3/2
PIC
C
PIC
(a) KS.
(b) Kolmogorov.
Figure 40: RMSD for the universal amortised model the KS (left) and Kolmogorov (right) datasets for varying P C' scenarios. For KS we use a window model, whereas for Kolmogorov we use windw 5 model.
44

---

I Frequency Analysis
Lippe et al. [42] showed that capturing high-frequencies in the spectrum is crucial for accurate long- term rollouts. Following their analysis compute the frequency spectra of one-step prediction for all the methods considered in Sec. for both KS and Kolmogorov. In Fig. 41, we compare the spectrum of one-step prediction at the beginning and end of the trajectory. the case of KS, the spectra of the states generated by our models do not vary significantly depending on how far away they are from the initial state. In contrast, PDE-Refiner significantly overestimates the amplitude of high frequencies as the states depart from the initial conditions. However, we can see that at the beginning of the trajectory, PDE-Refiner can capture higher frequencies more accurately than all the considered methods.
For Kolmogorov the frequency spectrum the coarser spatial discretisation used or the nature of the underlying PDE. We compute and plot the average frequency spectrum for the channels (i.e. the two velocity components). From Fig. 42, we do not see much difference between the spectra of different models all models are capable of capturing the ground truth spectrum well.
varied the one from KS. Possible reasons include
Example 1, 12Ar
Example 1, = 6014r
Amplitde
Amplinde
Wavenumber
Wavenumber
Example 2, 12AT
Example 2, 604AT
Rehoo
Goand Trabs
Satoer
caall U.N
Jest AS
MSE-nxsed U.Nks
Amplitude
Uxivenal anertsed
MST aa PNO
Amplitude
Usward conntioa
MSE eocee END
Wavenumber
Wavenumher
Figure 41: Analysis of the spectra of two KS trajectories for different methods (Joint AR, universal amortised, PDE--Refnner MSE-trained U-Net, MSE-trained FNO). Each row corresponds to different trajectory. The first column shows the spectra of states closer to the initial state, while the second corresponds to states further away.
45

---

Channel 1,3 = 62T
Channel 2, = 6Ar
Gouad Trch
208-Reteer Mstiro UiNe
Jone
Amplitude
Amplitude
MSE-naee
calad Mattenttad
Univrrl voomed
Wavenumber
Wavenumber
Channel 1, 51Ar
Channel 2, f 51A7
MSE-mzzed Uolk
Ampliaude
Amplitude
thovare mndd INO
Groead Thut
New AR
MSS Not
Cuirnnal azortacd
XSE PSO
Wavenumber
Wavenumber
Figure 42: Analysis of the spectrum of Kolmogorov trajectory for different methods (Joint AR, universal amortised, PDE-Refinet MSE-trained U-Net, MSE-trained FNO). Each row corresponds different time step (left closer to the initial state, right closer end of trajectory), and each column to a different channel.
J Stability of long rollouts
To investigate the long-term behaviour of the different diffusion models, we generate very long trajectory (I' = 2000 corresponding to 400 seconds) for the KS equation. We expect the samples to diverge from the ground truth after a certain number of states, but the question whether our models still generate physically plausible predictions. We illustrate in Fig. 43 an example of a long trajectory generated by the joint AR (middle) and universal amortised (bottom) models, showing that the generated states remain visually plausible throughout the trajectory. We also analyse the spectra different one-step predictions along the trajectories Fig. 44. We can see that both models manage to capture the low frequencies well, but overestimate the high-frequency components. This holds for all the different steps we considered- state close to the initial (known) ones, state in the middle of the long trajectory, and a state towards the end of the trajectory. This implies that the spectra do not change qualitatively along the trajectory length and hence, remain physically plausible throughout the trajectory.
True
Joint
Amortised
120
240
360
Time (in seconds)
Figure 43: Example of long trajectory (2000A7 steps, corresponding to 400s) generated by the joint AR (middle) and universal amortised (bottom) models. Note how, even if we are generating longer trajectories than what the models have been trained on, the states remain visually plausible.
46

---

? 10AT
1000
10a
Ground
10a.
Ground
Amplitude
10
Joint AR
Amplitude
10a
Joint AR
Universal amortised
Universal amortised
10-2
10-2
10
10-4.
10
10"55
ww
10
10
20
40
80
100
120
20
40
60
80
100
120
Wavenumber
Wavenumber
t = 2000AT
10a_
Ground Truth
Amplitude
109
Joint
10-2.
Universal annotiised
10
10-6
10-5
0
20
40
60
80
100
120
Wavenumber
Figure 44: Frequency spectra of a state close to the initial state (top left), in the middle of the trajectory (top right), and towards the end of the trajectory (bottom). In each plot, show the ground truth spectrum (blue), as well as the spectra of the trajectories generated by the joint AR (blue) and universal amortised models (orange).
K
Algorithms
Algorithm 1 Predictor-correctoor corrector (PC) sampling as implemented in [67, 60]. We will refer to this as the function predi corrector-eee step (N,CS)
Input:
Discretisation steps N for the reverse-time SDE, corrector steps CS
Output: Trajectory z0)
Initialise xT7) Pr
for i = N to ( do
att;) Predictor (x(4++1))
Using Eq. (5) as score
for j = 1 to CS do
((,) +- (orrector(z((4))
b Using Eq. (5) as score
end for end for return (0)
Algorithm 2 All-at-once (AAO) rollout in [60]
Input: Trajectory length L, conditioning states F1.C. discretisation steps N for the reverse-time SDE, corrector steps CS
Output: Trajectory Ø1.L (0) approximation of the score of the full trajectory and using guidance return z=0)
(N, CS) Using Algorithm 2 from [60] to get an
47

---

Algorithm 3 AR rollout with joint model
Input: Trajectory length L. Markov window M, initial true conditioning states 21:C. prediction states P p= .M C, discretisation steps N for the reverse-time SDE, corrector steps OS
Output: Trajectory 21:L
Initialize Zgen [1cc
while len( < L do
Define conditioning states as C-
D Selecting last C' states from Zgen This has length M Selecting the last P states
(0) +
Jgen < Tgen + (= :]]
end while
F1:L = Zgen
return Z1.L
Algorithm 4 AR rollout with a universal or classic amortised model
Input: Trajectory length L. Markov window M, initial true conditioning states & prediction states P M C, discretisation steps N for the reverse-time SDE, corrector steps CS
Output: Trajectory 21:L
Initialize Zgen [+1.c]
while len (ggaa) < L do
Define conditioning states as Zigen (-C :] (0) + wwweecctoororortttt.o Seweecttor-step-/... (N cS)
D Selecting last C states from Z'gen D This has length P
Fyen Tgen + (0)
end while
TT.L = Tgen
return Z1.L
48

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The experimental contributions are supported by Sec. 5 in the main text of the paper. We also provide extended results in App. and App. F. The theoretical findings are discussed in App. B.1 and Sec. 3.3.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitation of our work in the Sec. 6. In addition to that, we present runtime analysis of the different rollout strategies both the main paper, see Sec. 3.3 and in the supplementary material App. F.1. We report the computational cost in minutes for the DA task Sec. 5. One possible additional limitation of our work is that the results are computed using a single trained model with one random seed.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specifiiattin, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
49

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: In paper, theoretically justified the Markov score assumption in App. B.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
We present a full description of the data generation process for all the three datasets considered in the paper in App. D. App. D.4 we present all the architecture and training details for all three experiments. In the appendix, we also report results for different parameter settings. We include provisional code together with the submission and we plan to make the code publicly available upon acceptance.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
50

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Together with the paper submission we are including a starting code to run some of the experiments contained in the paper. Due memory constraints, able to share the exact data used for KS and Kolmogorov. However, we are providing all necessary instructions for data generation. In the README file included in the code, explain how to run the code. We plan to make the code publicly available upon acceptance.
not
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
SrridessCooeeSubee
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: // //nips wwtppcccppppppecccccccccccccccet wwwpc//gggs////eeee hewpgdee//eeeeeeiiiieeeeeeeeeeee www.desumisssiooooo...oo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We are reporting all the dataset and hyperparameters details in App. D and App. D.4. In addition to that, the submission code also contains these details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
51

---

Justification: In all the plots and tables we present in the main paper we explain in the caption how the error bars are computed, i.c we compute mean and standard error over all the test trajectories.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
In App. D.4, we summarised the hardware we used for training our models and running our experiments. Since relied cluster with multiple GPUs, it is challenging to provide precise mapping between the experiment and the specific GPU type used. However, everything was for the hyperparameters search.
Justification:
single GPU. We did not track how many runs used
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpss//narrrss..c/pbbbl///t/tt cc/ppbllic/Ehhhcecudllllllllllll
Answer: [Yes]
Justification: The experiments considered in this paper deal with PDE simulation. We used open-source code that we cited the paper to generate the data and we believe the paper conforms the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
52

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid-
eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: This work tackles the task of generating realistic trajectories dynamical systems with underlying physics described by partial differential equations (PDEs). PDEs are ubiquitous in Science and Engineering fields, as they are the foundation for describing and solving many practical problems. Fluid dynamics, which studies the flow of liquids and gases, is perhaps the field that relies the most on solving PDEs. Applications can range from profiling turbines with optimal performance when designing wind turbines hydroelectric dams, to optimising the aerodynamics of airborne objects such as planes or missiles. Our work, however, does not engage with any of these tasks specifically.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Since our work focuses on score-based models for PDE simulations, we helieve it does pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided relcasing unsafe images.
53

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: In the paper, we use open-source code for both data generation and the baselines. We credited them both in the paper and in App. D and App. D.4.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged to reach out the asset's 's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The assets introduced by this paper are in the form of models, i.c. code. We described in detail the architecture choices and the training details of the different models in App. D.4 and the data generation process in App. D. In addition to that, made the whole codebase publicly available.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
54

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
55


Reviews:
[{"review_id": "tzYOwM5Nbr", "reviewer": "Reviewer_D4vs", "summary": "- The paper tackles of the problem of ML-based PDE modelling. Specifically two sub-problems: (a) *forecasting*: to generate rollouts given initial observations; and (b) *data assimilation*: to refine a trajectory given partial and noisy observations.\n- The proposed approach extends a score-based diffusion model solution (notably, related to Rozet & Louppe NeurIPS '23, and Lippe et al. NeurIPS '23) that jointly tackles the two sub-problems. \n- First, by proposing a *joint model* that allows for auto-regressive sampling of the trajectory (in contrast to \"all-at-once\" of Rozet & Louppe). Second, to consider *amortization* using a conditional diffusion model that no longer requires multiple trained models for choices of prediction and correction steps.\n- The approach is evaluated on standard benchmarks 1D Kuramoto-Sivashinsky (\"KS\") and 2D Kolmogorov flow (\"Kolmogorov\") and compared with standard baselines (e.g., PDERefiner of Lippe et al.). For forecasting, the results are on-par or behind baselines. For data assimilation, the approach outperforms Rozet & Louppe.", "questions": "1. **Forecasting results**: can the authors comment on gains of the proposed approach for forecasting?\n2. **Experimental settings vs. AAO**: can the authors comment on how the setup/evaluation regime differs from that of AAO?", "limitations": "Yes, this is a dedicated and balanced discussion on this.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper proposes a common framework to tackle multiple PDE simulation related tasks of forecasting, offline and online data assimilation. This appears to constrast existing work that solely focus on one.\n2. The paper is overall well-written. I especially appreciated the related works discussions, which also highlight concurrent works.\n3. The evaluation is reasonable, and considers multiple relevant baselines (including very recent ones).\n4. The approach (although involves more computation) has some practical merits e.g., using \"universal amortized\" strategy to rely on a single trained model.", "weaknesses": "## Major concerns\n\nI have no major concerns with the paper.\n\n## Minor concerns\n\n**1. Results - Forecasting**\n- While a common framework for multiple PDE simulation tasks is appealing, I wonder the specific contributions for the forecasting problem given that many baselines outperform the proposed approach (and if I understand right, with lower computation cost).\n- Although, the paper claims (p8, L299) the approach is appealing for flexibility, I am unclear on scenarios that require this at the price of empirical gains.\n\n**2. Slightly unclear - experimental settings/results on offline DA vs. AAO**\n- Firstly, I have to preface that the AAO paper (Rozet & Louppe, 2023) doesn't appear to have an extensive quantitative analysis. \n- However, I am curious of the difference of the experimental setup between this and AAO paper: is the focus here sparse-observations, whereas AAO reports dense observations? \n- Moreover, in the dense observation setting, I see that errors (Fig. 3) are comparable. As a result, I wonder whether the results are in favour of the proposed approach only in the sparse observation regime (where the errors, although better than baselines, are still significantly large).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their positive consideration. We address the points raised in the Weaknesses and Questions sections below.\n\n**Forecasting results**\n\nIt is true that in the plain forecasting task, our proposed models do not achieve SOTA performance, as we also highlight in the limitations section of our work, with the main reason being that we choose to trade off forecasting performance with conditioning flexibility. Nevertheless, the paper brings two main contributions in the context of forecasting:\n\n1. With the proposed autoregressive (AR) sampling approach, we show that diffusion-based models can achieve results of comparable quality with other benchmarks (e.g., MSE-trained U-Net, MSE-trained FNO), while preserving flexibility. This is in contrast to the all-at-once (AAO) sampling approach proposed by [1] which predicts trajectories that diverge from the ground state early on.\n\n2. We propose a new training strategy for amortised models (which can also be applied to non-diffusion based models), which achieves stable performance for a wide range of conditioning scenarios. This is in contrast to other amortised approaches from the literature which report a decrease in performance with increasing history length [2].\n\n**Scenarios where flexibility is important**\n\nA setting where flexibility is important is the mix of forecasting and data assimilation (DA) considered in the paper, especially in the online DA case. More specifically, not only do we want to produce forecasts based on some estimated initial states, but we also gather some observations in an online manner that can help the model correct its previous forecasts to better match the observed data. \n\nOne important application of the above setting is *weather prediction*. The goal is to produce a forecast for a certain horizon (e.g. days to weeks for medium-range prediction). At the same time, we continuously acquire new observations from a variety of sources, such as weather stations, radiosondes, wind profilers, and satellites [3]. Thus, we would like to update our forecasts whenever these new observations arrive to better reflect reality. This is currently done in two stages, where the forecasting system first generates the forecast, and then the data assimilation system incorporates the observations. However, they are both very costly systems [4]. Most of the current ML-based approaches for weather modeling only tackle the forecasting component (e.g. PanguWeather [5], GraphCast [6]). A more preferable approach would be to perform both tasks at once, being able to flexibly condition on both initial states, as well as any incoming observations. And this is exactly what we are proposing in this paper, with the proposed online DA task being the one that most closely resembles this scenario. However, we do acknowledge that it is still in a toy-ish setting and a natural future direction involves investigating the potential of the technique to scale.\n\nFinally, another setting where flexibility is important is where the nature of the conditioning information might change over time, implying that during training one does not have access to examples of the conditioning information they might encounter at test time. In our case, we address this by conditioning the models using reconstruction guidance. In terms of practical applications, we can imagine an example in the weather modelling setting. For example, a new sensor might be added to the sources of information. Due to the flexibility of our models, we would not have to re-train the model to account for this new instrument.\n\n**Experimental settings vs. AAO**\n\nIndeed, in contrast to our work, [1] does not provide any quantitative results (on the Kolmogorov flow).\n\nTheir paper performed the following experiments: \n- Observing the velocity field every four steps, coarsened to a resolution 8 $\\times$ 8 and perturbed by Gaussian noise ($\\Sigma_y = 0.1 \\mathbf{I}$).\n- Observing a regularly sampled sparse velocity field at every time step, with factor $n$ ($n$ implies 1 observation for each $n \\times n$ region). They use $n$=2,4,8,16. $n$=16 implies 0.3% observed data. \n\nThey did not use the 1D KS dataset, only 2D Kolmogorov flow data. Moreover, in our work, we sample observations at random locations, and at each time step. The fraction of observed variables varies between 0.1% to 30%. \n\nIn our experiments, we show that the proposed method is able to perform at least on par with AAO across different settings, and it could, thus, be useful if we expect the fraction of observed variables to vary.\n\n**References**\n\n[1] Rozet, F. and Louppe, G. (2023). Score-based Data Assimilation. In Thirty-Seventh Conference on Neural Information Processing Systems.\n\n[2] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.\n\n[3] Keeley, S. (2022). Observations. https://www.ecmwf.int/en/research/dataassimilation/observations. \n\n[4] Bauer, P. et al. (2020). The ECMWF Scalability Programme: Progress and Plans.\n\n[5] Bi, K. et al. (2023). Accurate medium-range global weather forecasting with 3D neural networks. Nature, 619(7970):533–538.\n\n[6] Lam, R. et al. (2023). GraphCast: Learning skillful medium-range global weather forecasting.", "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": "i4bzpZXgxZ", "reviewer": "Reviewer_iiGR", "summary": "The authors perform an extensive study of generative diffusion models applied to the task of PDE forecasting and data assimilation (DA). Further, they introduce an autoregressive (AR) sampling strategy, and a universal amortised model based with variable context size, based on masking.", "questions": "1. What are the training times? How do they relate to model performance, i.e. validation metrics? \n2.  What is the long-term behaviour of the AR model? I.e. longer than seen during training.\n3. How do the Fourier modes of the predicted trajectories look like? Given that PDE-Refiner is a baseline it would be interesting to see how it compares.\n4. How well do the predicted trajectories match the physics? This is related to the question above (the energy spectrum, for example).\n\n -- End of questions --\n\nA brief explanation of the review and grades given, put here for lack of a better place: \n\nSoundness: 3 - extensive benchmarking\nPresentation: 2 - appropriate but limited validation metrics (see Weaknesses)\nContribution: 2 - extensive benchmarking, approach novel in this context, lacking some interesting metrics\n\nRating 6 - A technically solid paper with extensive tests. It is held back by the validations done, the scalability issues, and inference times.\n\nRating after rebuttal 7 - Added validation metrics and spectral studies. Scalability, accuracy in forecasting task and inference times remain an open question for future work.", "limitations": "Some of the limitations are properly acknowledged. The time scaling of the AR sampling, and the lower performance on forecasting tasks compared to SOTA approaches. The fact that while flexible, the universal amortised model has a maximum conditioning size (as opposed to AR models) should be mentioned.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The authors have done a large number of tests with different configuration, leading to useful insights for the community that go beyond the architectures introduced in this paper. The AR sampling does not suffer from the poor memory scaling of AAO sampling. By leveraging masking during training (similar to how LLMs are trained), their universal amortised model is flexible and robust with respect to the size of the conditioning window.", "weaknesses": "AR sampling scales poorly in time as each step costs (1 + c)*p evals, that is, a diffusion sampling process, known to be expensive, has to be performed at every timestep. The inference times for both AR and AAO models are large compared to the baselines in the forecasting literature or to Rozet et al. for DA. The validation metrics presented could be clearer. The RMSD and MSE (appendix) presented are not normalized (in some sense) to account for the magnitude of the data, making the absolute number hard to interpret. Further, insightful metrics are missing, such as the energy spectrum of the predicted trajectories or long term trajectory stability. There are only a few comparison trajectories for the forecasting tests. While qualitative results are always subjective, it is known to be hard to judge the quality of a predicted trajectory just from quantitative summary statistics. Similarly, for the offline DA results, the partial observations are not shown.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time taken to review our paper and the positive consideration of it. We address the comments below. \n\n**Alternative metrics**\n\nThanks for suggesting alternative metrics, such as the energy spectrum, as well as providing more qualitative examples. We agree that these are indeed helpful in understanding the quality of the generated PDE trajectories, Please find in the attached PDF:\n- Example energy spectra (Fig. 3 - KS and Fig. 9 - Kolmogorov).\n- More examples of trajectories from each of the evaluated models on forecasting (e.g. Fig. 1 - KS and Fig. 2 - Kolmogorov, besides the examples provided in Fig. 24 for KS and Fig. 25 for Kolmogorov in the Appendix)\n- Example of trajectories for DA offline, alongside the masks used and interpolation results, which we included as a baseline for offline DA (Fig. 5 - KS and Fig. 6 - Kolmgorov)\n\nWe will make sure to include them in the final version of the manuscript.\n\nFor the quantitative evaluation, we decided to use high correlation time for an easier comparison with the results provided in [1]. RMSD/MSE are metrics that we frequently encounter in the literature (e.g. [1], [2]). While we acknowledge they are not perfect in evaluating PDE dynamics, we believe the problem of finding the right metric is a research question on its own and was not the focus of this investigation.\n\n**Training times** - We found that for the KS dataset, the joint model requires less training (4k epochs) compared to the universal amortized model (15k) steps (please also see Table 3 p.21 for more training details). However, both models, universal amortised and joint, were trained for the same amount of epochs for Kolmogorov. All considered models take from a few days to a maximum of one week to train depending on the dedicated GPU hardware. Since we relied on a GPU cluster with different machines on it, it is challenging for us to report a more specific comparison about the training times.\n\nStandard diffusion validation loss (measured by how well the model predicts the added noise) is not comparable across joint and amortised models, as the amortised model predicts the conditional score, so validation loss is often smaller than the one for the joint model. \n\n**Long-term trajectory stability** - The KS dataset is already analysing this behavior, since training examples are shorter than validation and test (140 states / 28s for train vs. 640 states / 128s for validation and test). The evolution of MSE error can be found in Fig. 23, where the error starts to linearly increase after around 50 seconds.\n\nTo analyse even longer-term behaviour (e.g. does the model still produce plausible states), we generated a trajectory of 2000 steps (400s), clearly showing that the model produces plausible predictions (see Fig. 4 on the attached PDF). Moreover, we computed the energy spectra of a states close to the initial condition, a state in the middle of the trajectory, and of the last state. We observe that the energy spectra of these three states do not show different behaviour, indicating that our model does not start generating physically unplausible states if we roll it out for very long times.\n\n**Fourier modes** - Please see Fig. 3 and Fig. 9 on the attached PDF for example energy spectra.\n\n**Physics match** - We believe the energy spectrum + long-term behaviour already address this question, but we are open to investigate other ways of assessing how well the generated trajectories match the physics.\n\n**References**\n\n[1] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.\n\n[2] Georg Kohl et al. Turbulent flow simulation using autoregressive conditional diffusion models. 2023.", "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": "0NzOCbxkiX", "reviewer": "Reviewer_ofj7", "summary": "This paper studies the application of diffusion models (DMs) to PDE forecasting and data assimilation, focusing on different approaches to condition the DMs on initial conditions (forecasting) or sparse observations (data assimilation). The conditioning can either occur explicitly during training by conditioning the score network (amortised model), or at inference time via reconstruction guidance (joint model). The authors present an improved autoregressive sampling technique for the latter, which is shown to be important for forecasting tasks. The methods are evaluated on multiple PDE datasets, including on a hybrid forecasting+data assimilation problem setup.", "questions": "- Is the MSE-trained U-Net using the same modern U-net architecture as your diffusion-based methods do?\n- For the (non-universal) amortised models, you train one different model for each unique combination of P | C?  How do you explain that the amortised models trained for a specific P | C combination underperform the universal amortised model that might need to distribute capacity between all the P | C combinations that it was trained on? \n- You could also perform the \"universal\" training of the amortised models (i.e. train on any combination of P|C) for non-diffusion models, right? That would be very interesting to see (e.g. it may or may not work for a MSE-trained model).", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The studied models can tackle both forecasting and data assimilation jointly. Even if performance lags behind forecasting-only models, there might exist use-cases for this, and performance could be improved in future work.\n- There are a lot of careful ablations and insightful comparisons of the studied conditional diffusion models variants.\n- The universal amortised model is an interesting approach (taken from the ML for natural videos literature) that seems to perform promisingly.", "weaknesses": "- The forecasting evaluation could be enhanced with probabilistic metrics such as the CRPS and spread-skill ratios, by evaluating ensemble forecast of the diffusion models. It would also be beneficial to see evaluations of the spectra of the predictions. \n- I think that it would be useful to include algorithms or pseudocode for how the different methods (joint AR, joint AAO, amortised etc.) are trained and how they function at inference time. This will make it easier to compare the different approaches in a more detailed and self-contained way.\n- In fig. 1 for Kolmogorov: It would be nice to include some results for the amortised model, especially since you observe differences between the datasets.\n- I might have missed an ablation for the window size W=P+C of the universal amortised model. This would be especially interesting to see on the Kolmogorov dataset (i.e. basically repeat Fig. 1 right but with W \\in {3, 7, 9} etc.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and your comments, we aim to address them below.\n\n\n**Evaluation metrics** - Thank you for your suggestions, when we chose the evaluation metrics we followed [1], but we acknowledge that some other metrics can be used to assess the model. We provide some examples of energy spectra for the trajectories in the attached PDF (Fig. 3 and Fig. 9), showing that they are in agreement with the ground-truth spectra for low frequencies, and the differences mostly come from high frequency components, as also noted in [1].\n\n**Algorithm/Pseudocode** - Thank you for the suggestions, we agree that this would make it easier to compare the different methods and will include them in the updated manuscript should the paper be accepted.\n\n**Amortised Kolmogorov** -  Thank you for your suggestion. These models are still training at the moment because the cluster is currently very busy and they have been spending some time in the queue. We hope to be able to provide the results by the end of the discussion period. But this is definitely something we will include in the final version of the manuscript.\n\n**Ablation window size** - We launched models for window 3 and 7 for Kolmogorov for both the joint and amortised models. However, please note that because of the time constraints, their hyperparameters are not tuned and from the training curves, we can clearly see that the window 7 model would benefit from more training. We provide a few example results in terms of high correlation time below, as the other scenarios are still currently being queued on our cluster:\n\n- Amortised $1 \\mid 2$ (window $3$) - $9.09 \\pm 0.50$s\n- Amortised $1 \\mid 4$ (window $5$) - $8.85 \\pm 0.53$s - from paper\n- Amortised $1 \\mid 6$ (window $7$) - $8.63 \\pm 0.50$s\n- Amortised $3 \\mid 4$ (window $7$) - $8.01 \\pm 0.44$s\n- Joint $1 \\mid 2$ (window $3$) - $7.63 \\pm 0.41$s\n- Joint $1 \\mid 4$ (window $3$) - $8.16 \\pm 0.62$s - from paper\n- Joint $1 \\mid 6$ (window $7$) - $7.73 \\pm 0.50$s\n\nFor the joint model, we can see that the performance is comparable between all window sizes, although the best one is achieved with the window $5$ model. We believe the window $7$ one does not outperform it because of convergence issues.\n\nFor the universal amortised, we see a decreasing trend in performance with increasing window size. However, the performance of all models with $P=1$ is within error (and hence, comparable). We believe that this result might be influenced by hyperparameter tuning.\n\n**Architecture** - No, the MSE-trained U-Net and the modern U-Net architecture are not the same. We have experimented with several architectures and found that the architecture inspired by [2] (SDA) was the most stable one in the context of joint prediction. The MSE-trained U-Net is inspired by [1] and we found it better suited for amortised-style of models. As such, we optimised the architecture for each model and chose the optimal one, meaning that we used the SDA architecture for the joint and for the amortised models (for a one-to-one comparison), and the PDE-Refiner architecture for the baselines used in forecasting. \n\nWe will include a more detailed description of the MSE-trained U-Net in the Appendix in the updated version of the manuscript, thanks again for pointing this out.\n\n**Non-universal vs. universal** - You are correct and we also thought this finding is interesting. We believe the reason why the universal amortised models outperforms the plain amortised models (where we need to train a model for each $P \\mid C$ combination) is because the training task becomes more complex, which encourages generalisation. In effect, the universal amortised model also amortises over the predictive horizon $P$. We note, however, that this comes with the downside that in the KS case we had to train the universal amortised model for longer than the plain amortised models (which would start overfitting if they had been trained for the same amount of epochs as the universal amortised model).\n\n**Universal amortised for non-diffusion models** - This is indeed possible, as the framework we propose is not specific to diffusion-based models. We find your suggestion very interesting, but potentially not directly relevant to the topic of the paper, which aims to study diffusion-based models for PDE modelling.\n\n**References**\n\n[1] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.\n\n[2] Rozet, F. and Louppe, G. (2023). Score-based Data Assimilation. In Thirty-Seventh Conference on Neural Information Processing Systems.", "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": "0N8ZWeXFOy", "reviewer": "Reviewer_4n49", "summary": "The manuscript addresses the challenging problem of conditional diffusion modeling, towards accurate and efficient data assimilation in problems governed by PDE. The authors compare different diffusion modeling approaches in this setting, focusing on the auto-regressive setting, and also propose new approaches. The comparisons are done on relevant PDE benchmarks (chaotic, 1d and 2d spaces).", "questions": "1) l96: how can we divide by $\\mu_t$ if it is a vector? Is it done component-wise?\n2) It is probably my lack of experience with diffusion models, but how do we know the \"true\" probability density p to compare to (resp, $\\nabla \\log p$)? For example, in equation (3) and l.138, \"the local score network ... is trained by minimizing the following DSM loss\", which contains the expression $\\nabla \\log p$. This is probably not contained in the training data as a value for each x, right? How do we compute it? Equation 7 does not help, because it is also not clear to me how to compute the individual probability densities (for each i). In Equation 8 I can understand that the conditional is assumed to be Gaussian, which then leads to a simple expression, but it cannot be the case for the general density.\n\n3) l305: \"We assume to always observe some full initial states\" this is very confusing to me. What is \"some full\"? Does it mean the full initial state is observed? This would invalidate the entire approach, because then there is no need for \"data assimilation\", the PDE can just be solved with a classical solver, from this given initial condition. Looking at appendix G, it seems that only a portion of the initial state is \"always\" given, which would be ok - then the classical approach is not really possible - but then the sentence should be rewritten to clarify this.", "limitations": "Some limitations are discussed, but not in the direction of a comparison to classical solvers (as I mentioned in the \"weaknesses\" section). It is not clear if the diffusion based modeling can be better than classical solution methods for the particular problems, and if not, which other problems it can be used on that are hard to solve with classical methods.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Using conditional diffusion models for data assimilation is an interesting approach, and the results on the given PDE systems are impressive. The paper is relatively easy to read, and contains interesting suggestions for improvements. I very much appreciated the brief introduction to score-based modeling, as I am not working day-to-day with diffusion models (more with PDEs). Scalability of the approaches are also addressed, briefly in the main part and more detailed in the appendix. In general, the appendix is also extensive and contains a lot of details and additional studies related to the main paper.", "weaknesses": "1) The PDEs considered as examples have turbulent / chaotic dynamics. There is not discussion about this, but it is important to mention: comparing trajectory data is not a good metric, because the PDE itself is causing sensitivity to the initial conditions, not the models. Rather, the PDE itself could be used as a metric (i.e., evaluate the PDE on the generated trajectory and check if it is zero); alternatively, Lyapunov exponents or density measures on the attractors can be used for comparison. The 1D Burgers equation does not have these problems. The \"correlation time\" defined by the authors is probably a way to mitigate the issues with chaotic/sensitive dynamics, but is not ideal, because the divergence from the true solution depends on the initial condition. Regardless of this issue with the metric, it seems the trained models are capable of generating trajectories close to the attractor (i.e. they \"look reasonable\"), which is already impressive.\n\n2) There is no comparison to a classical scheme that solves the inverse problem. Especially looking at the long solution times (fig.3, between 25-100 minutes); I think it is possible to design a classical inverse problem solver that (a) starts with a random initial state of the PDE (b) solves the PDE with this guess, and then (c) updates the guess based on the error to the observations, and does all this iteratively in less than 25 minutes. It is ok to study how diffusion models perform in this setting, but it is not ok to not even mention this possibility of the classical approach in the related work. Many inverse problem solvers exist in this direction (also data-driven ones), a good overview is here:\n\nArridge, S., Maass, P., Öktem, O., Schönlieb, C.-B., 2019. Solving inverse problems using data-driven models. Acta Numerica 28, 1–174. https://doi.org/10.1017/S0962492919000059\n\nIt would be good to include a brief discussion of the benefits of the neural network / diffusion model approach compared to the more classical setting.\n\n3) There is no theoretical justification (meaning: no rigorous proofs) of the suggested new architectures and training methods. The new approaches are motivated plausibly in the main paper, with additional computations done in the appendix, but without theory.\n\n\nMinor:\n1) l103: the statement $\\mathcal{P}(u;\\alpha)=\\mathcal{P}(\\partial_z u,\\dots)$ is not correct. Does $\\mathcal{P}$ only depend on $u$ or only (?) on its derivatives? Are PDEs  like $<v,\\partial_z> u + u=0$ not included (because they include $u$ as well as $\\partial_z u$)? Probably this is just an abuse of notation here.\n2) l126: do the authors mean \"...time dependency [...] $=x_{1:L}(t)$\" instead of \"$x_{1:L}(0)$\"?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive consideration of our work and the useful feedback on it.\n\n**Turbulent dynamics + metrics** - Thanks for your suggestion, we will add a mention about the turbulent dynamics of the data. For a discussion about metrics and spectrum, please refer to common answer.\n\n**Classical solvers vs ML-based + runtime** - This is a great point and we will include a discussion in the updated manuscript. We believe ML-based approaches for PDE modelling are still fairly new and thorough comparisons between them and classical solvers will emerge. In the case of our paper we considered implementing such methods is out of the scope, since the focus is on ML-based solutions, with emphasis on diffusion models. But we agree that we should mention the possibility of implementing the classical approach.\n\nA motivating application of our work is ML-based approaches for weather modelling. For a more detailed description of the weather forecasting setup please see reply to Reviewer D4vs. Works in this field, such as Aardwark Weather [1] show that there is potential for ML-based methods to reduce the cost of medium-range weather forecasting (which is currently based on classical solvers) by orders of magnitude, while still achieving competitive performance. There have also been works proposing to tackle the forecasting component using diffusion-based models [2], where the probabilistic treatment is appealing - it allows one to predict a range of probable weather scenarios, which can be crucial for decision making, disaster prevention, etc.\n\nRegarding runtime, as pointed out in the common response, the universal amortised model can achieve a significant speed-up (8x faster) during sampling compared to what we report in Figure 3 in the main paper.\n\nWe will incorporate a subsection on classical solvers in the related work section, mentioning their applicability in the tasks considered. We will also add a discussion on how ML-based methods compare to classical solvers.\n\n**Theoretical justification** - We acknowledge that classical methods benefit from a much more mathematically-rigorous framework. Unfortunately, guarantees for ML-based methods are rarely available (not only in the context of PDE modelling). Diffusion models have mathematical foundations (e.g. [3] or [4]), but it is still non-trivial to develop guarantees when taking into account the complex network dynamics within diffusion models. As such, this is probably more of a limitation of the field rather than of this paper in particular, and it is unclear whether such guarantees will ever be available for ML-based methods.\n\n**W1minor** - Thanks for pointing this out. The dependence on $u$ should be included as well.\n\n**W2minor** - L126: We omit putting (0) as time dependency and just use $x_{1:L}$. The initial Markov assumption is introduced on the sequence coming from the original data distribution.\n \n**$\\mu_t$** - Yes, $\\mu_t$ is a scalar value and then we divide the numerator component-wise.\n\n**Denoising score matching** - Indeed, computing $\\nabla_x \\log p(x)$ is not easy as we do not know $p(x)$. The workaround is to use the denoising score matching loss (Eq. 3), where we use $\\nabla_x(t) \\log p(x(t)|x(0))$ instead of $\\nabla_x \\log p(x)$. For a justification of why this is the case, see Appendix of [5]. Then, the quantity of interest becomes $p(x(t)|x(0))$, which is referred to as the noising kernel in diffusion models. This can be derived from the governing noising SDE. In particular, if the drift and the diffusion coefficient are affine transformation of $x(0)$, then $p(x(t)|x(0))$ is Gaussian. In our case, we are using an Ornstein-Uhlenbeck process (see Eq. 1) as noising process and therefore this satisfies the above conditions making $p(x(t)|x(0))$ Gaussian.\n\n**Offline DA setup** - The initial states are fully observed in the offline DA setting - we found this to be a convenient setting to analyse the performance of the model in a variety of settings, going from almost pure forecasting (the sparsest setting) to a very dense observation regime. We understand that in the context of classical solvers, where the treatment is deterministic, this just collapses to plain forecasting. However, sparse observations can potentially be helpful in reducing accumulating approximation errors (resulting from e.g. space/time discretisation). Nevertheless, this was not the focus of the experiment, but rather to see whether our approaches can “correct” their forecasting mistakes (i.e. make use of the observations to bring the predictions closer to ground truth), and alleviate error accumulation which is an inherent issue of autoregressive approaches. \n\nIn the online DA experiment, the initial states were **not** fully observed, and thus, could not be tackled by a classical numerical solver. This is the task we believe best reflects the advantages of our approaches, and acted as the main motivation when developing them. The previous investigations on forecasting and offline DA were useful investigations to understand the workings of the models in a variety of conditioning scenarios, but we consider the online DA task to be the one with most relevance to upstream applications, such as in weather modeling.\n\n**References**\n\n[1] Vaughan, A. et al. (2024). Aardvark weather: end-to-end data-driven weather forecasting.\n\n[2] Price I. et al. Gencast: Diffusion-based ensemble forecasting for medium-range weather, 2024.\n\n[3] De Bortoli, V. Convergence of denoising diffusion models under the manifold hypothesis. TMLR (2022).\n\n[4] Benton, J. et al. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization. ICLR\n\n[5] Vincent P. A connection between score matching and denoising autoencoders, 2010.", "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": "eXjWEr1b8G", "reviewer": "Reviewer_D3JV", "summary": "The paper studies forecastiong temporal dynamics associated with forward PDE problems using diffusion models. Different strategies for approximating the score function of trajectories and how to parameterize them are proposed, ranging from sampling the whole trajectory all at once to sample it step by step autoregressively. In addition, the authors also propose to incorporate information from observation to refine and modulate the prediction process. The proposed methods are tested on two types of PDEs under different settings involving both forecasting and data assimilation.", "questions": "1. In its current form, the best model variant of the proposed method still falls behind PDE-Refiner. Given PDE-Refiner is also closely related to denoising diffusion model (with velocity prediction, starting from last time step instead of pure noise), it seems that it is also possible to apply many of the proposed techniques to PDE-Refiner. More specifically, according to my understanding, PDE-Refiner uses the amortized architecture (P=1 and C=1) and it is autoregressive, while the paper has showcased the proposed universal amortized method seems superior (can predict more time steps at once, flexible time window and can use longer context). Can the authors provide more comments related to this point?\n\n2. The major limitation of applying vanilla diffusion model to simulate PDE autoregressively is that it increases the NFE by several times (e.g. if the diffusion model needs k steps to sample one step then it will need kT NFE to sample T steps). The proposed method can alleviate this by doing joint sampling and predicting multiple future steps in one sampling loop, yet not too many study and comparison on NFE is presented in the main text. It will be clearer if computational cost statistics regarding the proposed method and existing methods like PDE-Refiner are also reported.\n\n3. Compared to pure data-driven forecasting models like using UNet trained with MSE target, how much performance gain does online data assimilation bring? (the authors have conducted experiments on varying portions of observed data but it can be helpful to see a comparison against method that cannot do data assimilation)", "limitations": "The authors have discussed the limitation in the conclusion section.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The writing of the paper is pretty clear and self-contained.\n\n2. Comprehensive empirical studies of different components and hyperparameters are carried out to support the claims.\n\n3. The proposed universal amortized model, while being straightforward, is quite effective and flexible. It resolves several issues existed in common amortized models, including fixed context windows and performance degradation when given a longer context.\n\n4. Applying reconstruction guidance to the diffusion process can be potentially useful for building hybrid forecasting system, e.g. combing data-driven forecasting model with some physics prior.", "weaknesses": "1. Many techniques in the proposed framework are not new. Using the local segment's score to approximate and decompose the score of global dynamics is introduced in Score-based Data Assimilation. Reconstruction guidance has been used in several previous works that applied diffusion model to inverse problems (e.g. Manifold Constrained Gradient).\n\n2. The empirical performance of the proposed method does not outperform many data-driven forecasting models. Nonetheless, as authors have stated, the objective of the proposed framework is not only to do pure data-driven forecasting but do flexible and versatile sampling tasks such as online data assimilation. However, I think flexibility does not necessarily have to be orthogonal to the performance (see question section)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful consideration of our work. We have taken the time to address all the points raised in the review.\n\n**Novelty** - While we acknowledge that the techniques per-se are not new, we were the first to: \n1) quantitatively evaluate the decomposition on physics data; \n\n2) propose AR sampling instead of AAO sampling, which works much better in scarce observed data regimes and is key to performing forecasting; \n\n3) use both reconstruction guidance and amortisation on different sets of variables simultaneously, which improved model performance and flexibility.\n\n**Flexibility vs performance** - We agree with the reviewer that flexibility does not necessarily have to be orthogonal to performance. One of the goals of the paper is to quantitatively evaluate the current status of diffusion models for PDE simulation compared to other deep learning approaches by highlighting pros and cons of each such approach on the considered tasks (forecasting, DA, and a mix thereof). Unsurprisingly, we found that in the case of pure forecasting, our methods lag behind SOTA methods that have specifically been designed for forecasting. However, we argue that even so, the forecasting performance of our methods is close to other commonly used ML-based benchmarks, while benefiting from a much more flexible conditioning framework.\n\n**Universal amortised training for PDE-Refiner** - Indeed, the “universal” amortized training and larger window size can be easily incorporated into PDE-Refiner. However, reconstruction guidance and data assimilation cannot be applied out-of-the-box, as PDE-Refiner first produces an estimate of the next state, and then refines this estimation using noise. This is in contrast to plain diffusion models that gradually build this estimation from pure noise, rather than from the neural operator prediction.\n\n**NFE** - We provide a more detailed analysis of the NFE below in the KS and Kolmogorov cases.\n\nFor KS we simulate $640$ states. For the joint AR and amortised models, we use $128$ diffusion steps in the paper with a window $9$ model, and the conditioning scenario $P \\mid C$ depends on the task. Overall, the computational cost is $(\\frac{640-9}{P} + 1) \\times 128$. In contrast, PDE-Refiner first produces the NO estimate, and then performs $3$ refining steps, leading to $4$ FE for each state. They always predict one state at a time, leading to NFE=$640 \\times 4$, indeed leading to a lower computational cost than our approaches.\nHowever, subsequent analysis on the amortised model indicates that we can achieve similar performance with the universal amortised model by performing as few as 16 diffusion steps (please also see Fig. 7 in the attached PDF) (This is not true for the joint AR approach). This would lead to NFE = $(\\frac{640-9}{P} + 1) \\times 16$, which becomes comparable (and depending on $P$, even lower) than PDE-Refiner cost. \n\nFor Kolmogorov, the same analysis holds, but we only need to simulate $64$ states, and found that for the universal amortised model we can also achieve good performance with $16$ diffusion steps.\n\nThus, the universal amortised approach is comparable to PDE-Refiner in terms of computational time, while we acknowledge that the joint AR approach is more computationally intensive.\n\n**MSE-based U-Net for offline DA** - In the online data assimilation case, we assume that in the first step we do not observe the initial condition (lines 346-348), but only some sparse observations within the $s$ first time steps. Hence, plain autoregressive U-Net trained with MSE target cannot be directly applied to this task, as it expects to have fully observed initial state.\n\nHowever, we have included a comparison in the offline DA scenario, where we report the RMSD for the MSE-trained U-Net based on the initial conditions alone (please see Fig. 8 in attached PDF).\n\nWe hope we have successfully addressed your concerns.", "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:
We would like to thank the reviewers for the time taken to review our paper, the overall positive consideration of our work, and their feedback and useful suggestions. We are content that we managed to get the main conclusions of our investigation across, providing an “insightful comparison of the studied conditional diffusion models variants”. We are also happy that the universal amortised approach we introduced was found to be “interesting”, “effective”, and “flexible”, with promising performance, especially on tasks that involve a mix between forecasting and data assimilation. Such an example is the online DA scenario, which was inspired from related real-life applications from weather modeling. In the following, we provide a general comment addressing the common concerns of the reviews. We are happy to clarify further during the discussion period if some concerns still remain.

**Performance metrics**

Several reviewers suggested alternative performance metrics, besides the ones used in the paper (high correlation time, RMSD, and per time-step MSE in the appendix). We chose these metrics following other works from the field ([1], [2]). However, we do acknowledge that they have their weaknesses, and the usefulness of complementing the quantitative metrics with some qualitative ones to better understand the behaviour of our models. As such, we provide in the attached PDF:

- In addition to Fig. 24 and Fig. 25 from Appendix F.5, we provide additional examples of forecasting trajectories for qualitative evaluation of our models (Fig. 1 - KS and Fig. 2 - Kolmogorov). For offline DA, we provide additional examples where we show the generated trajectories (joint AR), alongside the masks used, and an interpolation baseline for $10^{-2}$ and $10^{-1}$ proportion observed (Fig. 5 - KS and Fig. 6 - Kolmogorov).

- Energy spectra of predicted trajectories, to assess how they compare to the ground truth and how they compare to the spectra of other relevant models (e.g. PDE-Refiner). In Fig. 3 and Fig. 9 we show the spectra at the beginning of the trajectory and at the end of the trajectory. We can see that, in the KS case (Fig. 3), the spectra of the states generated by our models do not vary significantly depending on how far away they are from the initial state. In contrast, PDE-Refiner significantly overestimates the amplitude of high frequencies as the states depart from the initial conditions. For Kolmogorov (Fig. 9), we observe that all methods (except for MSE-trained FNO) generate similar spectra, that resemble the ground truth spectrum well. We do not observe much difference between states close and far away from the initial state.

- Example of very long KS trajectories (including energy spectra), to investigate the long-term behaviour of our models (Fig. 4). In particular, we expect the samples to diverge from the ground truth after a certain number of states, but the question is whether our models still generate physically plausible predictions. Fig. 4 shows that this is indeed the case. By investigating the energy spectra for different states along the trajectory, we do not observe much difference depending on how far away the states are from the initial one---all spectra indicate that the models capture the low frequencies correctly, but overestimate the high frequency components. Overall, these additional qualitative studies show that, even for longer-term predictions, the models still produce physically realistic states.

- As suggested by reviewer D3JV, we provide a modified version of Figure 3 (top) from the main paper, showing the RMSD for a plain forecasting baseline (MSE-trained U-Net). This aims to show the benefit of the additional observations in the offline DA setting, illustrating that the models can improve performance with relatively few observations (proportion observed higher than $\approx 4\times 10^{-3}$).

**Computational time**

Another common concern was regarding computational time of our approaches.

The joint AR approach does come with an increased computational cost as opposed to other techniques from the literature, but we believe that the rapid advances in diffusion model sampling could bring this cost down. However, the universal amortised model has a similar computational cost to SOTA methods (PDE-Refiner). This is not apparent in the paper and as such we included an extra figure in the PDF (see Fig. 7) that shows that while in the main text we used 128 diffusion steps, the amortised model can achieve similar performance with as few as 16 diffusion steps. We will include this finding in the updated version of the manuscript. For a more detailed analysis of computational cost, please also see response to reviewer D3JV.

We hope we addressed most of your questions and concerns and are looking forward to further engaging with you in the discussion period.

[1] Lippe, P., Veeling, B. S., Perdikaris, P., Turner, R. E., and Brandstetter, J. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.

[2] Georg Kohl, Li-Wei Chen, and Nils Thuerey. Turbulent flow simulation using autoregressive conditional diffusion models. 2023.

Meta-review: The reviewers unanimously recommend acceptance (7-7-6-6-7). The author-reviewer discussion has been constructive and has led to a number of improvements to the paper. The authors are asked to implement the changes discussed with the reviewers in the final version of the paper.
Decision: Accept (poster)