cs.LGApr 9, 2026

Bias-Constrained Diffusion Schedules for PDE Emulations: Reconstruction Error Minimization and Efficient Unrolled Training

Authors: Constantin Le Cleï, Nils Thuerey, Xiaoxiang Zhu

Organizations: Technical University of Munich, Munich, Germany.

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Conditional Diffusion Models are powerful surrogates for emulating complex spatiotemporal dynamics, yet they often fail to match the accuracy of deterministic neural emulators for high-precision 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.

Figures & tables

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\tau\pm 0.05 box, where τ∈[1.01,..,1.15]\tau\in[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,yB(own)≤τ\mathbb{E}_{\mathbf{x},\mathbf{y}}\mathcal{B}^{(own)}\leq\tau at convergence, indicating that there is a – data and model-dependent – minimal noise-level where the model can be bias-free.
Figure 2: Effect of schedule on final reconstruction error : Clean and Inference Input Errors during training. The diffusion sampling direction goes from right (high noise-levels) to left (low noise-levels). Schedules with more weight on low noise-levels obtain a better reconstruction performance early in the denoising process. However they are then unable to reduce the error further compared to schedules that put more focus on low noise-levels. At the end of sampling, schedules with higher clean-input reconstruction error start suffering from REB.
Figure 3: Samples from benchmark datasets
DatasetMethod1-step MSE10-steps MSEHigh-Correlation Time (s) / Final MSELast-step FSD
KolmoU-Net UT, U=2U=29.53e-72.79e-510.0 (6.5, 12.4)9.1e-1
 Linear TF1.30e-63.16e-49.7 (6.0, 12.3)2.1e5
 Sigmoid TF1.25e-69.81e79.5 (4.6, 12.1)2.4e6
 PDERefiner [ Lippe et al., 2023 ]1.12e-68.16e-510.2 (7.9, 12.2)1.05e1
 Adaptive (Ours)8.12e-72.18e-410.7 (7.6, 12.6)3.7e5
 Adaptive + Proxy UT, U=1U=1 (Ours)8.07e-71.72e-511.0 (9.0, 12.6)3.8e-1
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.
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\sigma_{0} , matching the observations made in [ Lippe et al., 2023 ] .
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

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.
Figure 6: Similarity of gradients obtained from the second step loss LP-2UT\mathcal{L}_{\text{P-2UT}} , using as input either the ground-truth data (= teacher forcing), the true previous step diffusion sample (= full unrolled training) or the proxy estimate Pθ(n)\mathcal{P}^{(n)}_{\theta} , as a function of the number of steps nn . 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=1n=1 ). For LogMinus5Uniform Schedule, the estimate still isn’t accurate after 4 steps.
Figure 7: Kolmogorov Flow
Figure 8: Transonic Flow
More in this topic

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Explore similar work

CardsList

  1. Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

    May 21, 2026Seo Taek Kong, Weina Wang, R. SrikantDiffusion SamplingDiffusion-Based Image Generation

    ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign

  2. ReAl: Refine neural operators via aligning dynamical flow

    Oct 14, 2026Jing Xu, Zenglin Xu, Li Yi +2PDE Surrogate ModelingPDE Operator Learning

    Shanghai Academy of AI for Science · Fudan University

  3. Kastor: An efficient fine-tuning strategy for generative emulation of PDE simulations

    Aug 6, 2026Guillaume Couairon, Alexis Jacq, Yu-Han Wu +4Surrogate ModelingFine-Tuning

    Google DeepMind