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

2026-02-14

AI Summary

This work addresses the degradation in reliability of autoregressive diffusion models for long-term partial differential equation (PDE) forecasting, which stems from error accumulation over time. To mitigate this issue, the authors propose an adaptive step-size selection mechanism grounded in predictive uncertainty. Specifically, they introduce the use of the standard deviation across multiple diffusion model samples as a proxy for confidence, enabling dynamic adjustment of the rollout step size to reduce reliance on inaccurate historical states. By integrating Monte Carlo–based uncertainty estimation with an adaptive planning strategy, the method significantly enhances both stability and accuracy in long-horizon PDE solutions. Experiments demonstrate that the approach effectively reduces prediction errors across multiple long-trajectory PDE benchmarks, yielding extended solution trajectories that closely align with ground-truth dynamics.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• long-horizon prediction
• error accumulation
• partial differential equations
• autoregressive diffusion models
• predictive uncertainty

Innovation

• DiffusionRollout
• uncertainty-aware planning
• autoregressive diffusion models
• PDE solving
• error accumulation mitigation