Generative Downscaling of PDE Solvers with Physics-Guided Diffusion Models

Abstract

Solving partial differential equations (PDEs) on fine spatio-temporal scales for high-fidelity solutions is critical for numerous scientific breakthroughs. Yet, this process can be prohibitively expensive, owing to the inherent complexities of the problems, including nonlinearity and multiscale phenomena. To speed up large-scale computations, a process known as downscaling is employed, which generates high-fidelity approximate solutions from their low-fidelity counterparts. In this paper, we propose a novel Physics-Guided Diffusion Model (PGDM) for downscaling. Our model, initially trained on a dataset comprising low-and-high-fidelity paired solutions across coarse and fine scales, generates new high-fidelity approximations from any new low-fidelity inputs. These outputs are subsequently refined through fine-tuning, aimed at minimizing the physical discrepancies as defined by the discretized PDEs at the finer scale. We evaluate and benchmark our model’s performance against other downscaling baselines in three categories of nonlinear PDEs. Our numerical experiments demonstrate that our model not only outperforms the baselines but also achieves a computational acceleration exceeding tenfold, while maintaining the same level of accuracy as the conventional fine-scale solvers.

Appendices

Appendix

Table 5 Table of notations

Neural Networks Architecture and Hyperparameters

Our diffusion models are based on the DDPM architecture [14], which uses U-Net [33] as the backbone (Table 5). During our experiments, we omit the use of self-attention, resulting in significant reductions in training time while maintaining similar sample quality. The base channel count, the list of Down/Up channel multipliers and the list of middle channel refer to the hyperparameters of the U-Net, which is detailed in Table 6. To accelerate sampling process using DDIM, we take skipped time steps $\tau$ be $[1, 5, 10, 15, 20, 25, \cdots , T-5, T]$. The linear noise schedule is configured from $\beta _0 = 0.0001$ to $\beta _{T}=0.02$.

During training, we utilize the Adam optimizer with a dynamic learning rate that linearly decays every 5000 steps with a decay rate of 0.05. The total number of training epochs is set to 10000.

Table 6 Table of DDPM hyperparameters

The architecture of FNO follows that described in [24]. The number of lifting channels, number of FFT truncation modes, and number of Fourier layers for different examples are specified in Table 7. During training, we utilize the Adam optimizer with a dynamic learning rate that linearly decays every 5000 steps with a decay rate of 0.05. Training continues until the loss drops below 1e-6 or reaches the maximum iteration number of 50000.

Our model training were performed on an NVIDIA RTX 3070 graphics card, while predictions and refinements with Gaussian-Newton were executed on an AMD Ryzen 7 3700X processor.

Table 7 Table of FNO hyperparameters

Levenberg–Marquardt Algorithm

In this part, we present the Levenberg–Marquardt (LM) algorithm for solving the nonlinear optimization problem (9) and (10) in Algorithm 3. In all of our numerical experiments, we fix $\lambda =0.5$ and $\eta =$1e-5.

Algorithm 3

Levenberg–Marquardt algorithm

Physics-Informed Diffusion Model

For the physics-informed diffusion model [35], our numerical test suggests that conditioning the diffusion model on both gradient information and the coarse solution yields better performance compared to the vanilla PIDM, which is conditioned solely on gradient information. In the application of the PIDM to 2D nonlinear Poisson equation, the conditioning information is defined as the gradient of the $L^2$ misfit, i.e.

$$\begin{aligned} \varvec{g}_t = \frac{\partial r_t}{\partial \varvec{x}_t}, \end{aligned}$$

where

$$\begin{aligned} r_t:= \Vert -0.0005 \varDelta \varvec{x}_t + \varvec{x}_t^3 - \varvec{a}\Vert _2^2. \end{aligned}$$

We employ the same architecture to construct and train the model using Algorithm 1, with a modified loss function:

$$\begin{aligned} \nabla _{\theta } \Big [ \sum _{t=1}^T \frac{1}{B}\sum _{j=1}^B \Vert \varvec{s}_\theta (\varvec{u}^f_{t, j},\varvec{u}_{i}^c, \varvec{g}_{t, j}, t) + \varvec{\epsilon }_{t,j} \Vert ^2 \Big ]. \end{aligned}$$

The gradient guidance strength is set to $w=1$. Various time-step locations $t_s$ in the backward diffusion process were tested ($t_s = [20, 100, 200, 400]$), and it was determined that $t_s = 20$ provides optimal performance, leading to its adoption in the model. As shown Fig. 3, DDPM outperforms PIDM, and we provide some heuristic explanations for this below. In fact, the inputs of the two score networks of the two methods are different. At a specific time t, the score of PIDM takes $\varvec{x}_t, t, \varvec{u}^c, \varvec{g}$ as the inputs, where $\varvec{g}$ is the output of a fixed problem-dependent function of $\varvec{x}_t$ and source term $\varvec{a}$. In contrast, the score of DDPM takes $\varvec{x}_t, t, \varvec{u}^c, \varvec{a}$ as the inputs. Intuitively, including gradient information as an additional input provides more comprehensive information than simply the source term. However, this significantly increases training complexity, especially when the residual function is complicated and the total time step $N_t$ is large, making training much more difficult and potentially leading to poor performance when the training data is limited.

Gaussian-Newton Algorithm

To refine the solution obtained from the coarse solver, diffusion model and the FNO, we introduce the one-step Gaussian-Newton refinement process, outlined in Algorithm 4.

Algorithm 4

One step Gaussian Newton update