Neural Multigrid Architectures
This work addresses the challenge of creating competitive neural solvers for linear problems in numerical linear algebra, though it is incremental as it builds on existing multigrid methods.
The authors tackled the problem of designing efficient neural network architectures for multigrid solvers, resulting in a method that reduces the spectral radius of the error propagation matrix by 2 to 5 times compared to a basic linear multigrid with Jacobi smoother.
We propose a convenient matrix-free neural architecture for the multigrid method. The architecture is simple enough to be implemented in less than fifty lines of code, yet it encompasses a large number of distinct multigrid solvers. We argue that a fixed neural network without dense layers can not realize an efficient iterative method. Because of that, standard training protocols do not lead to competitive solvers. To overcome this difficulty, we use parameter sharing and serialization of layers. The resulting network can be trained on linear problems with thousands of unknowns and retains its efficiency on problems with millions of unknowns. From the point of view of numerical linear algebra network's training corresponds to finding optimal smoothers for the geometric multigrid method. We demonstrate our approach on a few second-order elliptic equations. For tested linear systems, we obtain from two to five times smaller spectral radius of the error propagation matrix compare to a basic linear multigrid with Jacobi smoother.