Learning An Efficient-And-Rigorous Neural Multigrid Solver

Q10. Is there any nonlinearity (activation) at all in the UGrid submodule? If not, what is the point of using a neural network?

R10. Please refer to R3 (for not using activation) and R10 (for using convolution) to Reviewer pjtp (above).

Q11. Which parts in the multigrid are parametrized (take a two-grid scheme as an example)?

R11. (1) UGrid does not involve a direct solver at the coarsest level; it applies the smoothers directly. Thus, it's not precise to map UGrid to a two-grid cycle with an exact solver. (2) UGrid's learnable convolution operators are correction terms to grid transfer operators. Thus, we would state that $P (P^{\top} A P)^{-1} P^{\top}$ is parametrized as a whole.

Q12. If time for building MG hierarchy for AMGCL/AmgX is included, then training time for UGrid should also be included?

R12. (1) We do include MG hierarchy building time. Yet, it's not fair to include the training time for UGrid, because: (1.1) Training is required only once for one type of PDEs, and could be ignored over the solver's lifespan; (1.2) AMGCL/AmgX must reconstruct their MG hierarchy when input grid or boundary geometry changes; UGrid doesn't need retraining for these cases; (2) When MG hierarchy is constructed already, and only RHS changes, AMGCL/AmgX performs better. However, in fields like PDE-based CAD, grid and boundary-geometry changes as frequently as RHS, and UGrid will be a better choice.

Q13. "Sharp feature" in Fig. 3 (g) does not seem sharp; are there discontinuous effects? Are the region of the two circles included in the computational domain? Do they have a discontinuous diffusion coefficient?

R13. Please refer to R5 (above) and Fig. 5 (on pp.13).

Q14. No information on what the mesh is like. If the grid is uniform, then one can learn prolongation and restriction operators in the geometric multigrid. If it is AMG, the coarsening can be trained.

R14. We use sub-uniform unstructured mesh. The grid points are distributed uniformly inside an irregular region with non-trivial geometry and topology, so the grids are unstructured. It is far from trivial to apply GMG directly on our unstructured testcases. Existing GMG solvers either deal with structured boundaries or employ complicated grid transfer operators which are hard to implement on GPU (e.g. quad-tree-based operators). To the best of our knowledge, there is no GPU-based GMG solver which could handle our testcases, and UGrid outperforms existing AMG solvers.

Q1. The "optimal" smoother should damp "high frequency", while throwing everything into a loss makes the optimal smoother more likely to be damping all frequencies.

R1. (1) Our smoothing operator (Eq. 10) is not learnable. It serves as our mathematical backbone. UGrid's learnable part is a correction term to legacy grid transfer operators. Since we do not optimize smoothers, we won't "make the optimal smoother more likely to be damping all the frequencies". (2) UGrid does preserve high-frequency components, which is shown in testcase "Sharp Feature" (Fig. 5).

Q2. "Jacobi iterator converges for a huge variety of linear PDEs" is wrong.

R2. Sorry for the confusion. It's true that Jacobi does not converge in many situations. That’s why we state "a huge variety", not "most/all". Moreover, this paper handles Poisson & other simple PDEs, and Jacobi does converge to these PDEs. More specifically, under Dirichlet boundary condition, Jacobi generally converges to well-behaved elliptic PDEs & many other linear PDEs like time-variant diffusion/convection equations in homogeneous materials. Thus, this paper is compatible to a variety of PDEs.

Q3. There is no recursion defined in this algorithm at all.

R3. Please refer to R5 to Reviewer pjtp (above).

Q4. For "oscillatory behavior", Thm. 2 should prove $|r|_{\infty} >$ something.

R4. (1) Sorry for the misinterpretation. Thm 2 shows the possible max error when optimizing legacy loss, which explains the existance of solutions with small relative residuals but large legacy errors. Please also refer to R1 to Reviewer Ejzh (ablation study on residual/legacy losses). (2) Denote the prediction as $x$, the ground truth as $y$, $\varepsilon$s as their stochastic relative residuals, and $l_{max}$ as the legacy error between $x$ and $y$:

$$\varepsilon_{x} = \dfrac{||f - A \ x||}{||f||} \approx \dfrac{||f - A \ (y \pm l_{max} y)||}{||f||} = \dfrac{||(f - A \ y) \mp l_{max} \ A \ y)||}{||f||} \ge \left|\dfrac{||f - A \ y||}{||f||} - \dfrac{||l_{max} \ A \ y||}{||f||}\right| = \left|\varepsilon_y - \dfrac{||l_{max} \ A \ y||}{||f||}\right| > 0 \ .$$

Q5. In experiments, $u$ is harmonic (with RHS being zero) and smooth?

R5. Sorry for the confusion, but we should clarify that $u$ is neither harmonic nor smooth. (1) Not Harmonic. $u$ is harmonic only for our training data, and 1/10 of our testcases (Fig. 3 (i)). All other testcases have non-zero Laplacians. (2) Non-zero RHS. See Fig. 4. None of the testcases have zero Dirichlet boundary values. Fig. 4 (g-j) have non-uniform boundary values. (3) Not Smooth. Fig. 3 (c) (e) (h) (j) have jumps in Laplacians (thus C(1) only). For Poisson’s equations, $u$ must be C(1) in inner part (Laplacian operator involves second-order derivatives). $u$ could be discontinuous on the boundary: Testcase "Sharp Feature" (see Fig. 5) introduces "boundary elements" in the inner part of $u$, making $u$ only C(0) ($u$ is defined both inside & outside the circles, and $u$ is treated as a whole; we employ no boundary/interior segmentation.). Testcase "Noisy Input" (Fig. 4 (h)) has discontinous boundary values, making $u$ not even C(0).

Q6. "No strict mathematical guarantee on how fast legacy methods will converge" is wrong.

R6. Sorry for the misinterpretation, we meant to say that "there is no uniform upper bound which could be used to estimate the exact time needed without testing". This is because convergence guarantees are typically asymptotic, not always reachable in practice, e.g., conjugate gradient method often converges faster than what its $O(\sqrt{k})$ complexity suggests.

Q7. It's wrong to say that "nonlinear PDEs generally have no linear iterative solver".

R7. We are sorry for the misinterpretation. We would conduct literature reviews more thoroughly.

Q8. On page 4, which is the easily-invertible approx. of $A$, $P^{-1}$ or $P$?

R8. $P$ is the easily invertible approx. of $A$, and $P^{-1}$ is the inverse of $P$.

Q9. Why the authors make a big deal about the mask matrix M?

R9. (1) M enables us to treat all discrete grid points in the same way, without the need of segmentation. This simplifies the problem. (2) Without M, the Jacobi update kernel will change w.r.t. input geometry, adding more burden to the network. We did not have M in the earliest version of our work, and the network diverged.

Q10. Are you backpropagating the loss through the iterative method?

R10. Yes, and for stability during the training process, we are setting the number of iterations to a fixed number (64) during training.

Q11. Why learn convolutions?

R11. (1) Many differential operators are analogous to convs, e.g., discrete Laplacian operator is a 3x3 conv kernel. This leads to (2) Many components in legacy iterative solvers are also analogous to convolution and thus could be easily implemented with conv layers.

Q12. If the smoothing operation removes the high-frequencies, why not simply have a downsampling operation?

R12. The downsampling operator (grid-transfer operator) is a tricky component in the multigrid hierarchy as the optimal operator differs for different PDEs and boundaries. UGrid's learnable parts are correction terms on legacy grid transfer operators.

Q1. The paper is not well written; I feel the exposition could be greatly simplified. E.g., not all reviewers are familiar with multigrid. Explain in appendix.

R1. We have added more explanations on the multigrid method on pp.12 (Section A.1) of the rebuttal-revision pdf. Moreover, we would like to clarify that numerical approaches for PDEs, including the multigrid method, as well as their integration with AI-based methods, are of fundamental significance to scientific fields. Pointers regarding these topics are available in the main content of the paper, including pp.1 (first two paragraphs of Section 1), which explains the significance of these methods; and pp.6 (Section 4.2), which details our implementation of a multigrid-like structure with CNN.

Q2. Differences w.r.t. Hsieh et al. (2019).

R2. (1) Loss. UGrid is trained with residual loss; Hsieh et al. is trained with legacy loss, which is subject to numerical oscillations (Thm. 2); (2) Learnable Components. Hsieh et al. learns a linear operator H, which works on the finest grid only; UGrid has learnable components on all grid levels; (3) Training Procedure. Hsieh et al. executes their model for one iteration on a preconditioned intermediate result, and optimizes their update vs legacy MG; while UGrid is applied on a raw input for multiple iterations, and we optimize the final residual error; (4) MG Depth. Hsieh et al. only leverage a 3-level MG; UGrid reaches as deep as 6 levels. These differences endow UGrid with larger solution space (generalization power), higher efficiency, and higher precision, which are shown in our experiments.

Q3. "Neural" and "U-Net" mentioned, yet UGrid learns convolutions with no activation functions.

R3. UGrid has no activation functions. It mimics the linear property of legacy solvers and serves as our mathematical backbone. Moreover, Neural Networks could still learn from data without non-linear layers; we admit such a NN coallapses to one single layer. However, this is exactly what we want: We are learning to invert a system matrix ($y = A^{-1} x$).

Q4. What’s the link between convergence of UGrid and convergence of Eq. 8?

R4. (1) UGrid's "convergence guarantee" indicates "if the iterator converges, its result is correct", which is not observed in most neural methods. The UGrid iteration takes Eq. 8 as pre- and post-smoothers; the last operation in a UGrid iteration is post-smoothing, and convergence test of UGrid also evaluates the convergence of Eq. 8. Thus, convergence guarantee of Eq. 8 ensures the correctness of UGrid's output. The convergence guarantee makes our algorithm more mathematically-rigious. (2) Convergence itself (strictly on math) is guaranteed for Eq. 8 (Thm 1). For the whole UGrid routine, it's evaluated by our experiments.

Q5. More explanation on the recursive calls on UGrid submodule. Is the UGrid function called recursively on the same grid? Are there two different types of smoothers?

R5. (1) UGrid is an iterative solver. In its iteration step (Eq. 15), "smooth" refers to Eq. 10, "r" refers to Eq. 11, "UGrid" refers to UGrid submodule (Fig. 2). The UGrid submodule is a CNN and it's better to be visualized than to be expressed in math. (2) UGrid submudule is invoked recursively on a coarser grid, not the same grid (see the bottom of Fig. 2, note the preceeding down-sampling layers). (3) There is only one type of smoother (Eq. 10). UGrid's learnable part is not a smoothing operator, it is designed as a correction term to the legacy grid transfer operators (i.e., down- and up-sampling layers).

Q6. More details regarding the experimental section.

Q7. Please refer to: (1) supplemental materials; (2) R1 & R5 to Reviewer Ejzh (above).

Q8. How is the data generated?

R8. Our training data are Laplacian equations, thus we only need to generate boundary geometries and values. The geometries are generated from circles (with random distortions to mimic random simple closed curves), and we put an extra "hole" inside each curve. The boundary values are piecewise constants. For more details, please refer to "UGrid/script/generate.sh" in our anomynous repo. We would also like to clarify that our testcases have much more complicated boundary geometries/topology/values, and Laplacian distributions, which is illustrated in Fig. 3 and Fig. 4.

Q9. Ablation study.

R9. Please refer to R1 to Reviewer Ejzh (above).

Q1. It is unclear how the method can be extended to non-linear PDEs.

R1. Yes, and we leave the extension to non-linear PDEs as future work.

Q2. No quantitative experiments to support generalization to other types of linear PDEs.

R2. Please refer to R5 to Reviewer Ejzh (above).

Q3. UGrid takes longer to converge on small 'cat' and 'L-shape'. Why does it only happen in these two cases? Does 'cat'/'L-shape' have more similiar problem sizes to the training data than other small cases?

R3. (1) Not all small-scale cases take longer. UGrid performs well on all other eight small-scale testcases (Fig. 3 (b-c) (e-j)). The small "cat"/"L-shape" are outliers. (2) Problem size: Training data are all large-scaled systems of size 1,050,625x1,050,625; small-scale data are of size 66,049x66,049; and "cat" and "L-shape" do not differ from other small cases. (3) Numerical issues are always tricky. It’s hard to analyze efficiency when the lower-level problem, generalization power to unseen problem sizes, is itself more or less relying upon experience, but far less mathematically strict. We would leave this problem to future work.

Q4. Why UGrid performs better on large scale to small scale problems compared to the traditional numerical solvers?

R4. (1) Why better than legacy solvers: Legacy solvers are general-purpose solvers which do not exploit the special properties of specific PDEs, e.g., legacy multigrid solvers use sub-optimal general-purpose grid transfer operators for all PDEs. UGrid is trained on specific PDEs, and it adjusts the grid transfer operation at each level towards the optimal. (2) Why better on large scale: UGrid is trained on the large scale, so it reaches its best potential here. Generalization to small scale is tested by our experiments, not mathematically guaranteed in a strict sense. We leave the analysis of the mathematical details on the small-scale testcases as future work. (3) Traditional MG solvers may become faster if their components (e.g. grid transfer operator) are tailored for specific PDE formulations. It is not possible to ask a single conference paper to do a full justice on all of the above-mentioned criteria and their in-depth theoretical analysis and guarantees. We leave explorations on these possibilities as future work.

Q5. Have you tried applied UGrid on other linear PDEs?

R5. Please refer to R5 to Reviewer Ejzh (above).

Q1. Ablation study on the effects of UGrid & residual loss.

R1. We trained another UGrid with legacy loss (UGrid (L)), and a vanilla U-Net with residual loss, which regresses solution to Poisson's equations directly. The experimental results are as follows:

Residual loss endows UGrid with 2x speed up, and more importantly, stronger generalization power (UGrid (L) diverged on several cases). The U-Net model failed to converge, which demonstrates the significance of UGrid's mathematically-rigorous architecture. More details are available in the rebuttal-revision pdf (pp.14~).

Q2. How UGrid mimics multigrid components?

R2. (1) Smoother: UGrid implements it as masked convolutional iterator (Eq. 8); (2) Grid Transfer Operator: UGrid mimics it with legacy operators plus correction terms (learnable convolution layers); (3) Coarsest-level Solution: UGrid does not invoke direct solvers, as those are not as trivially-differentiable as iterative solvers. We invoke Eq. 8 directly on the coarsest level, and it works.

Q3. Generalization claims overstated. Geometry tests not enough to prove generalization over problem sizes, PDE types and discretization schemes.

R3. (1) We claimed that UGrid generalizes to: (1.1) Boundary geometries/values unseen during training phase (validated by our testcases with different geometries); and (1.2) New problem sizes (validated by the small-scale testcases; note that UGrid is trained at the large scale); and (2) If we have other types of PDEs, we must have UGrid retrained. Please refer to R5 (below).

Q4. Efficiency gain is modest. Added model complexity not worth 2-3x speedup (over Hsieh et al.). How does training time factor in?

R4. (1) Hsieh et al. failed to converge to most of our testcases and its time is the time taken before it reaches max number of iterations. When we compare UGrid with another solver, it must CONVERGE on the testcases. "2-3x speedup" does not reflect UGrid's efficiency gain. The overall gain should be 10-20x (vs AMGCL) and 5-10x (vs AmgX). (2) Training time. Training is required only once for one type of PDEs. It should be treated as preprocessing time and could be ignored over the long lifespan of the model.

Q5. Experiments on other PDEs (preferably non-self-adjoint equations).

R5. We have conducted experiments on two more non-self-adjoint equations: (1) Inhomogeneous Helmholtz equations with spatially-varying coefficients, i.e., $\nabla^2 u + k^2 u = f$, where $k^2$ is a non-zero spatially-varying coefficient, and $f$ is a non-zero right-hand side; (2) Inhomogeneous diffusion-convection-reaction equations, i.e., $\mathbf{v} \cdot \mathbf{\nabla} u - \alpha \nabla^2 u + \beta u = f$, where $\mathbf{v}$ is a varying vector velocity field, $\alpha$, $\beta$ are constants, and $f$ is the non-zero right-hand side.

Results for inhomogeneous Helmholtz equations (with spatially-varying coefficients) are as follows:

UGrid exhibits a 7-20x speedup vs AMGCL and 5-10x speedup vs AmgX on large-scale problems, and generalization power to unseen small-scale problem sizes, with comparable efficiency w.r.t. AMGCL/AmgX.

Results for inhomogeneous steady-state diffusion-convection-reaction equations are as follows:

UGrid exhibits a 2-12x speedup vs AMGCL and on average 2-6x speedup vs AmgX on large-scale problems, and generalization power to unseen small-scale problem sizes, with comparable efficiency w.r.t. AMGCL/AmgX.

These results demonstrate UGrid’s generalization power to elliptic PDEs, and in principle, as long as the Jacobi iterator converges, any linear PDE with a convolution-compatible differential stencil should be readily solvable by UGrid. For more details, please refer to our rebuttal-revision pdf (pp.16~).