PDE-regularized Dynamics-informed Diffusion with Uncertainty-aware Filtering for Long-Horizon Dynamics

Abstract: Long-horizon spatiotemporal prediction remains a challenging problem due to cumulative errors, noise amplification, and the lack of physical consistency in existing models. While diffusion models provide a probabilistic framework for modeling uncertainty, conventional approaches often rely on mean squared error objectives and fail to capture the underlying dynamics governed by physical laws. In this work, we propose PDYffusion, a dynamics-informed diffusion framework that integrates PDE-based regularization and uncertainty-aware forecasting for stable long-term prediction. The proposed method consists of two key components: a PDE-regularized interpolator and a UKF-based forecaster. The interpolator incorporates a differential operator to enforce physically consistent intermediate states, while the forecaster leverages the Unscented Kalman Filter to explicitly model uncertainty and mitigate error accumulation during iterative prediction. We provide theoretical analyses showing that the proposed interpolator satisfies PDE-constrained smoothness properties, and that the forecaster converges under the proposed loss formulation. Extensive experiments on multiple dynamical datasets demonstrate that PDYffusion achieves superior performance in terms of CRPS and MSE, while maintaining stable uncertainty behavior measured by SSR. We further analyze the inherent trade-off between prediction accuracy and uncertainty, showing that our method provides a balanced and robust solution for long-horizon forecasting.

• The paper introduces PDYffusion, a framework that integrates PDE regularization with an Unscented Kalman Filter to maintain physical consistency in long-horizon forecasts.
• It employs a Matérn-kernel-based interpolator and rigorous theoretical analysis to ensure convergence and improved accuracy compared to standard diffusion methods.
• Empirical evaluations on diverse datasets demonstrate enhanced CRPS, MSE, and SSR, underscoring the model's precision and stability in complex dynamical systems.

PDE-Regularized Dynamics-Informed Diffusion for Long-Horizon Dynamics

Introduction

Accurate long-horizon forecasting of spatiotemporal dynamical systems governed by PDEs is a central challenge in scientific machine learning and physical simulation. Classical numerical solvers are limited by computational costs when scaling to high-dimensional systems across extended time intervals. Recent work on data-driven models, such as RNNs, LSTMs, and diffusion models, demonstrates significant progress; however, these methods are often susceptible to error accumulation, loss of physical consistency, and instability, particularly for long-range prediction tasks. The work "PDE-regularized Dynamics-informed Diffusion with Uncertainty-aware Filtering for Long-Horizon Dynamics" (2604.09058) introduces PDYffusion, a novel probabilistic framework that tightly integrates PDE constraints and sequential uncertainty correction, addressing core deficiencies in prior approaches.

Framework Overview

PDYffusion is architected with two principal modules: a PDE-regularized interpolator and an uncertainty-aware forecaster based on the Unscented Kalman Filter (UKF). The interpolator enforces physically meaningful structure on intermediate state reconstructions by embedding a Matérn-kernel-based differential operator directly into the training loss. Subsequently, the forecaster leverages UKF to propagate uncertainty, carrying out nonlinear state estimation at each prediction step to mitigate error amplification common in deep autoregressive or generative models.

Figure 1: PDYffusion framework: The interpolator enforces structure via $(E-l^{-2}\Delta)^{\alpha/2}$; the UKF-based forecaster models and corrects uncertainty for long-horizon prediction.

The use of PDE-informed regularization differentiates the approach from standard MSE-optimized interpolation, ensuring the learned representation adheres to the physical laws characterizing the underlying dynamical process. The UKF forecaster is designed to explicitly capture and update the predictive distribution, correcting the trajectory through deterministic sigma-point propagation and weighted Bayesian update, thus maintaining uncertainty calibration and state fidelity over long-rollout regimes.

Theoretical Analysis

The authors provide a comprehensive theoretical analysis for both core modules. The PDE-regularized interpolator is shown to minimize the discrepancy between the predicted and groundtruth distributions, as quantified by squared Maximum Mean Discrepancy (MMD) in a Reproducing Kernel Hilbert Space defined by the Matérn kernel. The Matérn kernel is used precisely because it encodes the covariance structure induced by the governing SPDE of the dynamics. The interpolator loss is defined as the sum of a pointwise reconstruction term and a PDE regularization term, where the balance is modulated by a hyperparameter $\lambda$.

For the UKF-based forecaster, convergence is analyzed via the negative log-likelihood of the predicted terminal state under the UKF posterior. The objective encompasses both the Mahalanobis distance and the log-determinant of the covariance, penalizing both error magnitude and uncertainty miscalibration. Theoretical results confirm that, under standard regularity conditions, joint optimization of the interpolator and forecaster ensures convergence of predicted states, sigma points, and covariance estimates to their groundtruth values, rendering the combined loss asymptotically zero.

Empirical Evaluation

PDYffusion is evaluated on four datasets of increasing dynamical complexity: Sea Surface Temperature (SST), Navier–Stokes, Spring-mesh, and a synthetic Wave system, each emphasizing different aspects of physical dynamics, measurement noise, and long-range forecastability. Competing baselines include DYffusion, MCVD, DDPM, and traditional uncertainty-aware schemes (perturbation, dropout).

Across Navier–Stokes, Spring-mesh, and Wave datasets, PDYffusion achieves the lowest CRPS and MSE values, indicating the strongest distributional calibration and pointwise accuracy. Notably, CRPS values of 0.059, 0.0092, and 1.82e-3, and MSE values of 0.017, 4.01e-4, and 1.57e-5, respectively, represent strong improvements over all baselines. On SST, a fully observational, real-world dataset characterized by higher noise and environmental variability, the performance advantage is less pronounced—the Markovian MCVD outperforms in some cases, suggesting an interaction between modeling assumptions and data modality that warrants further investigation.

The SSR (Spread-Skill Ratio) metric remains close to the ideal value of 1 for PDYffusion, signifying stable spectral and temporal evolution even over long horizons. The accuracy-stability trade-off is explicitly revealed: tuning the magnitude of Gaussian noise injected into predictions adjusts the balance between MSE and SSR, confirming that improved accuracy and physical-mode stability compete, and must be carefully balanced using the PDE regularization parameter $\lambda$.

Impact of PDE Regularization and Boundary Conditions

Ablation on the regularization parameter $\lambda$ demonstrates that moderate values enable the strongest forecasting performance, while over-regularization degrades both accuracy and dynamical consistency. Analysis across Dirichlet, Neumann, and periodic boundary conditions reveals dataset-dependent preferences: periodic boundary conditions excel on fields with intrinsic periodicity (e.g., SST, Wave), while Dirichlet conditions are preferable for systems with boundary constraints (Navier–Stokes, Spring-mesh). This confirms the necessity of tailoring structural priors to the physics of the targeted system.

Qualitative Analysis

Visualization of wave propagation trajectories highlights the pronounced difference between PDYffusion and comparable MSE-driven methods such as DYffusion. While both models track the groundtruth evolution for initial steps, error accumulation in DYffusion leads to unrecoverable deviations and unstructured oscillations further into the forecast horizon. PDYffusion preserves both amplitude and phase, maintaining the physical shape and spectral content of the wavefield throughout extended simulation intervals.

Figure 2: Wave trajectories: PDYffusion maintains stable double-peak structure and suppresses error accumulation compared to DYffusion in long-horizon rollouts.

Practical and Theoretical Implications

PDYffusion establishes a viable paradigm for physics-informed generative modeling, bridging the gap between data-driven deep architectures and explicit dynamical constraints. By encoding PDE-regularization in the model's bias, it prevents physically implausible modes and enhances extrapolative reliability. The use of UKF-based uncertainty quantification provides a statistically grounded alternative to dropout/ensemble-based strategies, ensuring better calibration and interpretability for downstream decision-making.

Implications extend to a broad class of time-dependent forecasting problems in climate modeling, fluid dynamics, and engineering, where both physical consistency and uncertainty estimation are critical. The approach also presents a path to future integration of stochastic PDEs (SPDEs) and real-time Bayesian data assimilation for online prediction under model misspecification and sensor-induced noise.

Future Directions

The extension of the current framework to intrinsically stochastic systems governed by SPDEs is identified as a research trajectory, encompassing finance and high-dimensional stochastic control. Integrating Bayesian filtering with explicit stochastic sampling mechanisms could further enhance robust uncertainty transport and error management over indefinite horizons.

Conclusion

PDYffusion delivers a theoretically and empirically validated approach for long-horizon dynamical forecasting, embedding PDE-based physical constraints and principled uncertainty filtering within a diffusion generative model structure. The demonstrated improvements in both accuracy and spectral stability are enabled by the fusion of kernel-informed regularization and nonlinear Kalman updates. This framework advances the state-of-the-art in physically consistent, uncertainty-aware learning for complex dynamical trajectories and lays the foundation for further research at the interface of generative modeling and scientific computing.