DiffusionRollout: Uncertainty-Aware Rollout Planning in Long-Horizon PDE Solving

TL;DR

DiffusionRollout tackles the problem of error accumulation in long-horizon PDE predictions by leveraging the probabilistic nature of conditional diffusion models to quantify predictive uncertainty. It introduces a training-free adaptive rollout that uses sample-based uncertainty to selectively advance timesteps, balancing network approximation error and condition-induced error without retraining. Across Gray-Scott, turbulent flow, Cahn-Hilliard, and anisotropic diffusion benchmarks, it outperforms state-of-the-art baselines and maintains high cross-time correlations, demonstrating both improved accuracy and extended reliable horizons. The approach offers practical benefits for physics-informed simulations by providing reliable long-term forecasts with minimal additional computation beyond standard diffusion-based solvers.

Abstract

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

Figures (7)

• Figure 1: Adaptive rollout step size $s$ based on predictive uncertainty. The figure illustrates how predictive uncertainty $\hat{\boldsymbol{\varepsilon}}_t$ governs the rollout step size. On the left, low sample variance leads all future states to satisfy $\hat{\boldsymbol{\varepsilon}}_t<\tau$, allowing a full-step rollout ($s=F$). On the right, high variance (i.e., uncertainty) causes only the first state to pass the threshold, leading to $s=1$. Based on this, our diffusion-based PDE solver adaptively adjusts the rollout step size during autoregressive prediction.
• Figure 2: Analysis of predictive uncertainty $\hat{\boldsymbol{\varepsilon}}_t$ as a surrogate for prediction error. The proposed quantity exhibits a strong correlation with the true prediction error across multiple PDE benchmarks, including (from left to right) Gray-Scott Ohana:2024TheWell, Turbulent Flow Ohana:2024TheWell, Cahn-Hilliard Soares:2023Exponential, and Anisotropic Diffusion Koehler:2024APEBench. Best viewed when zoomed in.
• Figure 3: Qualitative comparison. Per-pixel error maps (columns 3–8) are visualized for each model’s prediction of the target (column 2), given the initial context (column 1). Darker regions in the error maps indicate lower prediction error.
• Figure 4: Qualitative comparison from the ablation study. Per-pixel error maps (columns 3-5 and 8-10) are shown for each variant of our method, visualizing the prediction error (columns 2 and 7), given the initial context (columns 1 and 6). Darker regions in the error maps indicate lower prediction error.
• Figure 5: Relative $\mathcal{L}_2$ errors of predicted solutions versus the average step sizes $s_{\text{avg}}$ under different uncertainty thresholds $\tau$ (shown as overlaid text). For comparison, the relative $\mathcal{L}_2$ errors of PDE-Refiner with constant step sizes are also visualized. As shown, selecting an appropriate threshold is crucial for balancing prediction accuracy and inference time, as reflected by both the error and the average step size.