MatRL: Provably Generalizable Iterative Algorithm Discovery via Monte-Carlo Tree Search

TL;DR

MatRL is a reinforcement learning framework that automatically discovers and plans iterative algorithms for matrix functions, ensuring provable generalization and improved performance over existing methods.

Contribution

The paper introduces MatRL, a novel RL-based approach that automates the design of matrix iteration algorithms with theoretical guarantees of generalization.

Findings

MatRL outperforms baseline algorithms in numerical experiments.

Learned algorithms generalize to larger matrices from the same distribution.

The approach effectively combines Monte-Carlo tree search with reinforcement learning.

Abstract

Iterative methods for computing matrix functions have been extensively studied and their convergence speed can be significantly improved with the right tuning of parameters and by mixing different iteration types. Handtuning the design options for optimal performance can be cumbersome, especially in modern computing environments: numerous different classical iterations and their variants exist, each with non-trivial per-step cost and tuning parameters. To this end, we propose MatRL -- a reinforcement learning based framework that automatically discovers iterative algorithms for computing matrix functions. The key idea is to treat algorithm design as a sequential decision-making process. Monte-Carlo tree search is then used to plan a hybrid sequence of matrix iterations and step sizes, tailored to a specific input matrix distribution and computing environment. Moreover, we also show that the learned algorithms provably generalize to sufficiently large matrices drawn from the same distribution. Finally, we corroborate our theoretical results with numerical experiments demonstrating that MatRL produces algorithms that outperform various baselines in the literature.