Learning to Relax: Setting Solver Parameters Across a Sequence of Linear System Instances

Mikhail Khodak, Edmond Chow, Maria Florina Balcan, Ameet Talwalkar

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: scientific computing, data-driven algorithm design, online learning, multi-armed bandits, contextual bandits, numerical analysis, learning-augmented algorithms, algorithms with predictions

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We show provable guarantees for learning the relaxation parameter of linear system solvers.

Abstract:

Solving a linear system ${\bf Ax}={\bf b}$ is a fundamental scientific computing primitive for which numerous solvers and preconditioners have been developed.
These come with parameters whose optimal values depend on the system being solved and are often impossible or too expensive to identify;
thus in practice sub-optimal heuristics are used.
We consider the common setting in which many related linear systems need to be solved, e.g. during a single numerical simulation.
In this scenario, can we sequentially choose parameters that attain a near-optimal overall number of iterations, without extra matrix computations?
We answer in the affirmative for Successive Over-Relaxation (SOR), a standard solver whose parameter $\omega$ has a strong impact on its runtime.
For this method, we prove that a bandit online learning algorithm—using only the number of iterations as feedback—can select parameters for a sequence of instances such that the overall cost approaches that of the best fixed $\omega$ as the sequence length increases.
Furthermore, when given additional structural information, we show that a _contextual_ bandit method asymptotically achieves the performance of the _instance-optimal_ policy, which selects the best $\omega$ for each instance.
Our work provides the first learning-theoretic treatment of high-precision linear system solvers and the first end-to-end guarantees for data-driven scientific computing, demonstrating theoretically the potential to speed up numerical methods using well-understood learning algorithms.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Primary Area: learning theory

Submission Number: 8824