Bias-Constrained Diffusion Schedules for PDE Emulations: Reconstruction Error Minimization and Efficient Unrolled Training Constantin Le Cleı̈ 1 Nils Thürey 1 Xiaoxiang Zhu 1
arXiv:2604.08357v1 [cs.LG] 9 Apr 2026
Abstract
estimates (Liu and Thuerey, 2024), can be used flexibly across tasks (Shysheya et al., 2024) for super-resolution and forecasting, and can sample multiple consistent scenarios (Price et al., 2023)
Conditional Diffusion Models are powerful surrogates for emulating complex spatiotemporal dynamics, yet they often fail to match the accuracy of deterministic neural emulators for highprecision tasks. In this work, we address two critical limitations of autoregressive PDE diffusion models: their sub-optimal single-step accuracy and the prohibitive computational cost of unrolled training. First, we characterize the relationship between the noise schedule, the reconstruction error reduction rate and the diffusion exposure bias, demonstrating that standard schedules lead to suboptimal reconstruction error. Leveraging this insight, we propose an Adaptive Noise Schedule framework that minimizes inference reconstruction error by dynamically constraining the model’s exposure bias. We further show that this optimized schedule enables a fast Proxy Unrolled Training method to stabilize long-term rollouts without the cost of full Markov Chain sampling. Both proposed methods enable significant improvements in short-term accuracy and long-term stability over diffusion and deterministic baselines on diverse benchmarks, including forced Navier-Stokes, Kuramoto-Sivashinsky and Transonic Flow.
However, Diffusion Models typically lag behind deterministic baselines in terms of pure reconstruction accuracy, which can lead to biased estimates despite their probabilistic formulation. Furthermore, because their sampling process requires tens to hundreds of function evaluations, standard stabilization techniques, such as multi-step unrolled training, are computationally prohibitive to apply directly. In this work, we argue that these limitations are intrinsically linked through the phenomenon of Diffusion Exposure-Bias (Li et al., 2023). This refers to the performance drop caused by the discrepancy between the ground-truth noisy targets seen during training and the estimated noisy states generated via ancestral sampling at inference. In the case of spatio-temporal forecasting, we demonstrate that standard conditional diffusion models suffer from this bias and obtain a sub-optimal error, mainly because their noise schedules are not designed to align the rate of error reduction with the intrinsic model capacity and task difficulty. In particular, we make the following contributions: • We define the Reconstruction Exposure-Bias, a special case of exposure-bias, which denotes the discrepancy between intermediate reconstruction errors during inference and training. We experimentally show that the main driver of this bias is the rate of decrease of the reconstruction error as a function of the noise level.
1. Introduction Machine learning-based emulators have achieved performance comparable to traditional numerical solvers for fluid dynamics tasks, operating at a fraction of the computational cost. Among these, Diffusion Models have recently demonstrated strong performance while resolving key limitations of deterministic approaches, such as instability (Kohl et al., 2023; Rühling Cachay et al., 2023; Lippe et al., 2023), and over-smoothing over long simulation horizons (Liu et al., 2025). Beyond stability, they allow to obtain uncertainty
• We propose a novel Adaptive Scheduling algorithm that treats schedule design as a final noise-level minimization problem under the constraint that the model stays within stability at every denoising step, thus jointly optimizing the Final Reconstruction Error and the Reconstruction Exposure-Bias • Connection between the exposure-biases : Reducing the diffusion exposure-bias naturally leads to a fast indistribution proxy for the full diffusion process, only requiring a few steps to provide an accurate sample which can be used for Unrolled training, therefore re-
1 Technical University of Munich, Munich, Germany. Correspondence to: Constantin Le Cleı̈ <[email protected]>.
Preprint. April 10, 2026.
1
Bias-Constrained Diffusion Schedules for PDE Emulations
ducing simulation exposure-bias.
developed to mitigate this effect, including scaling the predicted noise (Ning et al., 2023), shifting sampling timesteps (Li et al., 2023) or adding a regularization term in the diffusion loss (Daras et al., 2023). For PDE simulations however, the target posterior is highly concentrated, unlike the broad distributions typical of unconditional image generation. Consequently, this regime requires a tailored framework to address its specific constraints.
• Our proposed adaptive schedule systematically reduces first-step reconstruction error, while the fast proxy unrolled-training drastically mitigates artifacting effects, leading to an improvement multiple orders of magnitude in Fréchet Spectral Distance on Kolmogorov turbulent flow (see Table 1).
External Loss terms and Unrolled Training Constraining the output of diffusion models has become an active area of research, particularly through guidance mechanisms (Bansal et al., 2023; Ho and Salimans, 2022). In the context of trajectory generation, these constraints serve a dual purpose: facilitating accurate autoregressive predictions and improving robustness to error accumulation during unrolled training. In fluid dynamics tasks, physical consistency losses are typically enforced by approximating the clean output via DDIM shortcuts (Bastek et al., 2024) or linear interpolation (Amorós-Trepat et al., 2026). However, such approximations are insufficient for unrolled training; they lack the fidelity required to mimic the model’s actual inference behavior. Consequently, the model fails to learn how to correct its own generated artifacts, potentially leading to gradient inaccuracies. Specialized strategies to stabilize autoregressive diffusion have also emerged : (Chen et al., 2024) condition the model on noise-corrupted histories, while (Huang et al., 2025) relies on cached prior predictions rather than groundtruth. Alternative strategies for estimating the final output include Consistency Models (Song et al., 2023; Stock et al., 2025) or diffusion shortcuts (Shehata et al., 2025).
2. Related Work Diffusion for fluid flows Diffusion models (Ho et al., 2020; Sohl-Dickstein et al., 2015) have been shown to perform well for autoregressively generating videos (Ho et al., 2022), naturally expanding to other spatio-temporal tasks such as weather forecasting (Price et al., 2023) or time-series (Shen et al., 2024). In the context of fluid dynamics, there have been multiple applications : (Kohl et al., 2023) demonstrated that diffusion models can predict fluid states over extended horizons while preserving both sample quality and temporal stability. For fluid-flow reconstruction, (Shu et al., 2023) leveraged DDPMs for turbulent flow super-resolution, accurately reconstructing high-fidelity fields from low-resolution inputs, while (Rozet and Louppe, 2023) relied on score-based generative models for reconstruction. For 3D turbulent flows, (Lienen et al.) achieved fast spatio-temporal prediction. (Liu et al., 2025) merged the time and diffusion axis to produce fast physically plausible samples on very long horizons. Diffusion Hyperparameters Noise schedules are a central component of diffusion models, determining the rate of information destruction and the weighting of the learning objective (Kingma et al., 2021). The design of these schedules has evolved from heuristic linear and cosine formulations (Ho et al., 2020; Nichol and Dhariwal, 2021) to principled parameterizations based on the Signal-to-Noise Ratio (SNR) (Kingma et al., 2021; Choi et al., 2022) and the continuous noise level formulations of the EDM framework (Karras et al., 2022). However, these schedules are predominantly optimized for perceptual fidelity, prioritizing regimes where visual features emerge. This focus is ill-suited for high-precision fluid dynamics tasks, where the objective is not only to generate physically plausible fields, but also to remain tightly correlated with a specific groundtruth trajectory. In particular, (Lippe et al., 2023) has shown that schedules that focus on very low noise-levels tend to perform well on turbulence tasks.
3. Schedule, Reconstruction Error, and Exposure-Bias 3.1. Background: Conditional Diffusion Models for Autoregressive Fluid Simulations For Fluid Dynamics tasks, our goal is to construct a neural operator Mθ that emulates the ground-truth time evolution of a fluid. Given an initial condition x0 , we approximate the subsequent trajectory {x1 , . . . , xK } via autoregressive estimates: x̂1 = Mθ (x0 ), k
k−1
x̂ = Mθ (x̂
(1) )
∀k ∈ {2, . . . , K}.
(2)
Conditional Denoising Diffusion Probabilistic Models (CDDPMs) learn the conditional distribution p(xk+1 | xk ) by reversing a gradual noise-addition process. The sampling operator Mθ iteratively refines a noise sample yT ∼ N (0, I) into a prediction x̂k+1 conditioned on the previous state xk . A broader introduction can be found in Appendix A. To simplify notation for the remainder of this section, we
Exposure bias in diffusion models. This phenomenon arises from the mismatch between training, where the model sees ground-truth noisy data, and sampling, where it denoises its own predictions, causing errors to accumulate along the reverse trajectory. Different method have been 2
Bias-Constrained Diffusion Schedules for PDE Emulations
denote the condition as x and the target as y. We adopt the ϵ-prediction formulation for the diffusion process, where at each step a denoiser ϵθ predicts the Gaussian noise added to the target y = y0 . The reverse process starts by sampling ŷT ∼ N (0, I), and the following denoising steps are given by: q ŷt−1 = µ̃t (ŷt , ϵθ (ŷt , x, σt )) + β̃t ϵt−1 , (3)
model and a pair (x, y), and omit these dependencies from the notation; all quantities are implicitly functions of (x, y), and θ is reintroduced from Proposition 3.4 onwards where model comparisons are needed. To analyze how the schedule influences (9), we first define the concept of reconstruction exposure-bias: Definition 3.1 (Clean-Input vs. Inference-Input Error). The Clean-Input Error measures prediction error when the state input is the noised ground-truth state ỹt :
√ where t ∈ {1, . . . , T }, σt ≜ 1 − ᾱt is the noise level at time t, ϵt ∼ N (0, I), and the posterior mean and variance are: 1 βt µ̃t (yt , ϵθ ) = √ yt − ϵθ , (4) αt σt 1 − ᾱt−1 β̃t = βt . (5) 1 − ᾱt
E clean (t) = Eϵ ∥yest (ỹt ) − y∥ The Inference-Input Error measures error when the state input is the ancestor-sampling state ŷt : E inf (t) = Eϵ [∥yest (ŷt ) − y∥] . The divergence between these two via the Reconstruction Exposure-Bias (REB):
The model is trained to minimize the reconstruction loss: Ldiff (θ) = Ex,y,t,ϵ ∥ϵθ (ỹt , x, σt ) − ϵ∥2 , (6) √ where ϵ ∼ N (0, I) and ỹt = ᾱt y + σt ϵ is the noisy ground-truth target. Given any noisy input yt at noise level σt , the neural estimate of the clean signal is: yt − σt ϵθ (yt , x, σt ) √ yest (yt , x, σt ) = . ᾱt
REB(t) ≜
E inf (t) E clean (t)
3.3. Decomposing the REB: Two-Steps Bias as the Primary Driver
(7)
We now analyze the sources of REB by introducing the Two-Steps Bias, isolating the degradation that occurs in a single transition.
During training, this is evaluated on the ground-truth noisy target ỹt ; during inference, on the sampling iterate ŷt . Throughout this work, we use the following simplified onestep transition as a practical approximation to the full DDPM posterior (3): √ ŷt−1 = ᾱt−1 yest (ŷt , x, σt ) + σt−1 ϵt (8)
Definition 3.2 (Two-Steps Bias). Define the two-steps noisy state at time t as renoising the estimate of y obtained from a noised ground-truth state as input at time t + 1: (2S)
ŷt
i.e., the forward noising process applied to the predicted clean signal, bypassing the explicit dependence on ŷt in the mean. In the rest of this work, we will discard the dependency of x and σt in yest (·) to improve readability.
≜
√ ᾱt yest (ỹt+1 ) + σt ϵ,
(10)
The Two-Steps Bias is defined as the ratio between the error of the second estimate and the clean-input error at time t: (2S)
B (2S) (t) ≜
3.2. Problem Definition: Minimizing Reconstruction Error
Eϵ ∥yest (ŷt ) − y∥ E clean (t)
(11)
The following proposition shows that the REB is dominated by the individual two-steps biases.
While metrics like FID assess distributional quality, precision tasks such as PDE modeling require minimizing the specific trajectory error relative to the ground truth. In this setup, we aim to identify the schedule that leads to the optimal reconstruction error. We define the Optimal-Error Schedule S ∗ as the configuration that minimizes the final inference reconstruction error: S ∗ = argmin Ex,y,ϵ ∥Mθ∗ (S) (x) − y∥ , (9)
Proposition 3.3 (Re-noising Attenuation). When estimated errors are nearly aligned, the following recursive bound holds: REB(t) ≲ B (2S) (t) + λt REB(t + 1),
√ (ỹt+1 )−y∥ where λt := ∥Jt ∥ · ᾱt · ∥y∥yest . The REB is thus est (ỹt )−y∥ driven by the local two-step bias, and attenuated (resp. amplified) across steps when λt < 1 (resp. λt > 1). A full bound without the alignment assumption is given in Appendix B.
{S}
where θ∗ (S) denotes the model parameters converged via (6) under schedule S. In the following, we fix a trained 3
Bias-Constrained Diffusion Schedules for PDE Emulations
The proof is given in Appendix B, along with a visualization of the impact of the two-steps bias on the total REB for our practical datasets. This motivates enforcing B (2S) (t) ≤ τ at each timestep as a practical proxy for controlling the full REB. We further decompose the two-steps bias into two distinct contributions:
100
(12)
1. The Own-Prediction Bias B (own) measures the error amplification when the model is fed its own prediction instead of the ground-truth noisy input: √ (own) (13) ŷt ≜ ᾱt yest (ỹt ) + σt ϵ, (own) Eϵ ∥yest (ŷt ) − y∥ B (own) (t) ≜ clean E (t)
γ(σ , 1.1) γ(σ , 1.15)
10−2
10−3
1.8×10−1
10×10−1
Noise Level σ
(14)
KS
Stability Thresholds γ(σ , 1.01)
Clean-Input Error Eclean
γ(σ , 1.03)
2. The Propagation Residual δ(t) > 0 denotes the residual error due to additional error inherited from the step t + 1.
Minimizing B (own) is a necessary condition for stability. The following result establishes that own-prediction bias directly depends on clean-input error, and that for each noise level there exists a natural stability boundary. Proposition 3.4 (Stability Threshold). Assuming (A1) Wiener denoiser structure and (A2) Spectral bias of the (own) neural denoiser (detailed in Appendix C), Bθ (t) is an clean increasing function of Eθ (t). In particular, at each noise level σt and for any threshold τ ≥ 1, there exists a critical clean-input error γ(σt , τ ), increasing in σt and decreasing (own) in τ , such that if Eθclean (σt ) ≤ γ(σt , τ ) then Bθ (t) ≤ τ .
γ(σ , 1.05)
10−6
γ(σ , 1.1)
10−7
3.2×10−3
5×10−2
Noise Level σ
Figure 1. Stability Thresholds. We train a diffusion model on a log-uniform schedule. For a given training checkpoint, if the own-prediction bias at a given noise level falls in the τ ± 0.05 box, where τ ∈ [1.01, .., 1.15], we report its clean-input error. One sees that Clean-Input Error and Own-Prediction Bias correlate on both datasets. The red crosses correspond to the smallest noise-level for which the model achieves Ex,y B(own) ≤ τ at convergence, indicating that there is a – data and model-dependent – minimal noise-level where the model can be bias-free.
The proof is given in Appendix C. In other words, at each noise level, there is a minimal clean-input error to be achieved in order for the model to be stable. (A1) does not hold exactly in practice, and (A2) may not be true for all classes of neural networks, however we display the empirical correlation between the clean-input error and the own-prediction bias for a diffusion model in Figure 1, on different benchmark datasets (see Section 6.1 for setup).
The reason is that a model with large REB drifts outside of distribution, leading to unpredictable behaviour and making sub-optimal allocation of model capacity (reducing cleaninput error on noise-levels where it cannot be reduced during inference). We therefore expect that the true optimal schedule falls in this category.
3.4. Schedules control Clean-Input Error decrease rate
As illustrated in Figure 2, the schedule governs the decay rate of the reconstruction error, directly influencing the final prediction. In the absence of any stability constraint, one could simply set σ0 → 0, making the final denoising task trivially easy (E clean (0) → 0). However, this degenerate schedule doesn’t take into account that a larger input error is propagated during inference compared to training. Consequently, ŷ0 diverges from ỹ0 , leading to high REB. Using the stability threshold, we can formalize this:
In the rest of this work, we will restrict the schedule optimization problem to schedules that are stable. Given a small τ , we therefore restrict (9) to identifying the Optimal Stable Schedule: S ∗ = argmin Ex,y Eθclean (15) ∗ (S) (0) {S}
∀t < T,
γ(σ , 1.05)
10−1
10−4
A model is stable at a step t when B (own) (t) ≈ 1, meaning that repeated denoise-renoise iterations introduce no drift from the noised ground-truth state prediction. Conversely, B (own) (t) ≫ 1 signals instability.
s.t.
γ(σ , 1.01) γ(σ , 1.03)
Clean-Input Error Eclean
B (2S) (t) = B (own) (t) + δ(t),
Kolmo
Stability Thresholds
(2S) Ex,y Bθ∗ (S) (t) ≤ τ, 4
MSE with Distribution target Schedule Distributio Schedule
MSE own / with cleantarget
LowLevels-Focus Linear Cubic
101 Bias-Constrained Diffusion Schedules for PDE Emulations
100 1e 02 1e 03 1e 05
4. Bias-Constrained Schedule Construction via Error Constraints
Sampling direction, T = 100 Clean Input LowLevels-Focus Inference Input Linear Cubic
Based on decomposition (12), a stable schedule must ensure that the sum of the own-prediction bias B (own) and the propagation bias δ remain bounded, yielding the following principles:
101
100 1e 02 1e 03 1e 06 1e 05
Clean Input (0)= 1.2e-06 clean Inference Input inf(0)= 1.3e-06
clean(0)= 2.6e-06 inf(0)= 3.1e-06
• At every σt in S, clean-input error must be reduced enough to reach stability, i.e. B (own) (t) ≤ τ .
clean(0)= 1.7e-06 inf(0)= 2.2e-06
• The step size σt+1 − σt must be constrained to limit the propagation bias such that B (2S) (t) ≤ τ .
1.15
• Efficient Capacity Allocation: Driving the two-steps bias well below the stability boundary for a given noiselevel is wasteful. Capacity should be conserved for lower noise levels.
1.10 1.05 1.00 1e 06
clean(0)= 1.2e-06 inf(0)= 1.3e-06
clean(0)= 2.6e-06 inf(0)= 3.1e-06 1
10
Noise-level
clean(0)= 1.7e-06 inf(0)= 2.2e-06
100
4.1. Adaptive Scheduling Algorithm
1.15 Figure 2. Effect of schedule on final reconstruction error: Clean 1.10 and Inference Input Errors during training. The diffusion sampling direction goes from right (high noise-levels) to left (low noise1.05 levels). Schedules with more weight on low noise-levels obtain a better reconstruction performance early in the denoising process. 1.00 However they are then unable to reduce the error further compared to schedules that put more focus on low noise-levels. At the end0 of 10 1 sampling, schedules with higher clean-input reconstruction 10 error Noise-level start suffering from REB.
own / clean
Our algorithm contains two stages : an exploration stage, where the stability thresholds are discovered for the task at hand; and a schedule construction phase, where we greedily pick the noise levels such that the 2-steps bias remains bounded by τ at each step. While the ideal τ is 1, a value of τ slightly above 1 allows more flexibility in constructing the schedule and is shown experimentally to still lead to no error accumulation, potentially due to error correction at next step.
Proposition 3.5 (Slow Error Decrease Principle). Under the stability condition of Proposition 3.4, a finite capacity assumption (A1), and assuming Eθclean (σ) is strictly decreasing in σ (A2), the schedule minimizing (15) contains no unnecessary noise levels: every intermediate σk is necessary in the sense that removing it would violate the bias constraint, i.e. lead to B (2S) (k − 1) > τ . In particular, the greedy construction that always takes the largest feasible jump is optimal.
Phase 1: Exploration of Stability Thresholds. We initialize a dense log-uniform exploration schedule σ exp = {σ (1) , . . . , σ (N ) } covering [σmin , σT ] with uniform weight. We train the model on this schedule, periodically evaluating B (own) (i) for all active levels. Any level satisfying B (own) (i) ≤ τ is solved: we save the corresponding checkpoint θ∗ (σ (i) ) and remove σ (i) from the active schedule. Because solved levels are progressively removed, the active schedule shrinks continuously, making later epochs strictly cheaper than earlier ones. The exploration terminates once no new level is solved in a pass, ensuring the total epoch count does not exceed E (the number of epochs required to train a single baseline model to convergence). This yields a set of checkpoints {θ∗ (σ)}σ∈σsolved .
This follows from the observation that any unnecessary intermediate step represents wasted capacity that could be reallocated to reduce Eθclean (0) (a formal proof is given in Appendix D). Note that the bias constraint need not be tight at each step — it suffices that no larger jump is feasible. Since satisfying the bias constraint is increasingly easy at higher noise levels (as γ(σ, τ ) increases with σ), the schedule is naturally coarse at high noise and fine near σ0 . This motivates reformulating (15) as minimizing the final noise level σ0 subject to the bias constraint: min σ0
s.t.
∀t < T,
(2S)
Ex,y Bθ
Phase 2: Bias-Constrained Schedule Construction. Given the solved checkpoints, we construct the final schedule greedily, by setting σ0 = min σ solved and proceeding forward: n o (2S) σt+1 = max σ ′ ∈ σ solved : Bθ∗ (σ′ ),θ∗ (σt ) (σt , σ ′ ) ≤ τ , (17)
(x, y, t) ≤ τ (16)
(2S)
where Bθ∗ (σ′ ),θ∗ (σt ) (σt , σ ′ ) denotes the two-steps bias evaluated using checkpoint θ∗ (σ ′ ) for the first denoising step 5
Bias-Constrained Diffusion Schedules for PDE Emulations
and θ∗ (σt ) for the second. At each step we take the largest feasible jump, directly implementing the slow decrease principle (Proposition 3.5). This ensures (1) a schedule with the minimal number of diffusion steps and (2) no wasted model capacity on unnecessary intermediate steps. The number of steps T emerges naturally from the construction. This phase requires only forward passes to evaluate two-steps biases between checkpoint pairs (no gradient computation) and is therefore computationally negligible. The algorithm pseudo-code can be found in Appendix G. The model is trained from scratch on the constructed schedule using the standard diffusion loss (6).
Faithfulness of the proxy estimate. If the model does not suffer from exposure-bias up to step n, the distribution ∼ pθ · | xkn , z of latents obtained during sampling x̂k+1 n matches the ground-truth forward distribution x̃k+1 ∼ n (n) q · | xk+1 , and therefore the distribution of Pθ matches that of Mθ . We give an experimental visualization of the accuracy of the gradients in Appendix E. We consequently define the 2-steps proxy unrolled-training loss as: LP-UT (θ) = Ek Ldiff xk , xk+1 + Ek Ldiff Pθ (xk , xk+1 ), xk+2
5. Proxy Unrolled Training
Although gradient propagation is truncated beyond the n proxy steps, the use of a high-fidelity estimate provides a precise supervision signal, enabling the model to build resilience to its own prediction and correct potential artifact generation.
Training neural emulators via Teacher Forcing (TF) is notorious for leading to Simulation Exposure Bias (Schmidt, 2019; Brandstetter et al., 2022b; Chen et al., 2024): during inference, the model must condition on its own past predictions x̂k rather than the ground-truth used in training. Small errors in x̂k shift the input distribution for subsequent steps, causing errors to accumulate and trajectories to diverge.
6. Experimental Results 6.1. Experimental Setup Datasets. We evaluate our method on three distinct fluid dynamics datasets representing different physical regimes. Visualizations are provided in Figure 3
For fast-inference models, this limitation is typically mitigated via Unrolled Training (UT) (List et al., 2022), where the model is optimized over a horizon of U autoregressive steps. In the context of diffusion models, a naive unrolling takes the form: LUT (θ) = Et
U X u=1
Ldiff xk+u , M◦u xk θ
.
1D Kuramoto-Sivashinsky (KS) (Brandstetter et al., 2022a): Fourth-order nonlinear PDE, which governs flame front propagation and chaotic solidification dynamics. The scalar field u evolves according to ∂τ u + u∂x u + ∂x2 u + ν∂x4 u = 0. Numerical integration is performed on a periodic domain with 256 spatial points and a timestep ∆τ = 0.2. We use 512 training trajectories of length 140∆τ and 64 validation/testing trajectories of length 640∆τ . The models are trained with a step-size of 4∆τ .
(18)
However minimizing this loss is computationally prohibitive, as generating a single step x̂k+1 = Mθ (x̂k ) necessitates executing the full iterative sampling chain (e.g., T function evaluations). Therefore obtaining a simple pushforward estimate (Brandstetter et al., 2022b) is intractable, let alone backpropagating through the sampling chain.
2D Transonic Flow (Tra) (Kohl et al., 2023) : Simulated flow over a cylinder on a 128×64 grid. Time evolving fields are 2D velocity u, pressure p, and density ρ. The evaluations assess model performance over R = 60 timesteps.
5.1. Estimating the Model’s Output
2D Kolmogorov Flow (Kolmo) (Rozet and Louppe, 2023) : Incompressible fluid driven by sinusoidal forcing. The system obeys the Navier-Stokes equations subject to the incompressibility constraint ∇ · u = 0. 800 training trajectories of length 64, and 100 trajectories for validation/testing are simulated, with a spatial resolution of 64 × 64 and ∆τ = 0.2.
In this section, we assume that we have access to a diffusion model with low exposure-bias. We leverage this property for enabling computationally efficient unrolled training. In (n) particular we introduce a Proxy Estimate Pθ , which approximates the full model output using only the final n denoising steps :
Baselines. We train standard diffusion schedule baselines (Linear, Cosine, Sigmoid) using Teacher-Forcing loss. We further train deterministic U-Net baselines with TeacherForcing and Unrolled-Training with U = 2 and U = 8 (for each dataset we report the best performing U-Net training variant). Additionaly, we train PDE-Refiner (Lippe et al., 2023), a variant of diffusion that focuses on very low noiselevels.
(n) k Pθ (xk , xk+1 ) ≜ Denoisen...0 x̃k+1 . n , cond = x (19) √ √ Where x̃k+1 = ᾱn xk+1 + 1 − ᾱn ϵ. and Denoisen...0 n represents the sequence of n denoising steps from t = n down to t = 0. 6
Bias-Constrained Diffusion Schedules for PDE Emulations
Transonic (Pressure)
Kolmogorov (Vorticity) 2 1 0 1 2
0.8 0.6 0.4 0.2 0.0 0.2 0.4 0.6 0.8
Kuramoto-Sivashinsky (Velocity)
3 2 1 0 1 2
Time
Figure 3. Samples from benchmark datasets
Implementation Details. We use the same standard UNet architecture with attention mechanisms following (Kohl et al., 2023) for all models (all diffusion models and vanilla U-Net baselines). We employ continuous time-embeddings parameterized by the Signal-to-Noise Ratio (SNR) rather than discrete timesteps.
dominant factor driving performance gains. Regarding onestep MSE, our method consistently outperforms standard diffusion baselines by a significant margin. This improvement improves deccorelation time For Kolmogorov Flow and KS which are especially sensitive to initial error. We note that PDE-Refiner (Lippe et al., 2023) performs competitively on KS, achieving the best 1-step MSE and highcorrelation time among all methods; this is consistent with its design, which explicitly focuses capacity on low noise levels — a strategy that aligns with our Slow Error Decrease Principle. Our method achieves comparable performance on KS while also generalising to the other benchmarks where PDE-Refiner’s fixed low-noise focus is less effective. Visualisations of the obtained schedules. Figure 4 Left plot shows the reconstruction error landscape of the obtained schedule after running adaptive training over the 3 benchmark datasets. Inference-Input error nicely follows the Clean-input error, indicating that our proxy exposure-bias metrics yield consistent REB minimization. Furthermore, the obtained σ0 differ vastly for each dataset. Following observations made in (Lippe et al., 2023), a low optimal σ0 appears to be correlated with more high-frequency components in the energy spectrum in the fluid data. However, we argue that the Energy spectrum is not the only driver of σ0 . In particular, modifications in step-size naturally influence task difficulty, and could therefore lead to different errors and minimal sigma. Investigating the main drivers of noiselevel landscapes is of relevant importance, furthermore it remains to be investigated whether the instability thresholds γ(σ, τ ) are shared quantities across tasks.
Training strategy. The procedure contains two stages: (1) Pre-training: We first run the Adaptive Schedule Procedure to obtain a bias-constrained schedule for the given task. (2) Proxy Unrolled Fine-tuning: Since diffusion exposure bias is directly correlated with model error through the instability threshold, we only use the Unrolled proxy as a fine-tuning stage. This prevents introducing noisy gradients derived from inaccurate proxy estimates during the early phases of training, and saves computation time. We set the number of proxy steps to n = 1. Computational Cost & Hyperparameters. Each individual training run uses the same number of epochs E. For our adaptive schedule, the total number of epochs is therefore 2E as exploration and final training both contain a single training runs. All baseline methods are also trained to convergence over E epochs. For each dataset, E was picked as the largest budget which leads to improvement of the 1-step MSE. Full hyper-parameter configurations and architectural details are provided in Appendix G. Metrics. We report mean squared errors for the first prediction step and the short-term (10 steps) prediction. For KS and Kolmo, long-term accuracy is evaluated on the groundtruth correlation of the predicted signal and the vorticity ω = ∂x uy − ∂y ux , respectively, while for Tra we compute the MSE at the last time-step. To assess physical consistency in long-horizon rollouts for the 2D case, we report the Fréchet Spectral Distance (FSD) (Liu et al., 2025).
Impact of Proxy Unrolled Training. Consistent with the hypothesis that exposure bias degrades autoregressive rollouts, Proxy Unrolled Training significantly enhances both short-term precision and long-term stability. Notably, on Kolmogorov Flow, while the one-step error is similar as Teacher Forcing, our proxy method leads to an improved long-term correlation with the ground-truth trajectory. FSD improvements demonstrate the prevention of artifact formation and maintains bounded physical consistency for longrollouts, compared to baseline models. This validates that the proxy estimate is sufficient to sensitize the model to its own distribution shifts. We provide the temporal evolution
6.2. Results and Analysis Comparative results across the different benchmarks are displayed in Table 1. Impact of Adaptive Schedule. Our experiments show that the proposed schedule optimization framework is the 7
Bias-Constrained Diffusion Schedules for PDE Emulations 1-step MSE
10-steps MSE
High-Correlation Time (s) / Final MSE
Last-step FSD
U-Net UT, U = 2 Linear TF Sigmoid TF PDERefiner (Lippe et al., 2023)
9.53e-7 1.30e-6 1.25e-6 1.12e-6
2.79e-5 3.16e-4 9.81e7 8.16e-5
10.0 (6.5, 12.4) 9.7 (6.0, 12.3) 9.5 (4.6, 12.1) 10.2 (7.9, 12.2)
9.1e-1 2.1e5 2.4e6 1.05e1
Adaptive (Ours) Adaptive + Proxy UT, U = 1 (Ours)
8.12e-7 8.07e-7
2.18e-4 1.72e-5
10.7 (7.6, 12.6) 11.0 (9.0, 12.6)
3.7e5 3.8e-1
U-Net UT, U = 1 U-Net UT, U = 8 Linear TF Cosine TF
5.63e-5 4.15e-4 5.76e-5 4.90e-5
2.0e-5 2.54e-3 9.73e-4 8.79e-4
7.50e-1 1.08e-1 6.1e-1 1.22e-1
2.2e5 6.3e1 3.54e1 2.3e1
Adaptive (Ours) Adaptive + Proxy UT, U = 1 (Ours)
3.87e-5 4.05e-5
8.11e-4 6.66e-4
1.33e-1 1.25e-1
3.3e1 3.8e1
U-Net TF Cosine TF Sigmoid TF PDERefiner (Lippe et al., 2023)
1.60e-7 2.79e-7 1.55e-7 8.99e-8
1.50e-5 1.43e-5 1.57e-5 1.19e-5
68.4 (50.5, 80.0) 74.9 (48.8, 105.4) 78.4 (50.4, 113.4) 88.0 (52.5, 123.6)
– – –
Adaptive (Ours) Adaptive + Proxy UT, U = 1 (Ours)
9.55e-8 8.83e-8
9.74e-6 1.04e-5
85.8 (55.7, 117.4) 86.1 (55.0, 111.7)
– –
Dataset
Method
Kolmo
Tra
KS
Table 1. Comparison of Method Performance on Fluid Dynamics Datasets. We compare standard deterministic baselines and diffusion schedules against our Adaptive + Unrolling framework. For High-Correlation Time, the numbers in parenthesis are the high-correlation time of the worst and best 10 trajectories. The bolded numbers are the smallest value across methods in a (dataset, metric) pair. For FSD, we only consider improvements to be significant if they lead to an order of magnitude improvement. On each dataset, we report the best-performing models among baselines.
of the Fréchet Spectral Distance (FSD) in Appendix F.
10 2
KS (Clean) KS (Inference) Kolmo (Clean) Kolmo (Inference) Tra (Clean) Tra (Inference)
10 2
10 3
MSE
KS Kolmo Tra
100 10 1
Normalised E(k)
10 1
While our primary goal was to isolate the governing factors of sampling reconstruction error, one may be suspicious of the double training time required to idenfy a biasconstrained schedule. We claim, however, that our method can be considered a ”smart hyperparameter exploration”: given how large the space of schedules is, iterating over all possible schedules one by one would be much more computationally expensive. We showed that our method adapts well to different data statistics, with potential insights for larger-scale applications.
10 4 10 5
10 3 10 4 10 5 10 6 10 7
10 6
10 8
10 7 10 2
Noise Level
10 1
100
10 9 100
101
Wavenumber k
Figure 4. On the three benchmark datasets, the obtained adaptive schedules are, as shown in the left plot, exposure-bias free (inference error matches clean-input error). Interestingly, tasks which contain a higher high-frequency content tend to obtain a smaller minimal σ0 , matching the observations made in (Lippe et al., 2023).
A fundamental question raised by our findings is the balance between modeling uncertainty and maintaining reconstruction accuracy. Traditional diffusion schedules in computer vision are optimized for perceptual diversity, often at the expense of strict pixel-wise accuracy (Blau and Michaeli, 2018). Investigating the trade-off between those two objectives remains a compelling direction for future research.
7. Discussion and Future Work In this work, we have established an empirically grounded link between noise schedules, reconstruction error, and diffusion exposure bias for PDE diffusion models. We demonstrated that our adaptive schedule can yield near-optimal reconstruction, showing superior performance compared to both deterministic and probabilistic baselines. Furthermore, we proposed to leverage exposure-bias reduction and obtain a proxy unrolled training estimate, leading to considerable improvements in mitigating artifact formation. Beyond autoregressive simulations, our schedule holds promise for other reconstruction-bound scenarii, such as diverse inverse problems, data-assimilation or super-resolution tasks.
Finally, whether the reconstruction error landscape is influenced by the choice of denoising prediction parameterization (e.g., x0 -prediction (Hoogeboom et al., 2023) and v-prediction (Salimans and Ho, 2022)) is an important question. For instance, x0 -prediction typically yields smaller errors in early sampling steps, which may fundamentally alter the dynamics. Another potential direction would be to evaluate if the type of exposure bias we have shed light on, which is, different from the bias identified in (Ning et al., 2023), rather bound to allocation of model capacity, is also present in unconditional diffusion models. 8
Bias-Constrained Diffusion Schedules for PDE Emulations
8. Impact Statement
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.
This paper displays methodologies for enhancing probabilistic fluid models. While this could be of societal benefit for forecasting tasks such as weather and climate, we acknowledge that high-fidelity fluid simulations are general-purpose tools that can also be applied to military contexts, such as aerodynamics for defense technologies.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020.
References
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in neural information processing systems, 35:8633–8646, 2022.
Marc Amorós-Trepat, Luis Medrano-Navarro, Qiang Liu, Luca Guastoni, and Nils Thuerey. Guiding diffusion models to reconstruct flow fields from sparse data. Physics of Fluids, 38(1), 2026.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pages 13213–13232. PMLR, 2023.
Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 843–852, 2023.
Xun Huang, Zhengqi Li, Guande He, Mingyuan Zhou, and Eli Shechtman. Self forcing: Bridging the train-test gap in autoregressive video diffusion. arXiv preprint arXiv:2506.08009, 2025. Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565–26577, 2022.
Jan-Hendrik Bastek, WaiChing Sun, and Dennis M Kochmann. Physics-informed diffusion models. arXiv preprint arXiv:2403.14404, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021.
Yochai Blau and Tomer Michaeli. The perception-distortion tradeoff. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6228–6237, 2018.
Georg Kohl, Li-Wei Chen, and Nils Thuerey. Benchmarking autoregressive conditional diffusion models for turbulent flow simulation. arXiv preprint arXiv:2309.01745, 2023.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning, pages 2241–2256. PMLR, 2022a.
Mingxiao Li, Tingyu Qu, Ruicong Yao, Wei Sun, and MarieFrancine Moens. Alleviating exposure bias in diffusion models through sampling with shifted time steps. arXiv preprint arXiv:2305.15583, 2023.
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint arXiv:2202.03376, 2022b.
Marten Lienen, David Lüdke, Jan Hansen-Palmus, and Stephan Günnemann. From zero to turbulence: Generative modeling for 3d flow simulation, 2024. URL https://arxiv. org/abs/2306.01776.
Boyuan Chen, Diego Martı́ Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. Advances in Neural Information Processing Systems, 37:24081–24125, 2024.
Phillip Lippe, Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Processing Systems, 36:67398– 67433, 2023.
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11472–11481, 2022.
Björn List, Li-Wei Chen, and Nils Thuerey. Learned turbulence modelling with differentiable fluid solvers: physicsbased loss functions and optimisation horizons. Journal of Fluid Mechanics, 949:A25, 2022.
Giannis Daras, Yuval Dagan, Alex Dimakis, and Constantinos Daskalakis. Consistent diffusion models: Mitigating sampling drift by learning to be consistent. Advances in Neural Information Processing Systems, 36:42038– 42063, 2023.
Kunhao Liu, Wenbo Hu, Jiale Xu, Ying Shan, and Shijian Lu. Rolling forcing: Autoregressive long video diffusion in real time. arXiv preprint arXiv:2509.25161, 2025. 9
Bias-Constrained Diffusion Schedules for PDE Emulations
Qiang Liu and Nils Thuerey. Uncertainty-aware surrogate models for airfoil flow simulations with denoising diffusion probabilistic models. AIAA Journal, 62(8):2912– 2933, 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256–2265. pmlr, 2015.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162–8171. PMLR, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023. Jason Stock, Troy Arcomano, and Rao Kotamarthi. Swift: An autoregressive consistency model for efficient weather forecasting. arXiv preprint arXiv:2509.25631, 2025.
Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Onal Ertugrul. Elucidating the exposure bias in diffusion models. arXiv preprint arXiv:2308.15321, 2023. Ilan Price, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, et al. Gencast: Diffusion-based ensemble forecasting for medium-range weather. arXiv preprint arXiv:2312.15796, 2023. François Rozet and Gilles Louppe. Score-based data assimilation. Advances in Neural Information Processing Systems, 36:40521–40541, 2023. Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. Dyffusion: A dynamics-informed diffusion model for spatiotemporal forecasting. Advances in neural information processing systems, 36:45259–45287, 2023. Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022. Florian Schmidt. Generalization in generation: A closer look at exposure bias. arXiv preprint arXiv:1910.00292, 2019. Youssef Shehata, Benjamin Holzschuh, and Nils Thuerey. Improved sampling of diffusion models in fluid dynamics with tweedie’s formula. In The Thirteenth International Conference on Learning Representations, 2025. Lifeng Shen, Weiyu Chen, and James Kwok. Multiresolution diffusion models for time series forecasting. In The Twelfth International Conference on Learning Representations, 2024. Dule Shu, Zijie Li, and Amir Barati Farimani. A physicsinformed diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics, 478: 111972, 2023. Aliaksandra Shysheya, Cristiana Diaconu, Federico Bergamin, Paris Perdikaris, José Miguel HernándezLobato, Richard Turner, and Emile Mathieu. On conditional diffusion models for pde simulations. Advances in Neural Information Processing Systems, 37:23246– 23300, 2024. 10
Bias-Constrained Diffusion Schedules for PDE Emulations
A. Preliminaries A.1. Conditional Diffusion Models Denoising Diffusion Probabilistic Models (DDPM) are generative models that learn to approximate a data distribution p(y) by reversing a gradual noising process. We consider the conditional setting where the generation of a target y ∈ Rd is conditioned on an input x. Given a fixed variance schedule {βt ∈ (0, 1)}Tt=1 , the forward process q(ỹ1:T |y0 ) transforms a clean sample y0 ∼ p(y|x) into Gaussian noise ỹT ∼ N (0, I) via a Markov chain: p q(ỹt |ỹt−1 ) = N (ỹt ; 1 − βt ỹt−1 , βt I). (20) A notable property of this process is that any intermediate noisy state ỹt can be sampled directly from y0 : √ √ ỹt = ᾱt y0 + 1 − ᾱt ϵ, ϵ ∼ N (0, I), where αt ≜ 1 − βt and ᾱt ≜
Qt
s=1 αs . We characterize the noise intensity at step t by the standard deviation σt ≜
(21) √
1 − ᾱt .
The generative process pθ (y0:T |x) learns to reverse this corruption. A neural network yest is trained to predict the clean signal y0 (or equivalently the noise ϵ) given a noisy latent ỹt , the conditioning x, and the time step t. The model is optimized by minimizing the re-weighted squared error: Ldiff (θ) = Ex,y0 ,t,ϵ ∥yest (ỹt , x, σt ) − y0 ∥2 . (22) A.2. Autoregressive Neural Emulators In the context of fluid dynamics, we aim to construct a surrogate model Mθ that emulates the temporal evolution of a physical system. Given an initial state x0 , the objective is to generate a trajectory {x̂1 , . . . , x̂K } that approximates the ground-truth sequence {x1 , . . . , xK }. This is typically formulated as an autoregressive problem: for k ∈ {1, . . . , K},
x̂k = Mθ (x̂k−1 )
(23)
with x̂0 = x0 . Standard training relies on Teacher Forcing (TF), which minimizes the one-step prediction error conditioned on ground-truth states: LTF (θ) = Ek ∥Mθ (xk ) − xk+1 ∥2 . (24) However, TF is susceptible to exposure bias: during inference, the model accumulates small errors, leading to a distributional shift in the input x̂k that diverges from the training distribution. To mitigate this, Unrolled Training (UT) optimizes the model over a horizon of M autoregressive steps: LUT (θ) = Ek
M X m=1
∥x̂k+m − xk+m ∥2 ,
where x̂k+m is recursively generated from the model’s own prior prediction x̂k+m−1 .
11
(25)
Bias-Constrained Diffusion Schedules for PDE Emulations
B. Proof of Proposition 3.3 Proposition B.1 (Re-noising Attenuation, restated). When estimated errors are nearly aligned, the following recursive bound holds: REB(t) ≲ B (2S) (t) + λt REB(t + 1), √ (ỹt+1 )−y∥ where λt := ∥Jt ∥ · ᾱt · ∥y∥yest . The REB is thus driven by the local two-step bias, and attenuated (resp. amplified) est (ỹt )−y∥ across steps when λt < 1 (resp. λt > 1). (2S)
Full bound. Let ρ1 = cos θ1 where θ1 is the angle between yest (ŷt θ2 is the angle between yest (ŷt ) − y and yest (ỹt ) − y. Then: REB(t) ≤ 1 +
q
B (2S) (t) − ρ
1
2
+ (1 − ρ21 )
) − y and yest (ỹt ) − y, and let ρ2 = cos θ2 where
+ λt ·
q
REB(t + 1) − ρ2
2
(26)
+ (1 − ρ22 ).
When ρ1 , ρ2 ≈ 1, the square-root terms simplify to B (2S) (t) and REB(t + 1) respectively, recovering the statement above. Sampling definitions. Both inference and ground-truth paths follow one-step backward diffusion: (2S)
ŷt−1 = ŷt−1 =
√
√
ᾱt−1 yest (ỹt ) + σt−1 ϵ,
(27)
ᾱt−1 yest (ŷt ) + σt−1 ϵ,
(28) (29)
ŷT ∼ N (0, I), √ ỹt = ᾱt y + σt ϵ,
(30) (2S)
where ϵ ∼ N (0, I) is the same noise draw shared between equations (27) and (28). Similarly, ŷt and ŷt noise at their respective levels.
use the same
The two-steps and reconstruction bias are given as: (2S)
B (2S) (t) ≜
yest (ŷt ) − y , yest (ỹt ) − y
(31)
REB(t) ≜
yest (ŷt ) − y . yest (ỹt ) − y
(32)
Alternative biases definitions: We furthemore define an alternative version of the biases taking into account the distance of estimated values with each other, rather than only their distance to y. (2S)
Balt (t) ≜
yest (ŷt ) − yest (ỹt ) , yest (ỹt ) − y
(33)
REBalt (t) ≜
yest (ŷt ) − yest (ỹt ) . yest (ỹt ) − y
(34)
(2S)
Proof. We have that:
REBalt (t) =
yest (ŷt ) − yest (ỹt ) ∥yest (ỹt ) − y∥
(2S)
≤
yest (ŷt ) − yest (ỹt ) ∥yest (ỹt ) − y∥ 12
(2S)
+
yest (ŷt ) − yest (ŷt ∥yest (ỹt ) − y∥
)
(35)
Bias-Constrained Diffusion Schedules for PDE Emulations (2S)
Both ŷt and ŷt
are obtained by a backward diffusion step with the same noise draw, so: (2S)
=
√
(2S)
)
≲ ∥Jt ∥ ·
ŷt − ŷt
(36)
ᾱt yest (ŷt+1 ) − yest (ỹt+1 ) .
A first-order Jacobian bound then gives: yest (ŷt ) − yest (ŷt
√
ᾱt · yest (ŷt+1 ) − yest (ỹt+1 ) ,
(37)
where Jt := ∇yest |ŷ(2S) . t
By definition, ∥yest (ŷt+1 ) − yest (ỹt+1 )∥ = REBalt (t + 1) · ∥yest (ỹt+1 ) − y∥. Therefore (35) gives : (2S)
REBalt (t) ≲ Balt (σt , σt+1 ) + λt · REBalt (t + 1).
(38)
with λt := ∥Jt ∥ ·
√
yest (ỹt+1 ) − y yest (ỹt ) − y
ᾱt ·
(39)
(2S)
Relationship between the two definitions. Let a = yest (ŷt ) − y and b = yest (ỹt ) − y, so that the standard definition (2S) gives B (2S) = ∥a∥/∥b∥ and the alternative gives Balt = ∥a − b∥/∥b∥. Letting θ denote the angle between a and b, the law of cosines yields: (2S) 2
Balt
=
B (2S)
2
+ 1 − 2 B (2S) cos θ =
B (2S) − cos θ
2
+ (1 − cos θ2 ),
(40)
Combining equation (40) and the triangle inequality: B
(2S)
(t) − 1 ≤
(2S) Balt (t) =
q
B (2S) (t) − ρ1
2
+ (1 − ρ21 ),
(41)
2
(42)
The same inequalities applie for the REB: REB(t) − 1 ≤ REBalt (t) ≤
q
REB(t) − ρ2
+ (1 − ρ22 ),
where ρ2 = cos θ2 and θ2 is the angle between yest (ŷt ) − y and yest (ỹt ) − y.
Therefore we can express the standard REB in terms of the standard 2-steps bias and the next-step REB: (43)
REB(t) ≤ 1 + REBalt (t) ≤
(2S) 1 + Balt (t)
≤ 1 +
q
+ λt · REBalt (t + 1) q 2 2 B (2S) (t) − ρ1 + (1 − ρ21 ) + λt · REB(t + 1) − ρ2 + (1 − ρ22 )
13
(44) (45)
Bias-Constrained Diffusion Schedules for PDE Emulations
B.1. Visualization Input Type
Reconstruction Error
4 × 10−5 3 × 10−5
Clean input Inference input Own-Pred input Two-Steps input
2 × 10−5
10−5
6 × 10−6
10−1
Noise Level σ
Figure 5. Contributions of errors to the REB. Metrics are reported for the final diffusion steps of a linear schedule model. Two-steps error nearly corresponds to inference-input error even though it only contains errors propagated from current and previous step. Own-prediction error is the major driver of the two-steps bias.
14
Bias-Constrained Diffusion Schedules for PDE Emulations
C. Proof of Proposition 3.4 Proposition C.1 (Stability Threshold, restated). Assuming (A1) Wiener denoiser structure and (A2) Spectral bias of the (own) neural denoiser (detailed below), Bθ (t) is an increasing function of Eθclean (t). In particular, at each noise level σt and for any threshold τ ≥ 1, there exists a critical clean-input error γ(σt , τ ), increasing in σt and decreasing in τ , such that if (own) Eθclean (σt ) ≤ γ(σt , τ ) then Bθ (t) ≤ τ . The proposition follows from the more precise result below, which makes assumptions (A1) and (A2) explicit. Proposition C.2 (Bias reduction through error-spectral shift). Let θA and θB be two denoisers and fix a noise level σ. Let Jθ denote the Jacobian of yest with respect to its noisy input, and let Eθ (k) denote the per-wavenumber residual error variance Eθ (k) = E[|ε̂k |2 /N ], where ε̂k is the Fourier coefficient of the clean-input residual yest (ỹσ , x, σ) − y. 1. (Own-Prediction Bias √ under Wiener denoiser structure.) Assuming that Jθ is diagonal in the Fourier basis with eigenvalue λk = ( ᾱ/σ 2 ) Eθ (k) at each wavenumber k, then P ᾱ 2 P E (k)3 Eθ (k)2 2ᾱ (own) Pk θ Bθ (σ) ≈ 1 + 2 Pk + . (46) 2 E (k) σ σ θ k k Eθ (k) 2. (Spectral bias implies bias reduction.) Suppose both models satisfy the Wiener structure of (1), that EB (k) is decreasing in k (larger residual errors at low wavenumbers), and that 0 ≤ EA (k) ≤ EB (k) for all k with the ratio rk := EA (k)/EB (k) non-decreasing in k (the relative error reduction is larger at low wavenumbers). Then (own) (own) R(θA ) ≤ R(θB ), and hence BθA (σ) ≤ BθB (σ). Proof. Part (1). Let ε = yest (ỹσ , x,√σ) − y denote the clean-input residual. In the own-prediction√setting (??), the model receives the re-noised input ŷσ = √ᾱ yest (ỹσ , x, σ) + σ z instead of the ground-truth ỹσ = ᾱ y + σ z. The input perturbation is thus δ = ŷσ − ỹσ = ᾱ ε. By first-order Taylor Expansion around ỹσ , we can write: yest (ŷσ , x, σ) − y ≈ ε + Jθ δ = ε +
√
(47)
ᾱ Jθ ε.
Taking the squared norm and dividing by ∥ε∥2 , we obtain: √ √ ∥ε + ᾱ Jθ ε∥2 2 ᾱ ⊤ ᾱ (own) Bθ (σ) = = 1 + ε Jθ ε + ∥Jθ ε∥2 2 2 ∥ε∥ ∥ε∥ ∥ε∥2
(48)
√ Under the Wiener structure, Jθ is diagonal in the Fourier basis with eigenvalue λk = ( ᾱ/σ 2 ) Eθ (k) at wavenumber k. By Parseval’s theorem: 1 X 1 X 2 1 X ε⊤ Jθ ε = λk |ε̂k |2 , ∥Jθ ε∥2 = λk |ε̂k |2 , ∥ε∥2 = |ε̂k |2 . (49) N N N k
k
Since E[|ε̂k |2 /N ] = Eθ (k), we replace |ε̂k |2 by its expectation N Eθ (k) to obtain: √ P √ P 2 ᾱ k λk Eθ (k) E (k)2 2 ᾱ ε⊤ Jθ ε 2ᾱ P Pk θ = , ≈ 2 2 ∥ε∥ σ k Eθ (k) k Eθ (k) P P 3 ᾱ k λ2k Eθ (k) ᾱ 2 ᾱ ∥Jθ ε∥2 k Eθ (k) P P ≈ = , ∥ε∥2 σ2 k Eθ (k) k Eθ (k) √ where we substituted λk = ( ᾱ/σ 2 ) Eθ (k). Combining both terms with (48) yields (46). (own)
Part (2). By (46), Bθ
is determined (up to monotone transformations) by the ratio P Eθ (k)2 R(θ) := Pk , k Eθ (k) 15
k
(50) (51)
Bias-Constrained Diffusion Schedules for PDE Emulations
P P which is the Eθ (k)-weighted mean of Eθ (k). (The cubic term k Eθ (k)3 / k Eθ (k) behaves analogously; we focus on R for clarity.) It therefore suffices to show that R(θA ) ≤ R(θB ). Write EA (k) = rk EB (k) with rk ∈ [0, 1] non-decreasing. Since R(θA ) = P 2 r EB (k)2 R(θA ) = Pk k rk EB (k) Pk rk EB (k)2 ≤ Pk rk EB (k) Pk P EB (k)2 k rk EB (k) ≤ P · Pk EB (k) k EB (k) Pk 2 EB (k) ≤ · Pk k EB (k)
P
2 2 k rk EB (k)
k rk EB (k), we have :
P
(Jensen’s inequality) (Chebyshev covariance inequality) P k rk EB (k) P since ≤1 k EB (k)
= R(θB ) where Jensen’s Inequality states that rk2 ≤ rk for every k, given 0 ≤ rk ≤ 1 ; and Chebyshev covariance inequality states that for two sequences f (k) and g(k) that are oppositely monotone (one non-decreasing, the other non-increasing) and non-negative weights wk ≥ 0: ! ! ! ! X X X X wk wk f (k) g(k) ≤ wk f (k) wk g(k) , k
k
k
k
P or equivalently Ew [f (k) g(k)] ≤ Ew [f (k)] Ew [g(k)]. Given the probability weights wk = EB (k)/ j EB (j), since rk is non-decreasing in k and EB (k) is non-increasing in k by assumption, applying this with f (k) = rk and g(k) = EB (k) gives Covw (r, EB ) ≤ 0, i.e., Ew [rk EB (k)] ≤ Ew [rk ] Ew [EB (k)]. (own)
Finally, by (46), Bθ
(own)
is increasing in R, so R(θA ) ≤ R(θB ) implies BθA
(own)
(σ) ≤ BθB
(σ).
Remark. The decreasing-Eθ (k) condition in (2) is satisfied whenever P (k) is decreasing (e.g. turbulence spectra P (k) ∼ k −β ), since the MMSE residual of the Wiener denoiser E(k) = σ 2 P (k)/(ᾱP (k) + σ 2 ) is an increasing function of P (k). The spectral-bias condition (the tendency of gradient-trained networks to learn low-frequency components first) is necessary: without it, a uniform rescaling EA (k) = c EB (k) leaves R unchanged, and hence B (own) unchanged too.
16
Bias-Constrained Diffusion Schedules for PDE Emulations
D. Proof of Proposition 3.5 Proposition D.1 (Slow Error Decrease Principle, restated). Under the stability condition established in Proposition 3.4 and assumptions (A1)–(A2), the schedule minimizing (15) contains no unnecessary noise levels: every intermediate σk is necessary in the sense that removing it — i.e., going directly from σk−1 to σk+1 — would violate the bias constraint, i.e. B (2S) (k − 1) > τ . In particular, the greedy construction that always takes the largest feasible jump is optimal. The proof relies on the instability threshold (Proposition 3.4) and the following assumptions: (A1) Finite capacity trade-off: Given a model trained on schedule S, adding an extra noise level σ ′ to S and retraining to convergence necessarily increases Eθclean (σs ) for some σs ∈ S. (A2) Monotone error decay: Eθclean (σ) is strictly decreasing in σ: lower noise levels yield strictly lower clean-input error. Proof. The final inference error decomposes as Eθinf (0) = REB(0) · Eθclean (0). Minimizing this requires both REB(0) ≈ 1 and Eθclean (0) as small as possible.
Suppose the optimal schedule S ∗ contains an intermediate noise level σk (0 < k < T ) that is unnecessary, i.e., B (2S) (σk−1 , σk+1 ) ≤ τ so that σk can be skipped without violating stability. Removing σk from S ∗ yields a strictly shorter schedule. By (A1), retraining on this shorter schedule frees capacity, which the model can reallocate to σ0 , strictly decreasing Eθclean (σ0 ) — possible by (A2), since σ0 is the lowest noise level and thus has the most room for improvement relative to other levels. This strictly reduces the final inference error, contradicting the optimality of S ∗ .
Therefore every intermediate step in the optimal schedule is necessary. Note that the bias constraint B (2S) (σt , σt+1 ) ≤ τ need not be tight at each step — it suffices that no larger jump is feasible. Furthermore, satisfying the bias constraint becomes increasingly easy at higher noise levels, since the instability threshold γ(σ, τ ) is increasing in σ: the model can afford larger clean-input errors at high noise, allowing larger jumps. This justifies the greedy construction, which always takes the largest feasible jump and is therefore optimal.
17
Bias-Constrained Diffusion Schedules for PDE Emulations
Noise Level sqrt(1- ) Distribution
E. Accuracy of the Proxy Estimator
Linear Schedule
60 50 40 30 20 10 0
LogMinus5Uniform Schedule
15
15
10
10
5
5
0
0
Training input 10 1 Inference input
Training input 10 1 Inference input
10 3
10 3
10 3
10 4
10 4
10 4
10 5
10 5
10 5
10 6
10 6
10 6
10 1 MSE w/ ground-truth
PSD Schedule
10 2
10 2
10 1
10 2
Training input Inference input
10 2
10 1 100 10 2 100 10 2 Noise Level sqrt(1- ), High = Only Noise, Low = Only Image
Training Final Error: 2.81e-06 Inference Final Error: 3.43e-06
10 1
Training Final Error: 1.26e-06 Inference Final Error: 1.34e-06
Linear Schedule
100
Training Final Error: 6.54e-07 Inference Final Error: 2.01e-06
PSD Schedule
LogMinus5Uniform Schedule
1.000 0.995
Cosine similarity
0.990 0.985 0.980 0.975 0.970
GT vs. Diffusion sample D Estimate Pn vs Diffusion sample D D sample 0 vs. D sample 1 (upper-bound)
0.965 0.960
1.0
1.5
2.0
2.5
3.0
Number of proxy steps n
3.5
4.0
1.0
1.5
2.0
2.5
3.0
Number of proxy steps n
3.5
4.0
1.0
1.5
Similarity of gradients obtained from next autoregressive step loss, using different inputs
2.0
2.5
3.0
Number of proxy steps n
3.5
4.0
Figure 6. Similarity of gradients obtained from the second step loss LP-2UT , using as input either the ground-truth data (= teacher forcing), (n) the true previous step diffusion sample (= full unrolled training) or the proxy estimate Pθ , as a function of the number of steps n. PSD and LogMinus5Uniform Schedules are arbitrary schedules that we defined. The alignment of the gradients well reflect the different exposure-biases. In particular, with PSD Schedule, which doesn’t suffer from exposure-bias, the proxy estimate is already fully accurate in a single step (n = 1). For LogMinus5Uniform Schedule, the estimate still isn’t accurate after 4 steps.
18
Bias-Constrained Diffusion Schedules for PDE Emulations
F. FSD Evolution
Figure 7. Kolmogorov Flow
Figure 8. Transonic Flow
Our proxy unrolled training is particularly useful when the model suffers from artifacts, as is the case on the Kolmogorov Flow. The improvement is however marginal when the distribution shift is steady and without jumps, as in Transonic Flow.
19
Bias-Constrained Diffusion Schedules for PDE Emulations
G. Implementation details G.1. Algorithm Pseudocode Algorithm 1 Phase 1 — Exploration 1: Require: bias tolerance τ , exploration grid σ exp = {σ (1) , . . . , σ (N ) } (log-uniform), shared weights w = 1 2: Initialize: active schedule σ active ← σ exp , solved set σ solved ← ∅ 3: while σ active ̸= ∅ do 4: Train θ on σ active with uniform weights until convergence 5: for each σ (i) ∈ σ active do (own) (i) 6: if Bθ (σ ) ≤ τ then 7: Save checkpoint θ∗ (σ (i) ) ← θ 8: σ active ← σ active \ {σ (i) } 9: σ solved ← σ solved ∪ {σ (i) } 10: end if 11: end for 12: if no new level solved this pass then 13: break 14: end if 15: end while 16: return {θ ∗ (σ)}σ∈σsolved
Algorithm 2 Phase 2 — Greedy Schedule Construction 1: Require: checkpoints {θ ∗ (σ)}σ∈σsolved , bias tolerance τ 2: σ0 ← min σ solved , t ← 0, S ← [σ0 ] 3: while σt < σT do n 4:
σt+1 ← max σ ′ ∈ σ solved : Bθ∗ (σ′ ), θ∗ (σt ) (σt , σ ′ ) ≤ τ (2S)
o
5: S ← S ∪ [σt+1 ], t ← t + 1 6: end while 7: Fine-tune a single shared model θ on S, warm-started from θ ∗ (σ0 ) 8: return S, θ
G.2. Training Hyperparameters We set τ = 1.05 across all tasks as it constrains the two-steps bias while allowing for flexibility. A stricter value could potentially reduce the reconstruction error, but at the cost of making the optimization process more complex. If not mentioned otherwise, each baseline diffusion model uses T = 20 steps. A training run contains E = 1000 epochs on Kolmogorov Flow and KS, E = 2000 epochs on Transonic Flow. On the other hand, unrolled fine-tuning is done for 200 epochs on Kolmogorov Flow and KS, and 400 epochs for Transonic.
H. Dataset and Data Generation • Kolmogorov Flow : We use the generation pipeline from (Rozet and Louppe, 2023), to generate 800 trajectories for training, 100 for testing, 100 for validation. Generation scripts can be obtained in SDA. • Transonic Flow : We obtain the dataset from the ACDM benchmark (Kohl et al., 2023) (repository ACDM). The testing experiments are ran over the Extrapolate test case. • Kuramoto-Sivashinsky : We generate the data using the pipeline provide in (Brandstetter et al., 2022a) (repository LPSDA). We increase the number of validation and testing trajectories to 512, and fix the timestep ∆t to 0.2, as well as the grid spacing ∆x, following (Shysheya et al., 2024). 20