PDESpectralRefiner: Achieving More Accurate Long Rollouts with Spectral Adjustment

TL;DR

PDESpectralRefiner introduces spectral adjustments to diffusion-based models to improve the accuracy and stability of long-term PDE simulations, especially for complex equations like Navier-Stokes.

Contribution

The paper proposes a spectral adjustment method for diffusion models, enhancing PDERefiner's ability to accurately model high-frequency components in long PDE rollouts.

Findings

Spectral adjustment improves accuracy for complex PDEs.

PDESpectralRefiner outperforms baseline models in MSE and rollout loss.

Effective with various neural network backbones.

Abstract

Generating accurate and stable long rollouts is a notorious challenge for time-dependent PDEs (Partial Differential Equations). Recently, motivated by the importance of high-frequency accuracy, a refiner model called PDERefiner utilizes diffusion models to refine outputs for every time step, since the denoising process could increase the correctness of modeling high frequency part. For 1-D Kuramoto-Sivashinsky equation, refiner models can degrade the amplitude of high frequency part better than not doing refinement process. However, for some other cases, the spectrum might be more complicated. For example, for a harder PDE like Navior-Stokes equation, diffusion models could over-degrade the higher frequency part. This motivates us to release the constraint that each frequency weighs the same. We enhance our refiner model with doing adjustments on spectral space, which recovers Blurring diffusion models. We developed a new v-prediction technique for Blurring diffusion models, recovering the MSE training objective on the first refinement step. We show that in this case, for different model backbones, such as U-Net and neural operators, the outputs of PDE-SpectralRefiner are more accurate for both one-step MSE loss and rollout loss.