Paper: A Neural-preconditioned Poisson Solver for Mixed Dirichlet and Neumann Boundary Conditions (ICLR-2024)

Reviewer comment (Theory):
Without a positive definite preconditioner, it is likely that the simulations in Figure 2 become unstable (or unphysical) if run for long enough. I think the authors should carefully consider both the computational cost per time step and the numerical stability of the time-stepping scheme when using various preconditioners. This direction is mentioned as a future work, but I think the paper's contribution is compromised without ensuring symmetric positive definiteness.

Author response:
We appreciate the acknowledgement of the surprising effectiveness of our preconditioner. We reiterate our point from the main post that symmetry and positive definiteness, while likely to bring modest benefits, are not critical for our solver's efficacy--or the simulation's stability--in practice.

Reviewer comment (Theory):
The lack of theoretical convergence analysis of the preconditioner and iterative solver means that the solver is of limited general use. Since the preconditioner is not positive definite, I suspect that the theoretical convergence of the iterative scheme is extremely difficult to understand.

Author response:
Regarding the lack of theoretical convergence analysis and the high resource requirements for training, we note that these same limitations apply to the seminal work [Tompson et al. 2017] that has nonetheless made a large impact due to its practical benefits. Especially for the applications we target in real-time simulation and visual effects, the single upfront cost of training the network can pay great dividends over the lifetime of the solver's use.

Reviewer comment (Experiments):
Training the neural network with a large memory requirement takes a very long time before it can be successfully used in time simulations.

Author response:
Regarding the lack of theoretical convergence analysis and the high resource requirements for training, we note that these same limitations apply to the seminal work [Tompson et al. 2017] that has nonetheless made a large impact due to its practical benefits. Especially for the applications we target in real-time simulation and visual effects, the single upfront cost of training the network can pay great dividends over the lifetime of the solver's use.

Reviewer comment (Experiments):
The network does not yet leverage sparsity, so may not scale gracefully to extreme sparse problems with many empty grid cells.

Author response:
Our systems feature sparsity of two types: sparsity of the finite difference Laplacian matrix, and sparsity of the fluid domain occupancy. Matrix sparsity is already exploited by the linear algebra operations in our NPSDO solver (and by our use of small convolution kernels in our network), while domain sparsity is what we mentioned as future work. We note that in typical fluid simulations, the domain geometry is constantly evolving and can range from highly sparse to fully dense; we anticipate inventing a single architecture that works well in both occupancy regimes to be very challenging. Furthermore, efficient algorithms and implementations for sparse convolutions--especially with variable sparsity patterns--are still an active area of research with limited support in available libraries. Even support for sparse matrix formats like CSR and CSC is still in beta stage in PyTorch.

Reviewer comment (Experiments):
Only isotropic problems were considered - extending the approach to anisotropic or nonlinear problems is an open question.

Author response:
We agree that all of the extensions raised here are interesting avenues for future work, and we are especially excited to see how our spatially-varying kernel architecture performs for problems with heterogeneous and anisotropic constitutive laws. However, we do not see any of these investigations as essential to the completeness and impact of our paper.

Reviewer comment (Experiments):
Quantitative ablation of design choices like coarsening levels, memory usage, etc. could provide further insights.

Author response:
We will include performance statistics for networks with a range of coarsening levels $\mathcal{L}$ trained using the strategy described in our response to the second review.

Reviewer comment (Experiments):
While extensive, the benchmark set considers a subset of potential problem domains - generalizing to new classes would reinforce claims.

Author response:
We agree that all of the extensions raised here are interesting avenues for future work, and we are especially excited to see how our spatially-varying kernel architecture performs for problems with heterogeneous and anisotropic constitutive laws. However, we do not see any of these investigations as essential to the completeness and impact of our paper.

Reviewer comment (Experiments):
Can you exploit sparsity directly in the network/algorithm to better handle extremely sparse problems that arise in practice.

Author response:
Our systems feature sparsity of two types: sparsity of the finite difference Laplacian matrix, and sparsity of the fluid domain occupancy. Matrix sparsity is already exploited by the linear algebra operations in our NPSDO solver (and by our use of small convolution kernels in our network), while domain sparsity is what we mentioned as future work. We note that in typical fluid simulations, the domain geometry is constantly evolving and can range from highly sparse to fully dense; we anticipate inventing a single architecture that works well in both occupancy regimes to be very challenging. Furthermore, efficient algorithms and implementations for sparse convolutions--especially with variable sparsity patterns--are still an active area of research with limited support in available libraries. Even support for sparse matrix formats like CSR and CSC is still in beta stage in PyTorch.

Reviewer comment (Experiments):
Can you explore non-Euclidean/anisotropic problem settings to broaden the method's scope?

Author response:
We agree that all of the extensions raised here are interesting avenues for future work, and we are especially excited to see how our spatially-varying kernel architecture performs for problems with heterogeneous and anisotropic constitutive laws. However, we do not see any of these investigations as essential to the completeness and impact of our paper.

Reviewer comment (Experiments):
Can you develop cheaper training data generation techniques or consider self-supervised learning alternatives?

Author response:
Our training arguably is already self-supervised, since our loss function does not require a ground-truth solution for each linear system. As mentioned in the paper, we did pursue several less costly approaches for generating the training data that unfortunately did not work well enough. Our existing approach still could be accelerated by using a higher-performance implementation for generating the Ritz vectors, but we do not consider the current implementation prohibitively expensive, especially when compared to training time.

Reviewer comment (Experiments):
Can you broaden empirical evaluation to new problem classes to further validate generalization abilities?

Author response:
We agree that all of the extensions raised here are interesting avenues for future work, and we are especially excited to see how our spatially-varying kernel architecture performs for problems with heterogeneous and anisotropic constitutive laws. However, we do not see any of these investigations as essential to the completeness and impact of our paper.

Reviewer comment (Novelty):
The primary contributions of the paper, including the understanding of convolution as a smoother, pooling as the restriction, and the connection between multigrid and convolutional neural networks, have been extensively explored in the following reference:

  Juncai He and Jinchao Xu. "Mg-Net: A unified framework of multigrid and convolutional neural network." arXiv:1901.10415, 2019.

Author response:
## Primary Contribution
We consider our primary contributions to be our novel network architecture and our demonstrations that it can achieve state-of-the-art performance on problems with mixed boundary conditions. We agree that past work has understood the connections between components of traditional CNN architectures and their counterparts in a multigrid hierarchy.

Reviewer comment (Experiments):
Could the authors provide an explanation of what the neural network truly learns for an iterative solver to maximize its performance? Specifically, does it focus on learning the smoother or the prolongation operator, or both? Clarity on this aspect would be valuable for understanding the role of the neural network in enhancing solver performance.

Author response:
## What Does the Network Learn?
The network learns by optimizing the parameters defining (i) our spatially varying convolution kernel (i.e., the smoother) and (ii) the fixed convolution kernels that compute coefficients for mixing in the interpolated corrections. It does not learn restriction or prolongation operators, which are implemented by fixed pooling and upsampling blocks.

Reviewer comment (Experiments):
Is it necessary to fix the levels of coarsening in the proposed approach?

Author response:
## Number of Coarsening Levels
An instance of our network architecture is defined by its number of coarsening levels. This is a hyperparameter that is fixed to 4 in all examples. However, the weights of a network with $\mathcal{L}$ levels can be copied to a new network with $\mathcal{L + 1}$ levels to initialize its training, and we found this strategy beneficial for training deeper hierarchies; we will discuss this approach in the revised draft.

Reviewer comment (Theory):
The biggest weakness is perhaps that the "CNN" used is linear, then it is nothing but a single parametrized convolution, as stacking convolutions without nonlinearity is still a convolution just with a different kernel sizes.

Author response:
## Equivalence to One Large Convolution
Indeed, because the network output is a linear function of the input vector, it can be interpreted as a single convolution with one large spatially varying kernel (whose size depends on the number of layers) that is parametrized by our network weights. We do note that the network output is a nonlinear polynomial function of the input *image* since the coefficients of successive kernels end up multiplying each other. We see the simplicity and linearity of our network not as a weakness but as a contribution (showing a heavier-weight, nonlinear architecture is unnecessary) and a practical advantage (accelerating inference at runtime).

Reviewer comment (Theory):
Page 4: "The PSDO algorithm can be understood as a modification of standard CG that replaces the residual with the preconditioned residual as the starting point for generating search directions and, consequently, cannot enjoy many of the simplifications baked into the traditional algorithm". This sentence is somewhat too long in reviewer's humble opinion, and up to the authors' judgement, I suggest rewrite it a little. Meanwhile, please consider spending a few words to explain the reason why PSDO cannot achieve CG's automatic $A$-conjugate directions. It says that PSDO does "an exact line search", so it does not a subspace search like the CG does?

Author response:
## A-conjugacy of Directions
Despite using an exact line search as in PCG, enforcing conjugacy with just the previous search direction does not automatically ensure conjugacy with all past directions. This property enjoyed by CG hinges on symmetry, and we would be happy to provide detailed proofs in the appendix.

Reviewer comment (Experiments):
Page 5-6, the default behavior of upsampling using the interpolation in `nn.functional` is `nearest` which gives a piecewise constant function that is unfavorable in this context, please test `bilinear` which is MG prolongation.

Author response:
## Bi/trilinear Prolongation
Thank you for pointing this out! Yes, trilinear interpolation is a closer analogue for MG prolongation. Simply substituting it for the default upsampling in our old trained network achieves a measurable speedup (reducing the iteration count across all linear systems in our test set by an average of 5.3%). We have updated the statistics reported in our paper to reflect these improvements. We have also begun training a new model incorporating this trilinear interpolation, which we expect to achieve further improvements.

Reviewer comment (Experiments):
The single-precision for NN and double for the CG is a good practice. However, it would be interesting to see whether this combination can reach relative residual `1e-8` (cf. those in Figure 5).

Author response:
## Tolerance Achievable
The accuracy of our solver is not limited by the preconditioner's single-precision arithmetic and can easily surpass a relative residual of $10^{-8}$. To test this, we disabled the convergence test and executed each solve in our entire test set for a fixed number of iterations (100). We measured a median relative residual of $4.01\times10^{-15}$ and report a histogram in the newly added Figure 11 of the appendix.