Guided Autoregressive Diffusion Models with Applications to PDE Simulation

Authors: federico  bergamin
cristiana  diaconu
aliaksandra  shysheya
paris  perdikaris
josé miguel hernández-lobato
richard e turner
emile  mathieu
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 17:14:53

Abstract: Solving partial differential equations (PDEs) is of crucial importance in science and engineering. Yet numerical solvers necessitate high space-time resolution which in turn leads to heavy computational cost. Often applications require solving the same PDE many times, only changing initial conditions or parameters. In this setting, data-driven machine learning methods have shown great promise, a principle advantage being the ability to simultaneously train at coarse resolutions and produce fast PDE solutions. In this work we introduce the Guided AutoRegressive Diffusion model (GuARD), which is trained over short segments from PDE trajectories and a posteriori sampled by conditioning over (1) some initial state to tackle forecasting and/or over (2) some sparse space-time observations for data assimilation purposes. We empirically demonstrate the ability of such a sampling procedure to generate accurate predictions of long PDE trajectories.

Main paper: 
 introduction: 
 Partial differential equations (PDEs) are ubiquitous as they are a powerful mathematical framework for modelling physical phenomena. Solving PDEs relies on numerical methods like finite element methods (Liu et al., 2022) which require very small space and time discretisation to obtain accurate and reliable solutions resulting in significant computational resources. There has been a rising interest in leveraging deep learning methods to learn neural approximations of PDE solutions (Raissi et al., 2017) . However, since dynamics arising from PDEs can be particularly complex or chaotic (Hyman & Nicolaenko, 1986; Kevrekidis et al., 1990) , a major challenge is to predict long trajectories whilst preserving accuracy and stability (Kochkov et al., 2021) since ambiguities accumulate. Building on prior work, we propose to address these issues taking a probabilistic treatment of PDE dynamics (Yang & Sommer, 2023; Cachay et al., 2023) by leveraging diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021) . In particular, autoregressive amortised diffusion style of models, which work by additionally feeding previous states to the score network, are currently state of the art for long rollouts (Lippe et al., 2023; Kohl et al., 2023) . However, (a) prior work has found that higher Markov order harms performance, and (b) amortisation does not straightforwardly handle data assimilation.In this work we show how a diffusion model trained on short segments of joint states-coined GUARD-can solve both of these limitations. Our contributions are the following: (i) We introduce a joint diffusion model which can be used autoregressively for sampling variable length trajectories, by iteratively generating states conditioned on past ones via reconstruction guidance (Song et al., 2022) . (ii) We empirically demonstrate that this approach can perform on par with amortised models for generating long trajectories, and in particular that, unlike amortised models, using higher Markov order improves its performance. (iii) We show that this model can jointly tackle data assimilation, as well as forecasting, and in most settings outperforms the strategy introduced in Rozet & Louppe (2023a) which consists of sampling the full sequence at once. In summary, while this work does not introduce a novel methodology per se, we combine existing techniques to develop a PDE simulation method that is principled, accurate, and flexible. In addition to that, we perform a rigorous comparison between amortisation and guidance in the context of PDE modelling, especially on the forecasting and data assimilation tasks, that no prior work has considered. 
background: 
 Continuous-time diffusion models. In this section we briefly recall the main concepts underlying continuous diffusion models, and refer readers to Song et al. (2021) for a thorough presentation. We consider a forward noising process (x(t)) t≥0 associated with the following stochastic differential equation (SDE),EQUATIONwith w(t) an isotropic Wiener process, p 0 the data distribution, and g t a diffusion coefficient. The process defined by equation 1 is the well-known Ornstein-Uhlenbeck process, which geometrically converges to N(0, I). Importantly, under mild assumptions on p 0 , the time-reversal process ( ←x (t)) t≥0 also satisfies an SDE (Cattiaux et al., 2022; Haussmann & Pardoux, 1986) which is given byd ← - x (t) = -1 2 ← - x (t) -g 2 t ∇ log p t ( ← - x (t)) dt + g t dw(t),where p t denotes the density of x(t). In practice, the score ∇ log p t is unavailable, and is approximated by a neural network s θ (t, •) ≈ ∇ log p t , referred to as the score network. The parameters θ are learnt by minimising the denoising score matching (DSM) loss (Hyvärinen & Dayan, 2005; Vincent, 2011; Song et al., 2021 )L(θ) = E[∥s θ (t, x(t)) -∇ log p t (x(t)|x(0))∥ 2 ].Conditioning with diffusion models. The methodology introduced above allows for (approximately) sampling from x(0) ∼ p 0 . Sampling from the conditional p(x(0)|y), where y represents some observations with likelihood p(y|A(x)) and A a linear measurement operator, is more challenging. While by Bayes' rule we can rewrite it as p(x(0)|y) = p(y|x(0))p(x(0))/p(y), this is yet not available in closed form.Below we present two methods of simulating the conditional denoising process in Eq. ( 2) directly.d ←x (t) = -1 2 ←x (t) -g 2 t ∇ log p t ( ←x (t)|y) dt + g t dw(t).(2)Amortising over observations. The simplest approach is to directly learn the conditional score log p t ( ←x (t)|y) (Song et al., 2021; Ho et al., 2020) . This can be achieved by additionally feeding the observations y to the score network s θ , whose parameters are then learnt by minimising the following DSM loss L(θ) = E[∥s θ (t, x(t), y) -∇ log p t (x(t)|x(0))∥ 2 ], with (x(0), y) ∼ p 0×y .Reconstruction guidance. Alternatively, one can leverage Bayes rule, which allows to express the conditional score w.r.t. x(t) as the sum of the following two gradients ∇ log p(x(t)|y) = ∇ log p(x(t)) + ∇ log p(y|x(t)) ≈ s θ (t, x(t)) + ∇ log q θ (y|x(t)).(3)The conditioning comes into play via the ∇ log p(y|x(t)) term which is referred to as reconstruction guidance (Ho et al., 2022b; Chung et al., 2022; 2023; Meng & Kabashima, 2023; Wu et al., 2022; Song et al., 2022; 2023) . Assuming a Gaussian likelihood p(y|x(0)) = N(y|Ax(0), σ 2 y I) , we have p(y|x(t)) = p(y|x(0))p(x(0)|x(t)) dx(0) ≈ N y|Ax(x(t)), (σ 2 y + r 2 t )I ≜ q θ (y|x(t)) (4) where p(x(0)|x(t)) is not available in closed form and is approximated by a Gaussian q(x(0)|x(t)) ≜ N(x(0)|x(x(t)), r 2 t I) ≈ p(x(0)|x(t)) with mean given by Tweedie's formula (Efron, 2011 ) and variance Var[x(0)|x(t)] ≈ r 2 t being chosen as a monotonically increasing function (Rozet & Louppe, 2023a; Finzi et al., 2023; Song et al., 2022; Pokle et al., 2023) . 
score-based forecasting and data assimilation: 
 Problem setting We denote by P : U → U a differential operator taking functions u : R + × Z → X ∈ U as input. Along with some initial u(0, •) = u 0 and boundary u(τ, ∂Z) = f conditions, these define a partial differential equation (PDE) ∂u ∂τ = P(u; α) = P ∂u ∂z , ∂ 2 u ∂z 2 , . . . ; α , with coefficients α. We assume to have access to many accurate numerical solutions from conventional solvers, which are discretised both in space and time, and denoted by x 1:L = (x 1 , . . . , x L ) ∈ R D×L ∼ p 0 with D and L being the size of the discretised space and time domains, respectively. The overall goal is to train a machine learning model on these solved trajectories and use it to solve similar PDEs with different initial condition u 0 , boundary condition f or parameter α, as fast as possible while staying close to the true solution. We are generally interested in the problem of generating realistic PDE trajectories given observations, i.e. sampling from conditionals p(x 1:L |y). Specifically, we focus on two tasks: (a) forecasting: predicting future states x i:i+H given some past states x 1:i-1 , i.e. sampling from p(x i:i+H |x 1:i-1 ), with H the forecast horizon; (b) data assimilation: inferring the ground state of a dynamical system given some sparse observed states x o , i.e. sampling from p(x 1:L |x o ).Markovian dynamics. Learning a diffusion model over the full joint distribution p(x 1:L (t)) can become prohibitively expensive for long trajectories, as it requires a score network taking the full sequence as input-s θ (t, x 1:L (t)) ≈ ∇ x 1:L (t) log p(x 1:L (t)), with memory footprint scaling linearly with the length L of the sequence. One approach to alleviate this, as suggested by Rozet & Louppe (2023a) , is to assume a Markovian structure of order k such that the joint over noised data trajectories can be factorised into a series of conditionals p(x1:L ) = p(x 1 )p(x 2 |x 1 ) . . . p(x k+1 |x 1:k ) L i=k+2 p(x i |x i-k:i-1). Here we are omitting the time dependency x 1:L = x 1:L (t) for clarity's sake. The score w.r.t. x i can be written as∇ xi log p(x 1:L ) = ∇ xi log p(x i |x i-k:i-1 ) + i+k j=i+1 ∇ xi log p(x j |x j-k:j-1 ) = ∇ xi log p(x i-k:i+k ) with k + 1 ≤ i ≤ L -k -1, whilst formulas for i ≤ k and i ≥ L -k can be found in App. B.1.Consequently, instead of learning the entire joint score at once ∇ x 1:L (t) log p(x 1:L (t)), we only need to learn the local scores s θ (t, x i-k:i+k (t)) ≈ ∇ x i-k:i+k (t) log p t (x i-k:i+k (t)) trained on short sequences of size 2k + 1. When necessary, the full score s θ (t, x 1:L (t)) ≈ ∇ log p t (x 1:L (t)) can then be reconstructed from these local scores-see App. B.1 and Rozet & Louppe (2023a, Alg 2.).Sampling. We can generate samples x 1:L ∼ p(x 1:L ) by solving the time-reversal SDE over the sequence. Optionally, to avoid errors accumulating along the denoising discretisation, each predictor denoising step can be followed by a small number of Langevin Monte Carlo (corrector) steps.We now describe how to tackle the forecasting and data assimilation tasks by sampling from conditionals p(x 1:L |x o )-instead of the prior p(x 1:L )-with the trained local score s θ described in Sec. 3. The conditioning information y = x o ∈ R O in this case is a (possibly noisy) measurement of a subset of variables in the space-time domain, i.e. x o = A(x) = A vec(x) + η with vec(x) ∈ R N the vectorised trajectory, η ∼ N(0, σ 2 y I) and a masking matrix A = {0, 1} O×N . Plugging this in (4) and computing the score we get the following reconstruction guidance term∇ log p(x o |x(t)) ≈ 1 r 2 t +σ 2 y (y -Ax(x(t))) ⊤ A ∂ x(x(t)) ∂x(t). Summing up this guidance with the trained score over the prior as in (3), we get an approximation of the conditional score∇ x 1:L (t) log p(x 1:L (t)|x 0 ) = ∇ x 1:L (t) log p(x 1:L (t)) + ∇ x 1:L (t) log p(x 0 |x 1:L (t)) ≈ s θ (t, x 1:L (t)) + 1 r 2 t + σ 2 y (y -Ax(x 1:L (t))) ⊤ A ∂ x(x 1:L (t)) ∂x 1:L (t)and can thus generate conditional trajectories by simulating the conditional denoising process Eq. (2).All-at-once (AAO). We can generate entire trajectories x 1:L in one go as proposed by Rozet & Louppe (2023a). Assuming we observe some values x o ∼ N(•|Ax 1:L , σ 2 y I), the conditional score ∇ log p(x 1:L (t)|x o ) is obtained by summing the full sequence score s θ (t, x 1:L (t)) built by combining local scores as developed in Sec. 3, with the guidance term ∇ log p(x o |x 1:L (t)) estimated as above.Let's first focus on forecasting, i.e. sampling x C+1:L given x 1:C initial states. An alternative approach to AAO is to factor the forecasting prediction problem as p(x 1:L |x 1:C ) = the full trajectory by iteratively generating H states conditioned on the previous C ones, such that H + C = 2k + 1 with k being the Markov order. We refer to this approach as GUARD for Guided AutoRegressive diffusion model. Data assimilation is similarly tackled by further conditioning on additional observations appearing within the predictive horizon H at each AR step of the rollout.Modelling capacity. There are two main reasons why the AR sampling strategy outperforms the AAO approach. First, for the same score network with an input window size of 2k + 1, the AAO model has a Markov order k whilst the AR model has an effective order of 2k. Indeed, as shown above, a 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 autoregressive step. This twice greater Markov order implies that the AR approach is able to model a strictly larger class of data processes than the AAO sampling-see App. B.1. Second, in AAO the denoising process must ensure coherence between the start and end of the sequence. This requires information to be propagated between the two ends of the sequence. However, at each AAO denoising step, the score network has a limited receptive field of 2k + 1 which limits information propagation. In contrast, in AR sampling, there is no such limitation due to the nature of the autoregressive scheme. Scalability of the two methods is discussed in App. F.One can directly learn an approximation to the conditional score. Practically, conditioning on the previous states x i-C:i is achieved by feeding them to the score network along with the input as separate channels (Voleti et al., 2022) . 
experiments: 
 We parameterise the score network s θ with a modern U-net architecture (Ronneberger et al., 2015; Gupta & Brandstetter, 2023) with residual connections and layer normalization. For the full details of our experimental setup we will refer to App. D, where more results can be found.Data. In this experimental investigation, we look into the Burgers' and the Kuramoto-Sivashinsky (KS) equations since their dynamics are interesting yet challenging, and have been investigated in prior work (e.g. Zhang et al., 2019; Du & Zaki, 2021) . The Kuramoto-Sivashinsky is a fourth-order nonlinear one-dimensional PDE describing flame fronts, whose position is governed by ∂u ∂τ + u ∂u ∂z +∂ 2 u ∂z 2 + ν ∂ 4 u ∂z 4 = 0,where we used ν = 1. We refer to App. C and App. D for more details on the data generating processes and results on the Burgers' equation.Evaluation metrics. We use two sets of metrics to measure the accuracy of samples: (a) per time step-with the mean squared error MSE 1:L = E[(x 1:L -x1:L ) 2 ] and Pearson correlation ρ 1:L = Cov(x 1:L ,x 1:L ) Var(x 1:L )Var(x 1:L ) between model samples x1:L ∼ p θ and ground truth trajectories x 1:L ∼ p 0 , and (b) for the whole trajectory-via RSMD = 1 L L l=1 MSE l and high correlation time t max = l max × ∆t where l max = arg max l∈1:L {ρ l > 0.8}. Guided data assimilation. In this section, we compare the two guided diffusion sampling regimes, AAO and AR-introduced in Sec. 3with an amortized model on the data assimilation task. For each time step, we randomly sample a certain percentage of observed variables, i.e. these are sparse in space. We assume that we have some initial states that are fully observed.Our joint diffusion model is trained with a window of size 9 and in the AR case it generates 3 states at a time conditioning on the 6 previous ones. We vary the sparsity of observations, from almost fully observed to exclusively conditioning on the initial states, i.e. plain forecasting, and evaluate the different sampling regimes in terms of RMSD. We observe in Fig. 1  
discussion: 
 In this work, we have explored different ways to train and sample diffusion models for forecasting and data-assimilation tasks. Our proposed approach is a novel combination of existing techniques, such as learning a Markov score network (Rozet & Louppe, 2023b), amortising the score network (Song et al., 2021) using reconstruction guidance (Pokle et al., 2023) and autoregressive rollouts (Ho et al., 2022b) . The paper serves as a proof that the combination of these techniques succeeds in modelling long-range rollouts of PDEs. In particular, we have empirically demonstrated the long-range rollout ability of diffusion models trained on short trajectory segments, and autoregressively sampled.Limitations and future work Although we have explored different solvers and a number of discretisation steps, the autoregressive sampling strategy requires simulating a denoising process at each step, which is computationally costly. We believe this could be further improved by taking advantage of the advancements in diffusion model research. Additionally, the guided sampling strategy requires tuning the guidance strength, yet we suspect there exist some simple heuristics that are able to decrease the sensitivity to this hyperparameter. Finally, in this work we only explored linear observation operators A, yet reconstruction guidance can also be employed with non-linear observation processes. We will investigate this in future work. given the Markov assumption on the sequences. We then describe the data generating process in App. C for the datasets used in the experiments in Sec. 4. In App. D, we present additional experiments that, because of the space limit, did not fit in the main paper. After, in App. E, we assess the ability of the all-at-once sampling strategy-which 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. Then in App. F, we show that the same learnt local score can give rise to two different sampling strategies-all-at-once and autoregressive-and empirically assess their performance for forecasting and data assimilation. In App. G, we look at how the guided autoregressive sampling approach behaves when trading off the prediction horizon with the number of states being conditioned over. Then in App. H, we compare the amortised model with the guided approach in the context of forecasting, and investigate how the window size and the combination between prediction horizon H and number of conditioned frames C affect the performance of each model. In App. I, we empirically explore different guidance schedules and their impact on forecast accuracy. Finally, in App. J we examine how sensitive the amortised and guided models are to noise in the initial observed states.In this section we look at the structure of the score induced by a Markov time series. In particular we lay out different ways to parameterise and train this score with a neural network. In what follows we assume an 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 asEQUATIONB.1 WINDOW 2k + 1:Looking at the score of equation 5 for intermediary states we have∇ xi log p(x 1:L ) = ∇ xi log p(x i+1 |x i ) + ∇ xi log p(x i |x i-1 ) (6) = ∇ xi log p(x i+1 , x i |x i-1 ) = ∇ xi log p(x i+1 , x i |x i-1 ) + ∇ xi log p(x i-1 ) 0 = ∇ xi log p(x i+1 , x i , x i-1 ) = ∇ xi log p(x i-1 , x i , x i+1 ).(7) Similarly, for the first state we getEQUATIONand for the last stateEQUATION)Noised sequence Denoising diffusion requires learning the score of noised states. The sequence x 1:L is noised with the process defined by the forward equation introduced in Sec. 2 for an amount of time t: x 1:L (t)|x 1:L ∼ p t|0 , where p t|0 is the noising kernel. With x 1:L an AR-k Markov model there is no guarantee thatx 1:L (t) is still k-Markov. Yet, we can assume that x 1:L (t) is an AR-k ′ Markov model where k ′ > k, as in Rozet & Louppe (2023a).Although it 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. ( 7), ( 8) and ( 10) suggest training a score network to fits θ (t, x i-1 (t), x i (t), x i+1 (t)) ≈ ∇ log p t (x i-1 (t), x i (t), x i+1 (t)) by sampling tuples x i-1 , x i , x i+1 ∼ p 0 and noising them.Full score The first, intermediary and last elements of the score are then respectively given by• ∇ x1(t) log p t (x 1:L (t)) ≈ s θ (t, x 1 (t), x 2 (t), x 3 (t))[1] • ∇ xi(t) log p t (x 1:L (t)) ≈ s θ (t, x i-1 (t), x i (t), x i+1 (t))[2] • ∇ x L (t) log p t (x 1:L (t)) ≈ s θ (t, x L-2 (t), x L-1 (t), x L (t))[3]Joint sampling The full score over the sequence x 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 a segment of length 2k + 1 as input.Guided AR Alternatively for forecasting, having learnt a local score network s θ (t, x i-1 (t), x i (t), x i+1 (t)), i.e. modelling joints p(x i-1 , x i , x i+1 ), we can condition on x i-1 , x i to sample x i+1 leveraging reconstruction guidance-as described in Sec. 3. Effectively, this induces an AR-2-instead 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 k Markov assumption.Amortised AR Obviously we note that, looking at equation 6 and equation 9, one could learn a conditional score network amortised on the previous value such that s θ (t,x i (t)|x i-1 (0)) ≈ ∇ xi(t) log p(x i (t)|x i-1 (0)). This unsurprisingly allows forecasting with autoregressive rollout as in Sec. 3.2. Additionally learning a model over p(x 1 ) and summing pairwise the scores as in equation 6 would also allow for joint sampling.In App. B.1, we started with an AR-k model, but ended up with learning a local score taking 2k + 1 inputs. Below we derive an alternative score parameterisation. Still assuming an AR-1 model equation 5, we have that the score for intermediary states can be expressed asEQUATION)Local score From equation 11, we notice that we could learn a pairwise score network such that sθ (t,x i-1 (t), x i (t)) ≈ ∇ xi(t) log p(x i-1 (t), x i (t)), and sθ (t, x i (t)) ≈ ∇ xi(t) log p(x i (t)), trained by sampling tuples x i-1 , x i , x i+1 ∼ p 0 and noising them. Note that this requires 3 calls to the local score network instead of 1 as in App. B.1.Full score The first, intermediary and last elements of the score are then respectively given by∇ x1(t) log p(x 1:L (t)) = ∇ x1(t) log p(x 1 (t)) + ∇ x1(t) log p(x 2 (t)|x 1 (t)) = ∇ x1(t) log p(x 1 (t), x 2 (t)) ≜ sθ (t, x 1 (t), x 2 (t))[1]. ∇ xi(t) log p(x 1:L (t)) = s θ (t, x i-1 (t), x i (t), x i+1 (t)) ≜ sθ (t, x i (t), x i+1 (t))[1] + sθ (t, x i-1 (t), x i (t))[2] -sθ (t, x i (t)). ∇ x L (t) log p(x 1:L (t)) = ∇ x L (t) log p(x L (t)|x L-1 (t)) = ∇ x L (t) log p(x L (t), x L-1 (t)) = sθ (t, x L-1 (t), x L (t))[2].Joint sampling The full score over x 1:L (t) could therefore be reconstructed from these, enabling joint sampling.AR Additionally, autoregressive forecasting could be achieved with guidance by iterating over pairs Li et al. (2021) , where they consider (x, τ ) ∈ (0, 1) × (0, 1]. The dataset consists of 1200 trajectories of 101 timesteps, i.e. ∆τ = 0.01, and where the spatial discretisation used is 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 (Wang et al., 2021) using the code available in their public repository1 . Initial conditions are sampled from a Gaussian random field defined as N(0, 625 2 (-∆ + 5 2 I) -4 ) (although be aware that in the paper they state to be using N(0, 25 2 (-∆ + 5 2 I) -4 ), which is different from what they are doing in the code) and then they integrate the Burgers' equation using spectral methods up τ = 1 using a viscosity of 0.01. Specifically, they solve the equations using a spectral Fourier discretisation and a fourth-order stiff time-stepping scheme (ETDRK4) (Cox & Matthews, 2002) with step-size 10 -4 using the MATLAB Chebfun package (Driscoll et al., 2014) . Examples of training, validation, and test trajectories are shown in Fig. 3 .(x i-1 , x i ),We generate the groundtruth trajectories for the 1D Kuramoto-Sivashinsky equation Lippe et al. (2023) , which is based on the setup defined by Brandstetter et al. (2022) . In contrast to Lippe et al. (2023) , we generate the data by keeping ∆x and ∆τ fixed, i.e. 2022) solves the KS-equation 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. u 0 (x) = K k=1 A k sin(2πl k x/L + ϕ k ), where A k , l k , ϕ kk are random coefficients representing the amplitude of the different sin waves, the phase shift, and the space dependent frequency. The viscosity parameter is set to ν = 1. Data is initially generated with float64 precision and then converted to float32 for training the different models. We generated 1024 trajectories of 140∆τ , where ∆τ = 0.2s, as training set, while the validation and test set both contain 128 trajectories of length 640∆τ . The spatial domain is discretised into 256 evenly spaced points. Trajectories used to train and evaluate the models can be seen in Fig. 4 .∂u ∂τ + u ∂u ∂z + ∂ 2 u ∂z 2 + ν ∂ 4 u ∂z 4 = 0, following the setup ofIn this section we present additional experiments performed on the Burgers' equation dataset and additional results that did not fit into the main text of the paper. As we have already mentioned, we parameterise the score network s θ with a modern U-net architecture (Ronneberger et al., 2015; Gupta & Brandstetter, 2023) with residual connections and layer normalisation. The U-net we used closely follows the architecture used in (Lippe et al., 2023) . If not stated differently, we generate conditional samples by simulating the reverse diffusion ODE using the DPM solver (Lu et al., 2023) in 128 evenly spaced discretisation steps. As a final step, we return the posterior mean over the noise free data via Tweedie's formula. Unless specified otherwise, we set the guidance schedule to r 2 t = γσ 2 t /µ 2 t (Finzi et al., 2023; Rozet & Louppe, 2023a) , tune γ via grid-search and set σ y = 0.01.We now present the results obtained on the Burgers' equation dataset. The Burgers' equation is a second-order nonlinear PDE describing the motion of viscous fluid or gas with speed u given by ∂u ∂τ + u ∂u ∂z = ν ∂ 2 u ∂z 2 , where ν > 0 is the viscosity. In our experiments we consider ν = 0.01. As done in Sec. 4, we train a local score on contiguous segments of size 9 randomly sampled from training trajectories. We evaluate the model on test trajectories of length 101, which is the same length as the training set trajectories in the case of the Burgers' equation dataset. We assume we have access to the true noiseless initial observations for forecasting. Fig. 5 shows examples of sampled trajectories from both the guided and the amortised models. We can see that for the Burgers' equation both models are able to generate accurate forecasts.By comparing one of the best guided models (1 | 8) with one of the best amortised models (7 | 2) 3 , we can see that the former leads to better results in terms of MSE and correlation on the considered test trajectories. As in the KS dataset, we see that the guided model leads to better results when we condition on many previous states. On the other hand, the amortised model seems to perform better when the number of conditioned states C is small. Similar to what we discovered in the KS dataset, we can see that the guided model is more robust to the number of conditioning states C used at sampling time. Using one of worst performing H | C guided configurations still leads to good performance (dashed blue line in Fig. 5 Example of a sampled trajectory. In Fig. 6 , we present an example test trajectory and the corresponding trajectories generated by the guided and the amortised models when conditioned on the first six states. We observe that both models generate similar trajectories up to 60 seconds. After that, the error accumulation makes the trajectories diverge from the true one.Observation noise and guidance ablation. We further explore different variants of the amortised and guided AR models. Specifically, the reconstruction guidance term∇ x(t) log p(y|x(t)) ≈ 1 γr 2 t +σ 2 y ∇ x(t) ∥y -Ax(x(t))∥ 2 ,involves a choice of guidance schedule r 2 t ≈ Var[x(0)|x(t)], and a guidance scaling γ to be tuned. We explore several heuristics proposed in the literature: 3 The guided and amortised models showed similar performance for several configurations (within error bars). This is why we state that we show the results for one of the best configuration for each corresponding model.Value Error Guided Amortised Var[x(0)] = σ 2 0 and that the variance should be monotonically increasing with t. For SDA and ΠGDM, we have that r 0 = 0, whilst r 0 = 1 for Video diffusion and r t ---→ t→0 ∞ for DPS (assuming the error goes to 0). Also, for SDA and Video diffusion, r t ---→ t→∞ ∞ whilst r t ---→ t→∞ 1 for ΠGDM. These properties make the ΠGDM and SDA guidance schedules reasonable heuristics to choose, and indeed from Fig. 7 we observe that these perform best.For the amortised model, we also consider adding Gaussian noise η ∼ N(0, s 2 ) to the conditional frames during training, with a fixed noise s = 0.01. On Fig. 7 we see that noising the conditional information hurts the amortised model's performance, although one might think that it could improve robustness w.r.t. error accumulation. Also, App. J shows that the amortised model trained without noise is much more sensitive to noisy initial condition when sampling compared to GUARD and the amortised model trained with s = 0.01. One approach for generating long rollouts conditioned on some observations using score-based generative modelling is to directly train a joint posterior score s θ (t, x 1:L (t)|y) to approximate ∇ x 1:L (t) log p(x 1:L (t)|y). The posterior can be composed from a prior score and a likelihood term, just as in 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 significantly.To deal with this issue, in this section we develop the approach proposed by Rozet & Louppe (2023a), whereby the prior score (s θ (t, x 1:L (t))) is approximated with a series of local scores (s θ (t, x i-k:i+k (t))), which are easier to learn. The size of these local scores is in general significantly 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 (EI) discretization scheme (Zhang & Chen, 2023), while the second one uses the DPM++ solver (Lu et al., 2023) . For each strategy, the predictor steps can be followed by a few corrector steps (Langevin Monte Carlo) (Song et al., 2021; Mathieu et al., 2023) . For the results below, we consider the case of conditional generation, where we generate trajectories of 101 states, conditioned on a noisy version of the initial 6 states (with standard deviation σ y = 0.01). The results are computed for the Burgers' dataset, based on 30 test trajectories, each with different initial conditions. The sampling procedure uses 128 diffusion steps and a varying number of corrector steps.Influence of solver and time scheduling We show the results using two solvers: EI (Zhang & Chen, 2023) and DPM++ (Lu et al., 2023) . For each, we study two time scheduling settings, which determine the spacing between the time steps during the sampling process • Linear/uniform spacing:t i = t N -(t N -t 0 ) i N • Quadratic: t i = ( N -i N t 1/κ 0 + i N t 1/κ N ) κwhere N indicates the number of diffusion steps, t N and t 0 are chosen to be 1 and 10 -3 and κ is chosen to be 2. Fig. 8 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 a large number of corrections (25), the quadratic time schedule leads to lower RMSD. 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. When using no/only a few corrector steps, the RMSD is not highly impacted by the choice of solver and time schedule. However, at 25 corrector steps, the quadratic time schedule seems to lead to lower RMSD.Influence of corrector steps and window size Figs. 9 and 10 illustrate the effect of the window size and number of corrector steps on 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 be 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, but it comes with increased memory requirements. Unsurprisingly, the window 9 model performs better than the window 5 model (when using the same number of corrector steps). Performing corrector steps, alongside predictor steps, is crucial for preventing errors from accumulating. However, even with 25 corrector steps the performance is not very good. Finally, the findings from above are also confirmed in Fig. 11 , 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 a negative impact on the sampling time. Figure 11 : RMSD for Burgers' as a 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.Once the joint model s ϕ (t, x 1:L (t)|y) ≈ ∇ x 1:L (t) log p(x 1:L (t)|y) 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 as 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 H (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.1. Sampling time: the AAO procedure is 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 to generate the samples.• AAO:NFE = (1 + c) × p • AR: NFE = (1 + c) × p ×(1 + L-W H ) 2.Memory cost: the AAO procedure is more memory-intensive than the AR one, as the entire generated trajectory length needs to fit in the memory when batching is not performed. If 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 regenerates the states it conditioned upon, leading to computational redundancy.4. Quality of generated samples: the AR procedure manages to capture long-term temporal dependencies between states better than the AAO procedure. One of the main reasons likely has to do with the fact that AR sampling induces an AR-2k Markov model, as opposed to a AR-k model in the AAO case (see App. B).This section aims to illustrate a fair comparison between the AAO and AR sampling methodologies using the Burgers' dataset in the context of forecasting. As mentioned above, the AR method comes at an increased computational cost at sampling time, but generates samples of higher quality and manages to 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 p and no batching for AAO), we can use c = L-W H corrector steps.Window = 5 Fig. 12 shows the results for a trajectory with L = 101, using a model with W = 5.The results are expressed in terms of the RMSD, calculated based on 30 test trajectories. For AAO sampling c = 48, while for AR sampling c = 0. In the AR sampling, we generated 2 states at a time conditioned on the 3 previously generated states. In both cases we conditioned on a noisy version of the initial 3 true states (with standard deviation σ y = 0.01).For all experiments, we used the EI discretization scheme, with a quadratic time schedule. The AR procedure is clearly superior to the AAO, even with 48 corrector steps. This can be easily observed by comparing the generated samples to the ground truth as shown in Fig. 13 . Window = 9 For a window of 9, we used the following number of corrector steps for each conditioning scenario:• H = 2, C = 7 : c = 101-9 2 = 46 • H = 3, C = 6 : c = 101-9 3 = 31 • H = 4, C = 5 : c = 101-9 4 = 23 • H = 5, C = 4 : c = 101-9 5 = 19Moreover, the number of states we initially condition upon is equal to C (i.e., when comparing AAO with 46 corrections to AR with H = 2, C = 7 (2 | 7), we condition on the first 7 initial noised up states). Figure 14 : RMSD for Burgers' for AAO and AR generation with a window of 9 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 (H). This is also indicated on the x-axis in the format H | C (e.g., 2 | 7 implies we generate two states at a time and condition on 7). AR is clearly superior to AAO for forecasting.Once again, the AR-generated trajectories achieve significantly better metrics, which also reflects in the quality of the generated samples. Fig. 15 shows examples for AAO: 46/AR: 2 | 7.Performing data assimilation using the AAO sampling procedure simply corresponds to conditioning at generation time on all observed variables regardless of the time-step τ they appear at using reconstruction guidance. This is possible because in AAO all the states of the trajectory are generated simultaneously. In the AR case, instead, we sample a full trajectory in multiple rollout steps. a forecast step. If instead some observed variables are present in the predictive horizon, we have a likelihood for the conditioning states and one for the true observations within H. In this case, we also keep the resampled C states to ensure more coherency between nearby states. In this section, since we follow the data assimilation set-up from (Rozet & Louppe, 2023b), we have observations at each time step and therefore we never perform any plain forecast rollout in the data assimilation experiments.Results on Burgers' equation dataset In Sec. 4 we presented the results for the data assimilation task on the KS dataset. A similar behaviour can be also found when performing data assimilation on Burgers'. Burgers' equation, while still being non-linear, is less chaotic than the Kuramoto-Sivashinsky, resulting in an easier task for data assimilation when the setup is the same as the one presented in Sec. 4. This can be noticed in Fig. 16 , where the main performance differences between the considered approaches are mostly in the very sparse observation setting (percentage observed of up to 20%). In addition to that, the use of correction steps is not as important as in the KS dataset.Runtime comparison between AAO and AR for DA In Sec. 4, we presented results for the data-assimilation task on the KS equation dataset. In Fig. 17 , we report again the results considering but mostly focusing on the two guided approaches, AAO and AR. We compare them also in terms of their runtime. While using corrections is beneficial and needed for AAO, this leads to an increased runtime than the AR procedure.Impact of the guidance scaling in the data assimilation task The performance in the DA task seems to be slightly influenced by the value of the guidance scaling parameter, denoted with γ.To analyse this behaviour, we run the DA experiment on the KS dataset by using different scaling parameters. From Fig. 18 we can see that the value of γ significantly influences the results, with the most significant differences seen when performing AAO sampling with 0 corrections. In this setting, we found that using corrections was crucial for making the AAO approach more robust (i.e. less sensitive to the value of γ).The AR sampling scheme also seems to be impacted by the choice of γ. In particular, when the percentage of observed values is low (up to 20%), a smaller γ (higher guidance strength) seems to be preferred. In contrast, when a lot of observations are available, if the guidance strength is too high (i.e. γ is too small), the results get significantly worse. Therefore, it is clear that the best approach would be to fine-tune the guidance scaling parameter for each percentage of observed values. This opens up the possibility to investigate heuristics to automatically tune the scaling parameter depending on the number of observed values, but we leave this as future work.As already eluded to in App. F, the autoregressive sampling method offers some flexibility in terms of the number of states generated at each iteration (H). From a computational budget perspective, the higher H, the fewer iterations need to be performed. However, this might come at the cost of a decrease in the quality of the generated samples. In this section, we investigate this trade-off for the Burgers' and KS datasets.All results are shown for a window of 9, with varying predictive horizons H: 1, 2, 3, 4, 5, 6, 7, and 8. For each setting, the number of previously generated states we condition upon is C = W -H = 9 -H: 8, 7, 6, 5, 4, 3, 2, and 1. The trajectories are conditioned upon the first 9 -H true states.The results are presented in terms of the MSE and correlation (Fig. 19 ), which are computed based on 200 test trajectories for the Burgers dataset, and on 128 test trajectories for the KS dataset. The number of diffusion steps used is 128. For all experiments, we used the DPM++ solver, with a linear time discretisation scheme. At the end, we also performed an additional denoising step, and adopted the dynamic thresholding method for the Burgers' dataset.Fig. 20 also shows the RMSD for all the different studied settings for Burgers'. The performance is not very sensitive to H for low to medium values of H, although the model seems to perform1 | 8 2 | 7 3 | 6 4 | 5 5 | 4 6 | 3 7 | 2 8 | 1 H | C 20 30 RMSD 1 | 8 2 | 7 3 | 6 4 | 5 5 | 4 6 | 3 7 | 2 8 | 1 H | C 50 100High corr. time (s)Figure 22 : Experimental results on the Kuramoto-Sivashinsky equation. RMSD (left) and high correlation time (right) for different values of H. For RMSD, the smaller, the better, while the reverse is true for the high correlation time. Outliers have been removed.• by learning a diffusion model over the prior, and sampling from the conditional a posteriori.Conditioning is introduced through reconstruction guidance, hence we will call this model the guided model.We analyse two datasets: Burgers' and KS. For each dataset, the sampling in the guided model is performed autoregressively. We set the predictive horizon and the number of conditioning frames in such a way as to allow for a fair comparison between the amortised and the guided models.In this case we study the following setting:• Guided model with a window of 5, a prediction horizon H of 2 and 3 conditioned frames. Amortised model, where we generate 2 samples at a time conditioned on the previous 3.For the guided model, the conditioning from the first iteration uses a noisy version of the first three states (with standard deviation σ y = 0.01), while for the amortised model, we use the first three true states. In the sampling process we use 128 diffusion steps. For the discretisation scheme, we study both the EI scheme (Zhang & Chen, 2023), as well as DPM++ solver with order 1 (Lu et al., 2023) . The main difference between the two is that the first uses the ϵ prediction model, whereas the DPM++ solver uses the data prediction model. This allows us to also perform dynamic thresholding during the sampling process, which was shown to stabilise sampling with large guidance scales in the case of images (Ho et al., 2022a) . Moreover, we can optionally take a final denoising step, specified through the parameter denoise to zero. Finally, we study the two time discretisation schemes introduced in App. E: quadratic and linear. The results are determined based on 200 test trajectories.• Final denoising step: This was only studied for the DPM++ solver. The performance is clearly improved by using an extra denoising step at the end, as shown in Fig. 23 . This is true for all models. In the amortised case, using the final denoising step helps avoid the jump in MSE around the initial states.• Solver and time discretisation: The choice of solver and time discretisation schedule does not affect the results drastically, but it does lead to some differences. In the Burgers' dataset, the best setting seems to be to use the DPM++ solver with the linear time discretisation schedule when using 128 diffusion steps. However, this might depend on the dataset and on the number of diffusion steps used.• Guided vs. amortised: Fig. 23 clearly shows the that guided model outperforms the amortised one both in terms of MSE (lower MSE), as well as correlation (higher correlation). The amortised model is better at modelling the initial states, as it does not contain any observation noise (unlike the guided one). However, errors in modelling initial states propagate and make the amortised model diverge faster than the guided one. Thus, the amortised model suffers from the main drawbacks of autoregressive models -because the conditioning is done through previously generated states, errors propagate rapidly and lead to divergent states. This is true in the guided model as well, but to a lesser extent. Moreover, the guided model has two components: the prior joint score and the conditional information from reconstruction guidance. It is possible that the influence from the prior model also helps to avoid divergence. Results for model without denoise to zero A summary of these findings can be found in Fig. 24 , where we can easily see the positive impact of the final denoising step. Moreover, Fig. 24 shows that the guided model outperforms or performs on par with the amortised model in terms of RMSD. However, note that we only looked at the setting 2 | 3 in this analysis and the findings might not hold for other settings. In this case we study the following settings:• Guided model with a window of 5, a prediction horizon H of 2 and 3 conditioned frames. Amortised model, where we generate 2 samples at a time conditioned on the previous 3. • Guided model with a window of 9, a prediction horizon H of 3 and 6 conditioned frames.Amortised model, where we generate 3 samples at a time conditioned on the previous 6. • Varying window sizes (H and C) for guided and amortized models.For the KS dataset we did not use dynamic thresholding, but still used the final denoising step. We also just investigated the DPM++ solver with either the quadratic or linear time schedule. The results are derived on 128 test trajectories, for a trajectory length of 640 time steps, corresponding to 128s.As a summary metric, besides the RMSD, we also compute the high correlation time, defined as the time until the correlation drops to below 0.8 (the higher, the better).Fig. 25 shows that the guided model performs better than the amortised one on average, although the metrics of both models are within error bars. However, we only analysed the setting 2 | 3 and the findings might not hold for other settings. In particular, we observed that the guided model tends to perform better with more conditioning information, whereas increasing C harms the performance of the amortised model. The summary metrics over the entire trajectory are shown in Fig. 26 In this case, we generate 2 states conditioned on 3 (with a window of 5 for the guided model). For RMSD, the lower the value, the better the model, while for the high correlation time a larger value is preferred.Similar conclusions as in the previous section can be drawn if we generate 3 states at a time conditioned on the previous 6. Fig. 28 suggests that the difference between the two models is more pronounced in this setting, with the guided model outperforming the amortised one. 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 H the model is trained with a better learning task compared to H = 1, as larger H provides more information to the model. For H ≥ 4, the model performs similarly for any H, being slightly more sensitive to larger C when H is smaller. Regardless of H, the performance of the model deteriorates significantly with larger C = 8, similar to the behaviour shown in Lippe et al. (2023) . In contrast, the performance of the guided model is not that sensitive to the window size, as shown in Fig. 30 . The best performance is observed for medium window sizes (W = 9); if the window is too small, the effective Markov order assumed might be too low. As the window increases, so does the effective order, but there is probably a limit to how much increasing the effective Markov order can help performance. While the noised up states within the blanket might have infinite Markov order, the correlation between them likely shows a sharp decrease, so there is little gain between using a window of 9 or a window of 11. Moreover, increasing the window size might make the optimisation harder, explaining why we don't see improvement in performance beyond W = 9. EQUATIONSec. 4 introduces the four types of guidance studied in this paper. This section provides a comparison between these guidance formulations, applied on the Burgers' and KS datasets in the context of forecasting. The results are derived using AR sampling, using models with a window of 9, with H = 1 and C = 8 (1 | 8). Additionaly, we also condition on the first initial 8 states. In all experiments from this section 128 diffusion steps were employed, and the solver used was DPM++ with a linear time discretisation schedule. At the end, one final denoising step was taken. For both datasets, we also studied the influence of dynamic thresholding, focusing on the stability of the methods at high guidance strenghts.We analysed several settings of the γ hyperparameter: [0.01, 0.03, 0.05, 0.1, 0.5, 1] for SDA, [0.001, 0.01, 0.03, 0.1, 1] for DPS, [0.0002, 0.001, 0.01, 0.2, 1.0, 10.0] for VideoDiff, and [0.02, 0.1, 0.2, 2.0, 10.0, 100.0] for ΠGDM. Note that we have studied different γ ranges -for each guidance mechanism the object γ is multiplied by has different scales, so the most appropriate γ values will vary from one mechanism to another.The results in Fig. 31 correspond to the best setting for each guidance type, indicating that in terms of RMSD, SDA has the best performance, followed by ΠGDM, DPS, and then Video diffusion. Outliers have been removed. SDA and ΠGDM show the best results, with a similar performance in terms of RMSD, followed by DPS and then Video diffusion.Stability of each guidance type Fig. 32 shows the RMSD for each guidance type when applied to the Burgers' data, both with and without using dynamic thresholding.By comparing the four plots we can see that• The lowest RMSD is achieved when using dynamic thresholding for all four types of guidance. However, SDA is the only guidance mechanism that is able to achieve similar performance (within error bounds) without the use of dynamic thresholding. For all the other three guidance types, not using dynamic thresholding resulted in a significant increase in RMSD for the best γ setting.• The extent of the increase in RMSD resulting from not employing dynamic thresholding (i.e., lowest RMSD achieved with vs without dynamic thresholding) varies across guidance mechanisms -SDA: negligible increase; -DPS: significant, but relatively small increase; Averaged across all test samples, the mean RMSD increased more than five times (from 0.8 × 10 -2 to 4.5 × 10 -2 );-Video diffusion: big increase in RMSD, with the average across the 200 test samples increasing by almost two orders of magnitude;-ΠGDM: big increase in RMSD, with the average across the 200 test samples increasing by more than two orders of magnitude.• For the investigated guidance strengths, DPS is the only guidance mechanism that did not lead to numerical overflow. For high guidance strengths the results were highly inaccurate, but not unstable. We hypothesise this is due to the ∥y -Ax(x(t))∥ term that probably promotes stability.These findings suggest that while the SDA guidance leads to similar results regardless of whether other tricks are employed (such as dynamic thresholding), the performance of other guidance types is more sensitive to them. In particular, dynamic thresholding was introduced to stabilise sampling at high guidance strengths (low γ), and it seems to be crucial for all guidance types but SDA. Note that a lower γ implies a higher guidance strength. Also note that the scale of the x-and y-axes is logarithmic.For each γ, if the plot does not show a result for the case in which dynamic thresholding was not employed (green X), this is because the results were unstable. The models used in Sec. 4 were trained on the original KS data (i.e., no pre-processing was applied).Given that the training data values lie in the range [-3.61, 3.56] , dynamic thresholding could not be applied for this model. One requirement for dynamic threhsolding to work is for the data values to lie within [-1, 1] . Thus, we investigate the performance of each guidance type without using dynamic thresholding.Fig. 34 shows that SDA has a significantly better performance as compared to all other guidance types. This is in line with the findings from the Burgers' dataset, where we found dynamic thresholding to be crucial for DPS, Video diffusion, and ΠGDM to achieve a similar performance to SDA.among the best two, and not significantly different (i.e., within error bounds). Video diffusion was the only guidance type where the best γ setting was dataset-dependent. The same conclusion can be drawn from Fig. 36 , where we show that employing dynamic thresholding does not improve the RMSD significantly for SDA, but it has a crucial effect for the other three guidance mechanisms. For each γ, if the plot does not show a result for the case in which dynamic thresholding was not employed (green X), this is because the results were unstable.Finally, one test trajectory for the best setting of each guidance mechanism is shown in Fig. 37 .In this section we study the robustness of the guided AR and amortised AR models to noise in the initial observations in the context of forecasting. More specifically, for each model we analyse its forecasting performance (in terms of high correlation time) in the following two settings 1. Condition on the true initial observations. 2. Condition on noised initial observations, with noise with standard deviation σ y = 0.01.The models we analyse are the same as the ones in Sec. 4. For the amortised case, we study both a model that was trained on noiseless conditional frames (s = 0), as well as one that was trained on noisy frames with fixed noise s = 0.01, with the hypothesis being that the latter is more robust to noise in the initial observations.1. Guided AR model, generating 1 state at a time and conditioning on 8 (1|8). We also condition on the initial 8 true or noised up states. 2. Amortised AR models, generating 7 states at a time and conditioning on 2 (7|2). We also condition on the initial 2 true or noised up states.• Trained on noiseless conditional frames s = 0 • Trained on noisy conditional frames s = 0.01We used the DPM++ solver with a linear time discretisation schedule for all models from this section. We performed 128 diffusion steps with no Langevin corrector steps, and performed a final denoising step, as discussed in Sec. 4.The results on the KS dataset are shown in Fig. 38 . We can see that the guided model (using the SDA guiding mechanism) is very robust to noise in the initial observations, showing only a slight decrease in high correlation time when we condition on noisy initial states, as opposed to the true ones (less than 2s). In comparison, the performance of the amortised model that was trained on noiseless conditional frames (s = 0) suffers more, with a decrease in high correlation time of more than 7s (see right-most green bars in Fig. 38 ). We can also equip the amortised model with a similar robustness to noise in the initial states by training it on noisy conditioned frames (see left-most green bars in Fig. 38 ). This strategy clearly leads to less sensitivity to the amount of noise present in the initial observations, with a decrease in high correlation time of less than 2s. However, although more robust to noise, it is significantly inferior to the original amortised model (trained on noiseless conditioned frames), managing to achieve a high correlation time of only about 50s.Thus, it seems that the guided model is inherently less sensitive to the amount of noise present in the initial states. While there are strategies to make the amortised model similarly robust to noise, they significantly harm performance. i p(x i |x i-C:i-1 ) with C ≤ k. As suggested inHo et al. (2022b), it then follows that we can samplehttps://github.com/PredictiveIntelligenceLab/Physics-informed-DeepONets/ tree/main/Burger/DataThis just requires changing lines 166-171 of the generate data.py file in the public repository ofBrandstetter et al. (2022). 
 
 Appendix:  


Decision: This paper presents  the Guided AutoRegressive Diffusion model (GuARD), which is able to generate accurate predictions of long PDE trajectories. The concerns of clarification raised by both reviewers are valid. Those concerns can be addressed in the camera ready version.

Reviews: review: title: title: The authors present interesting experiments and results, however due to the denseness of the content several aspects remain unclear.
main_review: main_review: The paper introduces GUARD, a diffusion model which is trained on short PDE trajectories and aims to forecast longer trajectories and perform data assimilation given sparse spatial data. 

Pros 
- The approach and comparisons are interesting, experiments are well-described. 
- The results presented seem promising (especially for forecasting) and should be investigated further.


Quality: The quality of the paper is good. \
Clarity: While the paper is written clearly it sometimes lacks in logic as crucial parts are omitted/in the appendix. \
Originality: The basis of GUARD is AR (Ho et al.) and the distinction to this work remains unclear. Due to a lack of conclusion it is also not clear what exactly the authors claim to be new.  \
Significance of this work: The results presented are promising. However, to determine the significance a more elaborate benchmark would be needed. 

Major issues
- The paper lacks a discussion and conclusion. Overall the background section and description of the setup take up most space of the paper and sections on the experiments performed and their results remain very short. Therefore, many things seem to be omitted and the appendix is referenced often. 

- The experiments performed for data assimilation and forecasting are not comparable to each other. It remains unclear why for data assimilation only the two approaches AAO and AR are compared while for forecasting GUARD is compared to an amortized model. Also, the amortized model is not described and there's a lack of further benchmarks.


Minor issues 
- Math equations are difficult to read when inline. 
- Error metrics: Why are some evaluations performed using MSE and others using RMSD? 2a vs 2c. 
- Runtime comparison of GUARD and amortized model?
- Conditioned on less states the amortized model performs better than GUARD (Figure 2). Why? Could this be a reason to use the amortized model in some cases and GUARD in others?
title: title: Exhaustive evaluation of a marginally superior method
main_review: main_review: ### Summary:
A guided autoregressive diffusion model is introduced to forecast and assimilate PDE dynamics. The model makes use of conditional frames that provide context of the problem and can be trained on short slices of long sequences.

### Strengths:
- Very appealing to see that GuARD benefits from more context frames that are provided (i.e., using higher Markov orders). This actually matches people's intuition, but seems not to be the case for Amortized.
- Detailed information in the appendix (I did not check it thoroughly)

### Weaknesses
- Despite the concrete second paragraph of the introduction, I could not quite extract the contribution and novelty of this work. Can you elaborate more specifically what is new and in what way?
- GuARD only marginally improves over the standard Amortized technique. Nevertheless, GuARD is faster and slightly better.

### Questions:
1. Is $x(0)$ the noise-free state? Maybe make more explicit in the **Continuous-time diffusion models** paragraph.
2. How it the density $p_t$ of $x(t)$ calculated, or what does it concretely represent?
3. How do Langevin Monte Carlo (corrector) steps look like?
4. In Figure 2, bottom right, why is the Amortized AR model worse if conditioning on many frames an producing a single frame (1|8) than when conditioning on a single frame and producing many frames (8|1)? At the same time, why is GuARD worse than Amortized in the (7|2) configuration?

### Minor comments:
- Equal sign missing before $\Pi$ in second line of page 3?