One-shot acceleration of transient PDE solvers via online-learned preconditioners

TL;DR

This work introduces PCGBandit, an online-adversarial bandit framework that configures SPD solvers on the fly to accelerate transient PDE simulations within OpenFOAM. By treating preconditioner selection as a bandit problem and using variance-reduced Tsallis-INF updates with wallclock feedback, the method achieves substantial one-shot wallclock speedups without pretraining. Evaluated on six OpenFOAM/MHD problems, PCGBandit often matches or surpasses fixed baselines, with break-even typically around a thousand linear solves. The study also explores deterministic cost estimates for reproducibility and analyzes the impact of the configuration space size on performance, highlighting practical guidelines for deploying online-learned solver optimization in scientific computing.

Abstract

Data-driven acceleration of scientific computing workflows has been a high-profile aim of machine learning (ML) for science, with numerical simulation of transient partial differential equations (PDEs) being one of the main applications. The focus thus far has been on methods that require classical simulations to train, which when combined with the data-hungriness and optimization challenges of neural networks has caused difficulties in demonstrating a convincing advantage against strong classical baselines. We consider an alternative paradigm in which the learner uses a classical solver's own data to accelerate it, enabling a one-shot speedup of the simulation. Concretely, since transient PDEs often require solving a sequence of related linear systems, the feedback from repeated calls to a linear solver such as preconditioned conjugate gradient (PCG) can be used by a bandit algorithm to online-learn an adaptive sequence of solver configurations (e.g. preconditioners). The method we develop, PCGBandit, is implemented directly on top of the popular open source software OpenFOAM, which we use to show its effectiveness on a set of fluid and magnetohydrodynamics (MHD) problems.

Figures (13)

• Figure 1: Illustration of a learning-enhanced numerical simulation setup. In transient simulations the wallclock cost is the sum of costs incurred at each in a sequence of discrete timesteps. For many practical cases, the cost at each step is dominated by solving one or more linear systems, with each solve generating data such as the solver settings used and the time needed to converge. By passing this feedback to a learning algorithm that configures the solver, we hope to reduce the cost of solving the linear systems by more than the cost of the extra computations performed by the learner. If so the total wallclock cost will be reduced in a one-shot manner, i.e. where learning occurs solely within the simulation itself.
• Figure 2: Plots of the relative variation $\|\mathbf{A}_t-\mathbf{A}_1\|_F/\|\mathbf{A}_1\|_F$ of the linear system matrices $\mathbf{A}_t$ from the initial matrix $\mathbf{A}_1$ across simulation timesteps $t\in[T]$. Note that this just one indicator of change, as $A_t$ can vary significantly without significant changes in $\|\mathbf{A}_t-\mathbf{A}_1\|_F$. Indeed it does not fully capture the change in difficulty of solving the systems, as Figure \ref{['fig:costs']} demonstrates that the costs of the Shercliff flow and fringing B-field momentum equations do change over time despite their matrices not changing significantly in Frobenius norm. By comparing with Figure \ref{['fig:plots']}, we see that PCGBandit can improve simulation performance regardless of the magnitude of variation in $\|\mathbf{A}_t-\mathbf{A}_1\|_F$.
• Figure 3: PCG solve costs using two static baseline preconditioners, IC(0) and multigrid, plotted across the timesteps of the six simulations we consider. These plots track the mean, minimum, and maximum costs across ten simulations, demonstrating that the trajectories are largely non-stationary, undermining the main assumption of stochastic bandit methods.
• Figure 4: Performance profiles dolan2002profiles comparing bandit algorithms to two standard baselines: incomplete Cholesky and geometric-algebraic multigrid with the defaults described in Section \ref{['sec:eval']}. The curves report the fraction of the six evaluation tasks from Table \ref{['tab:simulations']} for which the corresponding method is $\tau$-suboptimal, i.e. has $\tau\times$ greater wallclock cost, relative to the best method of the three on each task. The PCGBandit (Tsallis-INF (RV)) curve being in the upper left implies it is optimal or near-optimal on all evaluated tasks.
• Figure 5: Separate wallclock comparisons of PCGBandit and baseline methods on six evaluation tasks. The bold component of each bar represents the time used during the simulations to solve SPD linear systems, while the remaining light component comprises all other costs.