Paper: Adversarial Adaptive Sampling: Unify PINN and Optimal Transport for the Approximation of PDEs (ICLR-2024)

Reviewer comment (Theory):
The authors mention that this methodology achieves variance reduction; what is the proof for that? Is it similar in flavor to Tang et al (2023)?

Author response:
In AAS, in the training procedure, the residual will approach the uniform distribution, with the balance from the growth of the adversarial sampling density function $p$. This will guarantee a variance reduction of residual, because if the residual is uniformly distributed, it is constant in the domain, and a constant function has $0$ variance with respect to any probability measure. This is why we didn't explain it in the paper.

Reviewer comment (Theory):
In Assumption 1, the authors refer to $r$ as an operator. Is that an operator, or the residual between the loss? If the authors mean that $r$ is an operator, then from which function space to what other function space? (Presumably it is U x F → U).

Author response:
Note that $r$ is the residual of PDEs. Given the partial differential operator $\mathcal{L}$, and the right-hand term $s$, for any function $u \in C^{\infty}(\Omega)$, then the residual defines an operator, which is interpreted as
    $r(u) = \mathcal{L}u - s$.

Reviewer comment (Writing):
In Theorem 1, $\nu$ is used to define the density of the residuals, whereas in the derivation under Equation 6, it is defined by $\mu_r$. It seems they are the same quantity; however, if not, what is $\nu$ and what is $\mu_r$?

Author response:
In theorem 1, $\mu$ represents the uniform distribution. $\nu$ represents the residual-induced distribution. To avoid confusion, we have revised the manuscript, where $\nu$ is replaced by $\mu_r$.

Reviewer comment (Writing):
It is unclear what is being approximated by a KRNet, since the authors first use $p_\alpha$ to define the density that they are trying to train; however, there is no mention of it after Section 3.2.

Author response:
Here, $p_{\boldsymbol{\alpha}}$ is a PDF modeled by KRnet. Solving PDEs with neural networks needs to generate collocation points in the computational domain for training. These collocation points are drawn from $p_{\boldsymbol{\alpha}}$ defined by KRnet—different sets of collocation points correspond to different KRnets. So, the maximization step in equation 8 is the training procedure for the KRnet modeled pdf, see section 3.3.

Reviewer comment (Novelty):
KRNet is not cited in the last paragraph of page 3.

Author response:
We here use a bounded KRnet, and the citation is on page 4: ``The bounded KRnet can be achieved by adding a logistic transformation layer (Tang et al., 2023) or a new coupling layer proposed in (Zeng et al., 2023)". We have also added a reference on page 3 in the revised manuscript.

Reviewer comment (Writing):
In Equation 9, what is p*?

Author response:
For a fixed $u_{\boldsymbol{\theta}}$, equation (9), the Euler–Lagrange equation for the functional in equation (8), gives the necessary condition for the existence of optimal $p*$ of the maximization step in equation (8), which is a standard technique called the calculus of variations in analysis. We write this equation to illustrate that the maximization step is well-posed after adding a gradient penalty term (see equation 8), which is a weakened condition for the Lipschitz norm condition. The regularization of it is then guaranteed, for otherwise, the maximization step will yield a delta measure (see section 3.1).

Reviewer comment (Reproducibility):
What are the implementation details (in terms of number of parameters, etc.) of the networks approximating the solution $u_\theta$ and approximating the density network, i.e., $p_\alpha$?

Author response:
Due to the page limitation, we do not present the details of the network structure. All the models are set to be the same as those in DAS-PINNs. We have stated this issue in the revised manuscript.

Reviewer comment (Writing):
While the methodology provided by the authors seems to do better than the baselines, given the presentation and the lack of clarity of the precise steps, the paper would benefit from a revision.

Author response:
Thanks for your reply. We will check the clarifications in the revised manuscript and proofread it again.

Reviewer comment (Novelty):
"[...] which is the first time to minimize the residual and seek the optimal training set simultaneously for PINN.": Analogous to DAS-PINNs, it seems that the final algorithm is still alternating between optimizing the two networks, see also Question 2) below. In general, it should be made more explicit what the novelty of the present work is and how it theoretically compares to related work.

Author response:
The current algorithm allows the two models (PINN and KRnet) to evolve at the same time. In the work of DAS-PINNs, the style of adaptive sampling is like that of the classical methods such as adaptive finite element methods. That is, the procedure of adaptive sampling in DAS-PINNs is like $\mathsf{solve} \rightarrow \mathsf{refine} \rightarrow \mathsf{solve} \rightarrow \ldots$. Simply speaking, given a training set, we need to train PINN with a large number of epochs to ensure that we have ``solved" the PDE, and then we use the KRnet-induced distribution to approximate the residual-induced distribution by minimizing the KL divergence between them. However, in this work, we do not require the minimization step to have a large number of epochs. Moreover, the min-max formulation in AAS implicitly pushes the residual-induced distribution to a uniform one without any KL divergence-like formulation, while there is an explicit minimization step about KL divergence in DAS-PINNs. In summary, AAS's formulation and training style are different from DAS-PINNs.

Reviewer comment (Evaluation):
For the numerical results, training times and standard deviations (w.r.t. different seeds) are missing. Especially given that the adversarial training slows down training.

Author response:
Thanks for the suggestion. Based on our available computational resources, we have taken three runs with different random seeds for the proposed AAS method and computed the mean error and the standard deviation (std) of the three runs. Also, we record the mean training time of the three runs for the three test problems presented in the manuscript. The results are as follows.  
|   | Mean error with std   | Training time |  
|  ----  | ----  | ----  |   
| Equation 12  | 3.1754e-05 $\\pm$ 1.1272e-05 |   4.82 hours |  
| Equation 13  |  6.4194e-04 $\\pm$ 4.6891e-04 |  5.41 hours |  
| Equation 14  |  0.0024 $\\pm$ 0.0015 |  15.74 hours |

Reviewer comment (Theory):
The loss for learning the normalizing flow seems to use the REINFORCE trick, which, however, is known to suffer from high variance.

Author response:
We are not quite sure about what the reviewer means by ``use the REINFORCE trick". The normalizing flow model generates samples for training PINN. Still, we do not require the flow model to have a high-accuracy performance since our goal is to solve PDEs instead of density estimation. It is enough if the normalizing flow model can capture where the low regularity of the solution is. In addition, other deep generative models with an explicit PDF can be applied here.

Reviewer comment (Novelty):
Since the residual continuously changes over the course of optimization, it should be better motivated what the advantage of learning a generative modeling is (as compared to just using a method to sample, e.g., according to squared residual)?

Author response:
Sampling data points according to the residual is suitable for low-dimensional problems. For high-dimensional problems, using sampling methods such as MCMC is not efficient, and this is why deep generative models (such as GAN, VAE, Flow, and diffusion models) are proposed.

Reviewer comment (Theory):
It seems that the precise connection to optimal transport remains a bit unclear: "So the minimization step will reduce not only the residual but also the Wasserstein distance between $\mu_r$ and the uniform distribution". Since there is only an upper bound shown on page 6, it is unclear why the minimization step is guaranteed to decrease the Wasserstein distance.

Author response:
This is just an illustrative analysis to show how the idea works. For more details, see Appendix A.1-A.4. We have provided the detailed analysis about this question.

Reviewer comment (Theory):
It seems that the precise connection to optimal transport remains a bit unclear: "The evolution of the residual-induced distribution has a clear path": According to the theorem, there only *exists* such a path, and it is not clear whether this path is taken by GD. It would be good to have at least plots of the residual evolution for all experiments.

Author response:
The theorem tells us that the residual-induced distribution converges to the uniform distribution. As long as the maximization of the sequence $\{p_n\}$ is efficiently solved, by substituting it into the integral, the minimization process is to force the residual to converge to the uniform distribution. Then, GD is just one way to make the parameterized network function to achieve this process, as in any other machine learning methods. In addition, the variance of the residual function values will converge to zero, i.e., the residual at the collocation points are almost the same. We have plotted the variance of the residual in the manuscript (see Figure 2(b), the left plot of Figure 4, and Figure 5(b)).

Reviewer comment (Presentation):
Minor issue — use `\citep`.

Author response:
Thanks for pointing out this issue. We have fixed it in the revised manuscript.

Reviewer comment (Presentation):
Minor issue — In Eq. (4), min-max should be written on both sides of the equation.

Author response:
Thanks for pointing out this issue. We have fixed it.

Reviewer comment (Writing):
Minor issue — What is a 'proper' constraint?

Author response:
The proper constraint for the PDF model $p_{\boldsymbol{\alpha}}$ is used to make the min-max formulation well-posed. Otherwise, the maximization step will simply yield a delta measure (see section 3.1). To remedy this issue, we can enforce a constraint on $p_{\\boldsymbol{\\alpha}}$. Here, the constraint is $p_{\\boldsymbol{\\alpha}} \\geq 0, \\int_{\\Omega} p_{\\boldsymbol{\\alpha}} d \\boldsymbol{x} =1, ||p_{\\boldsymbol{\\alpha}}||_{Lip} \\leq 1$, see section 3.2.

Reviewer comment (Writing):
Minor issue — It is unclear why we would consider $I^D$ instead of $\Omega$ when the KRnet is first introduced.

Author response:
Thanks for pointing it out. We focus on the problem that the solution has low-regularity, which is used to demonstrate the important role of adaptive sampling in the framework of PINN. In this work, the computational domain is restricted to $[-1,1]^D$, i.e., $\\Omega = [-1,1]^D$. This is because we want to rule out the other factors such as complex boundaries to restrict our attention to adaptive sampling. We have added the comments (``In this work, we consider $\\Omega = I^D$ for simplicity") to the revised manuscript.

Reviewer comment (Writing):
Minor issue — typo: resultd

Author response:
We have fixed this typo and proofread the manuscript again.

Reviewer comment (Experiments):
Question — Why is RAR not tested for the PDEs in (12) and (13)?

Author response:
We mainly use low- dimensional and low-regularity test problems to demonstrate that the sampling strategy affects significantly the performance of neural network approximation if the residual is strongly localized. For low-dimensional problems, RAR is simple, easy to implement, and sufficient. In RAR, candidate points to add into the current training set are also drawn from the uniform distribution. The candidate points who have large residuals are selected to add into the training set to refine the collocation points. However, for high-dimensional problems, generating the candidate points by the uniform distribution is not efficient, sometimes failed, but AAS and DAS can efficiently generate the data points to refine the collocation points thanks to the deep generative model. This is also discussed by the review paper \emph{A comprehensive study of non- adaptive and residual-based adaptive sampling for physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 403:115671, 2023.} See the conclusion of this paper. The authors cite the work of DAS-PINNs: “In this study, we sample residual points in RAD and RAR-D by using a brute-force approach, which is simple, easy to implement, and suﬀicient for many PDEs. However, for high-dimensional problems, we need to use other methods, such as generative adversarial networks (GANs) [46], as was done in Ref. [41]”, where Ref. [41] indicates the DAS method.

Reviewer comment (Theory):
Question — Why can we not use a single backward pass on the loss in (11) and update both $\theta$ and $\alpha$ (using decent and ascent, respectively)?

Author response:
It will lead to divergence. The training style is similar to WGAN, where the two model parameters are updated alternatively to ensure convergence.