Parallel Preconditioning via Reinforcement Learning

The spectral deferred correction method is an iterative solver for time-dependent partial differential equations. It can be interpreted as a preconditioned Picard iteration for the collocation problem. The key to an efficient method is the choice of the preconditioner: it defines the speed of convergence and the level of parallelism. While the de-facto standard is a fast-converging, serial preconditioner, our goal is to find a fast-converging, diagonal and therefore parallel one. To achieve that, we employ reinforcement learning enabled by differentiable programming, using differentiable loss based on spectral radius value obtained from suggested preconditioner initialization. We look at Dahlquist’s equation and train a network to select favorable, diagonal preconditioners depending on the parameter of the equation. The trained network can then be used directly to suggest parallel preconditioners for more complex problems: Using spectral methods, we can apply this to partial differential equations with a linear stiff part that we treat decoupled and implicitly with tailored preconditioners chosen by the network, whereas nonlinear parts are treated explicitly in real space. In this talk we present the overall approach, the training technique and show first results for linear and nonlinear PDEs.