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

TL;DR

DiffusionRollout introduces an uncertainty-aware planning method for long-horizon PDE predictions, adaptively selecting step sizes to reduce error accumulation and improve prediction accuracy in physical system modeling.

Contribution

It presents a novel adaptive rollout strategy leveraging uncertainty quantification to enhance long-term PDE predictions, addressing error propagation issues.

Findings

Lower prediction errors on long trajectories

Longer predicted trajectories with high accuracy

Uncertainty measures correlate strongly with prediction errors

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.