DAS-PINNs: A deep adaptive sampling method for solving high-dimensional partial differential equations

We outline our basic ideas on variance reduction in this section and more details about the algorithm will be presented later. We first consider the importance sampling technique. For simplicity, we only consider $J_{r}(u(\boldsymbol{x};\Theta))$ in equation (2), which is the loss induced by the residual. We have

where the set $\{\boldsymbol{x}_{\Omega}^{(i)}\}_{i=1}^{N_{r}}$ is generated with respect to the probability density function (PDF) $p(\boldsymbol{x})$ instead of a uniform distribution as in equation (3). If the variance of $r^{2}(\boldsymbol{X})p^{-1}(\boldsymbol{X})$ in terms of $p(\boldsymbol{x})$ is smaller than the variance of $r^{2}(\boldsymbol{X})$ in terms of the uniform distribution, the accuracy of the Monte Carlo approximation will be improved for a fixed sample size $N_{r}$. The optimal choice for $p(\boldsymbol{x})$ is

where $\mu=\int_{\Omega}r^{2}(\boldsymbol{x};\Theta)d\boldsymbol{x}$. Although the optimal choice is useless in practice since $\mu$ is the quantity to be computed, it suggests that variance reduction can be achieved if $p(\boldsymbol{x})$ is close to the residual-induced distribution.

When $p(\boldsymbol{x})$ is close to $p^{*}(\boldsymbol{x})$, the ratio $r^{2}(\boldsymbol{x})/p(\boldsymbol{x})$ should be close to $\mu$. The second term on the right-hand side of inequality (13) is related to the tail probability $\mathbb{P}(|r^{2}/p-\mu|>a;p)$ in terms of $p(\boldsymbol{x})$. It is seen that if the tail probability is small enough, a variance reduction can be achieved by choosing a small $a$.

Another option for variance reduction is to relax the definition of $J_{r}(u)$ as:

where the set $\{\boldsymbol{x}_{\Omega}^{(i)}\}_{i=1}^{N_{r}}$ is sampled from the PDF $p(\boldsymbol{x})$ and $p(\boldsymbol{x})>0$ on $\Omega$. It is easy to see that the minimizer of $J_{r,p}(u(\boldsymbol{x};\Theta))$ is also the solution of problem (1) when the boundary conditions are satisfied. To reduce the error induced by the Monte Carlo approximation, we may adjust $p(\boldsymbol{x})$ such that the residual $r^{2}(\boldsymbol{x};\Theta)$ does not vary dramatically. This is similar to the classical adaptive finite element method, where mesh refinement/coarsening is supposed to make the approximation error nearly uniform. For our case, we may sample a PDF that is close to the residual-induced distribution and add more samples from the region of large residual into the training set.

Two options for variance reduction are considered in the previous section. To make both ideas practical, the key issue is to generate samples efficiently from a PDF $p(\boldsymbol{x})\approx\mu^{-1}r^{2}(\boldsymbol{x};\Theta)$ for a fixed $\Theta$. The first option is based on importance sampling, meaning that an explicit PDF model $p(\boldsymbol{x})$ is needed, The second option is based on refinement of the training set, which only needs samples from the region of large residual and does not need the likelihood of the samples. The second option has been employed in the recent literature to improve the neural network approximation, which is either ad hoc [42] (only for low-dimensional problems) or based on traditional sampling strategies such as MCMC [43]. To the best of our knowledge, there does not exist a generic algorithm that can be effectively adapted to both options. We intend to fill this gap in this work. It is seen from Lemma 2 that the tail probability $\mathbb{P}(|r^{2}/p-\mu|>a;p)$ should be small enough for the effectiveness of $p(\boldsymbol{x})$, which means that $p(\boldsymbol{x})$ must be close enough to the distribution induced by $r^{2}(\boldsymbol{x})$. However, the approximation of PDF is a challenging task especially in high-dimensional spaces. Classical explicit PDF models such as the exponential family of distributions and Gaussian mixture models are in general not sufficient as an approximator for the PDF induced by $r^{2}(\boldsymbol{x})$. To alleviate this difficulty, we resort to deep generative modes. In particular, we employ a recently developed deep generative model called KRnet for both probability approximation and sample generation [22, 27]. KRnet is one type of normalizing flow [44], which provides an invertible transport map between a prior distribution and the target distribution. Unlike other types of deep generative models such as GAN [45, 46, 47] and VAE [48], normalizing flow provides an explicit likelihood. KRnet can be used to address both options for variance reduction while other deep generative models may also be employed if only the second option for variance reduction is considered. However, one computational issue faced by all deep generative models is that the distribution induced by $r^{2}(\boldsymbol{x})$ is usually defined on a compact domain while deep generative models are in general defined on the whole space. We subsequently address this issue without going into details about the structure of KRnet.

Let $\boldsymbol{X}\in\mathbb{R}^{d}$ be a random vector associated with a given data set, and its PDF is denoted by $p_{\boldsymbol{X}}(\boldsymbol{x})$. The target is to estimate $p_{\boldsymbol{X}}(\boldsymbol{x})$ from data or to generate samples that are consistent with a given $p_{\boldsymbol{X}}(\boldsymbol{x})$. Let $\boldsymbol{Z}\in\mathbb{R}^{d}$ be a random vector associated with a PDF $p_{\boldsymbol{Z}}(\boldsymbol{z})$, where $p_{\boldsymbol{Z}}(\boldsymbol{z})$ is a prior distribution (e.g., Gaussian distributions). The flow-based generative modeling is to seek an invertible mapping $\boldsymbol{z}=f(\boldsymbol{x})$ [31]. By the change of variables, we have the PDF of $\boldsymbol{X}=f^{-1}(\boldsymbol{Z})$ as

Once the prior distribution $p_{\boldsymbol{Z}}(\boldsymbol{z})$ is specified, equation (15) provides an explicit density model for $\boldsymbol{X}$. The inverse of $f(\cdot)$ provides a convenient way to sample $\boldsymbol{X}$ as $\boldsymbol{X}=f^{-1}(\boldsymbol{Z})$. The basic idea of KRnet is to define the structure of $f(\boldsymbol{x})$ in terms of the Knothe-Rosenblatt rearrangement. Let $\mu_{\boldsymbol{Z}}$ and $\mu_{\boldsymbol{X}}$ be the probability measures of two random variables $\boldsymbol{Z},\boldsymbol{X}\in\mathbb{R}^{d}$ respectively. A mapping $\mathcal{T}$: $\boldsymbol{Z}\mapsto\boldsymbol{X}$ is called a transport map such that $\mathcal{T}_{\#}\mu_{\boldsymbol{Z}}=\mu_{\boldsymbol{X}}$, where $\mathcal{T}_{\#}\mu_{\boldsymbol{Z}}$ is the push-forward of $\mu_{\boldsymbol{Z}}$ such that $\mu_{\boldsymbol{X}}(B)=\mu_{\boldsymbol{Z}}(\mathcal{T}^{-1}(B))$ for every Borel set $B$ [49]. The transport map $\mathcal{T}$ given by the Knothe-Rosenblatt rearrangement [49, 30] has a lower-triangular structure

Simply speaking, KRnet integrates the triangular structure of the K-R rearrangement into the definition of the invertible transport map. More details can be found in [27, 22, 29].

Let $f_{\mathsf{KRnet}}(\cdot;\Theta_{f})$ indicate the invertible transport map induced by KRnet, where $\Theta_{f}$ includes the model parameters. An explicit PDF model $p_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ can be obtained by letting $f=f_{\mathsf{KRnet}}$ in equation (15), i.e.,

The samples of $p_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ is given as $\boldsymbol{X}=f^{-1}_{\mathsf{KRnet}}(\boldsymbol{Z})$ by sampling $\boldsymbol{Z}$. A common choice for the distribution of $\boldsymbol{Z}$ is the standard Gaussian distribution. Depending on the prior knowledge of the problem, a more general model such as Gaussian mixture model can also be used as the prior distribution. The KRnet $f_{\mathsf{KRnet}}$ does not have any constraint on the range of the mapped data, meaning that both $\boldsymbol{X}$ and $\boldsymbol{Z}$ are defined on $\mathbb{R}^{d}$. Let $\boldsymbol{Z}$ be Gaussian. Due to the invertibility, $|\det\nabla_{\boldsymbol{x}}f_{\mathsf{KRnet}}|>0$ for any $\boldsymbol{x}\in\mathbb{R}^{d}$, which implies that $p_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})>0$ for any $\boldsymbol{x}\in\mathbb{R}^{d}$. So $p_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ is not consistent with the distribution induced by $r^{2}(\boldsymbol{x})$, which is equal to zero on $\mathbb{R}^{d}\backslash\Omega$. To deal with this issue, we propose the following strategy.

Without loss of generality, we can assume that $\Omega=[-1/2,1/2]^{d}$. Let $B=(-(1/2+\delta/2),1/2+\delta/2)^{d}$ with $0<\delta<\infty$ such that $\Omega\subset B$. For each dimension of $\boldsymbol{x}$, we define the following logarithmic mapping

with $s>0$ being a scale parameter, which defines a one-to-one correspondence between $x\in(-(1/2+\delta/2),1/2+\delta/2)$ and $y\in(-\infty,+\infty)$. Let $\boldsymbol{\ell}(\boldsymbol{x}):B\mapsto\mathbb{R}^{d}$ be a $d$-dimensional mapping such that

Then the following invertible mapping

defines a PDF

where the support of $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ is $B$.

We now consider a modification of $r^{2}(\boldsymbol{x};\Theta)$. Define a cutoff function as

where $h_{\delta}(x)$ is a piecewise linear function

We consider a modified PDF for any $\Theta$

Note that both $\hat{r}_{\boldsymbol{X}}(\boldsymbol{x})$ and $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ have the support $B$. We then solve the following optimization problem

where $D_{\mathsf{KL}}(\cdot\|\cdot)$ indicates the Kullback-Leibler (KL) divergence between two distributions. We finally use

as an approximation of the PDF induced by $r^{2}(\boldsymbol{x};\Theta)$. If $\delta\ll 1$, most the samples $\boldsymbol{\ell}^{-1}\circ f^{-1}_{\mathsf{KRnet}}(\boldsymbol{z}^{(i)};\Theta_{f}^{*})$ will be located in $\Omega$. Since $\hat{r}_{\boldsymbol{X}}(\boldsymbol{x})\propto r^{2}(\boldsymbol{x};\Theta)$ on $\Omega$, $p_{\boldsymbol{X}}(\boldsymbol{x})$ approximates the $r^{2}(\boldsymbol{x};\Theta)$-induced PDF well when $\delta$ is small. In our numerical experiments, we set $\delta=0.01$ and $s=2$.

We now look at the approximation of $\Theta_{f}^{*}$. The KL divergence in the optimization problem (21) can be written as

The first term on the right-hand side corresponds to the differential entropy of $\hat{r}_{\boldsymbol{X}}$, which does not affect the optimization with respect to $\Theta_{f}$. So minimizing the KL divergence is equivalent to minimizing the cross entropy between $\hat{r}_{\boldsymbol{X}}$ and $\hat{p}_{\mathsf{KRnet}}$ [50, 51]:

Since the samples from $\hat{r}_{\boldsymbol{X}}$ are not available, we approximate the cross entropy using the importance sampling technique:

where $\hat{p}(\boldsymbol{x};\hat{\Theta}_{f})$ is a PDF model with known parameter $\hat{\Theta}_{f}$ and its samples $\{\boldsymbol{x}_{B}^{(i)}\}_{i=1}^{N_{r}}$ can be generated efficiently as

with $\boldsymbol{z}^{(i)}$ being sampled from the prior distribution. We then minimize the discretized cross entropy (25) to obtain an approximation of $\Theta_{f}^{*}$. The choice for $\hat{\Theta}_{f}$ will be specified in section 4.3 when the adaptive sampling method is defined.

We are now ready to present our algorithms. In this work, we mainly focus on the adaptivity of $\mathsf{S}_{\Omega}$ for simplicity. The key step of our adaptivity strategy is to improve the effectiveness of the random samples in the training set $\mathsf{S}_{\Omega}$, and we provide two algorithms corresponding to the two options discussed in section 4.1.

1. I.

   We replace all the collocation points in the current training set using new samples from the probability measure for importance sampling. This corresponds to equation (11).

2. II.

   We gradually add more collocation points to the current training set. This corresponds to equation (14), where the new samples are mainly from the region of large residual.

We first present strategy I. Let $\mathsf{S}_{\Omega,0}=\{\boldsymbol{x}_{\Omega,0}^{(i)}\}_{i=1}^{N_{r}}$ and $\mathsf{S}_{\partial\Omega,0}$ be two sets of collocation points that are uniformly sampled from $\Omega$ and $\partial\Omega$ respectively. Using $\mathsf{S}_{\Omega,0}$ and $\mathsf{S}_{\partial\Omega,0}$, we minimize the empirical loss (3) to obtain $u(\boldsymbol{x};\Theta_{N}^{*,(1)})$. With $u(\boldsymbol{x};\Theta_{N}^{*,(1)})$, we minimize the cross entropy (25) to get $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(1)})$, where we simply use uniform samples for importance sampling. To refine the training set, a new set $\mathsf{S}_{\Omega,1}=\{\boldsymbol{x}_{\Omega,1}^{(i)}\}_{i=1}^{N_{r}}$ is generated by $\boldsymbol{\ell}^{-1}\circ f^{-1}_{\mathsf{KRnet}}(\boldsymbol{z}^{(i)};\Theta_{f}^{*,(1)})$ (see equation (18)). Then we continue to update the approximate solution $u(\boldsymbol{x};\Theta_{N}^{*,(1)})$ using $\mathsf{S}_{\Omega,1}$ as the training set. In general, we use $\mathsf{S}_{\Omega,k}=\{\boldsymbol{x}_{\Omega,k}^{(i)}\}_{i=1}^{N_{r}}$ to obtain $u(\boldsymbol{x};\Theta_{N}^{*,(k+1)})$ as

where $u(\boldsymbol{x};\Theta)$ is initialized as $u(\boldsymbol{x};\Theta_{N}^{*,(k)})$ and $J^{\mathsf{IS}}_{N}$ is defined as

Starting with $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(k)})$, the density model $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$ is updated as

where we let $\hat{\Theta}_{f}=\Theta_{f}^{*,(k)}$ in equation (25), i.e., the previous PDF model $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(k)})$ is used for importance sampling when computing the cross entropy. A new set $\mathsf{S}_{\Omega,k+1}=\{\boldsymbol{x}_{\Omega,k+1}^{(i)}\}_{i=1}^{N_{r}}$ of collocation points is then generated. As detailed in section 4.2, the support of data points generated by KRnet is $B=(-(1/2+\delta/2),1/2+\delta/2)^{d}$, while the computation domain is $\Omega=[-1/2,1/2]^{d}$. So we need to deal with the collocation points located in $B\backslash\Omega$. Instead of neglecting these points, we project them onto $\partial\Omega$. We define an entry-wise projection operator $\mathcal{P}(\boldsymbol{x}):B\mapsto\Omega$ as

For a sequence of i.i.d. samples $\boldsymbol{z}^{(j)}$ generated from the standard Gaussian with $j=1,2\ldots$, we compute $\boldsymbol{x}_{B}^{(j)}=\boldsymbol{\ell}^{-1}\circ f^{-1}_{\mathsf{KRnet}}(\boldsymbol{z}^{(j)})$. if $\boldsymbol{x}_{B}^{(j)}=\mathcal{P}(\boldsymbol{x}_{B}^{(j)})$, we assign $\boldsymbol{x}_{B}^{(j)}$ to $\mathsf{S}_{\Omega,k+1}$; otherwise, we add $\mathcal{P}(\boldsymbol{x}^{(j)}_{B})$ to $\mathsf{S}_{\partial\Omega,k}$. The updated training set $\mathsf{S}_{\Omega,k+1}$ and $\mathsf{S}_{\partial\Omega,k+1}$ will be used for the next training stage. This procedure is repeated until the stopping criterion is satisfied (see Algorithm 1). Since the collocation points in $\mathsf{S}_{\Omega,k}$ will be completely replaced at the next stage, we call this type of deep adaptive sampling strategy DAS-R for short. The alternative approach given in Remark 1 can also be used to obtain $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(k)})$.

We now look at strategy II. Unlike DAS-R, the number of collocation points in the training set $\mathsf{S}_{\Omega}$ increases gradually. So we denote this type of deep adaptive sampling strategy by DAS-G for short. Starting with an initial set of collocation points $\mathsf{S}_{\Omega,0}=\{\boldsymbol{x}_{\Omega}^{(i)}\}_{i=1}^{n_{r}}$ (as well as $\mathsf{S}_{\partial\Omega,0}$) drawn from a uniform distribution defined on $\Omega$, we minimize the empirical loss (3) on the training set $\mathsf{S}_{\Omega,0}$ (as well as $\mathsf{S}_{\partial\Omega,0}$) to obtain $u(\boldsymbol{x};\Theta_{N}^{*,(1)})$. Once we have $u(\boldsymbol{x};\Theta_{N}^{*,(1)})$ in hand, we can seek $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(1)})$ using the residual $r^{2}(\boldsymbol{x};\Theta_{N}^{*,(1)})$. We here use uniform samples to approximate and minimize the cross entropy (25). Similar to DAS-R, a new set of collocation points $\mathsf{S}^{g}_{\Omega,1}=\{\boldsymbol{x}_{\Omega,1}^{g,(i)}\}_{i=1}^{n_{r}}$ is generated by $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(1)})$ while the main difference is that we update the training set as $\mathsf{S}_{\Omega,1}=\mathsf{S}_{\Omega,0}\cup\mathsf{S}^{g}_{\Omega,1}$, in other words, $\mathsf{S}_{\Omega,0}$ is augmented rather than replaced by $\mathsf{S}_{\Omega,1}^{g}$. We continue to update $u(\boldsymbol{x};\Theta)$ using $\Theta_{N}^{*,(1)}$ as the initial parameters and $\mathsf{S}_{\Omega,1}$ as the training set, which yields a refined model $u(\boldsymbol{x};\Theta_{N}^{*,(2)})$. Staring from $k=2$, we seek $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f}^{*,(k)})$ using the approach given in Remark 1. We repeat the procedure to obtain an adaptive algorithm (see Algorithm 2).

Our two adaptive training methods are summarized in Algorithm 1 (DAS-R) and Algorithm 2 (DAS-G), where $N_{\rm adaptive}$ is a given number of maximum adaptivity iterations, $m$ is the batch size for stochastic gradient, and $N_{e}$ is the number of epochs for training $u(\boldsymbol{x};\Theta)$ and $\hat{p}_{\mathsf{KRnet}}(\boldsymbol{x};\Theta_{f})$. The algorithms consist of three steps: solving PDE, training KRnet and refining the training set.