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

TMLR Paper6097 Authors

Abstract:

We propose DiffusionRollout, a novel selective rollout planning strategy for autoregres-
sive 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 effective-
ness 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.

Submission Length: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Geoff_Pleiss1

Submission Number: 6097