Learning to Optimize Multigrid PDE Solvers

TL;DR

This paper introduces a neural network-based framework to learn optimal prolongation operators for multigrid PDE solvers, improving convergence rates across a class of 2D diffusion problems.

Contribution

It presents a novel, unsupervised learning approach to automatically construct multigrid prolongation matrices for parameterized PDEs.

Findings

Enhanced convergence rates over traditional Black-Box multigrid

Successful learning of rules for prolongation matrix construction

Applicable to a broad class of 2D diffusion problems

Abstract

Constructing fast numerical solvers for partial differential equations (PDEs) is crucial for many scientific disciplines. A leading technique for solving large-scale PDEs is using multigrid methods. At the core of a multigrid solver is the prolongation matrix, which relates between different scales of the problem. This matrix is strongly problem-dependent, and its optimal construction is critical to the efficiency of the solver. In practice, however, devising multigrid algorithms for new problems often poses formidable challenges. In this paper we propose a framework for learning multigrid solvers. Our method learns a (single) mapping from a family of parameterized PDEs to prolongation operators. We train a neural network once for the entire class of PDEs, using an efficient and unsupervised loss function. Experiments on a broad class of 2D diffusion problems demonstrate improved convergence rates compared to the widely used Black-Box multigrid scheme, suggesting that our method successfully learned rules for constructing prolongation matrices.