RL4RLA: Teaching ML to Discover Randomized Linear Algebra Algorithms Through Curriculum Design and Graph-Based Search

2026-05-18

AI Summary

This work addresses the reliance of existing randomized linear algebra (RLA) algorithms on manual expert design and the absence of automated discovery mechanisms. The authors propose a framework that integrates reinforcement learning, curriculum learning, and Monte Carlo tree search to automatically synthesize interpretable and efficient RLA algorithms from basic linear algebra primitives via symbolic representations. The approach incorporates an RLA-oriented numerical curriculum to inject domain-specific inductive bias and employs graph-based search to merge equivalent sub-algorithms, substantially improving exploration efficiency. Empirically, the method successfully reproduces state-of-the-art algorithms—including Sketch-and-Precondition, randomized Kaczmarz, and Newton Sketch—and enables customizable trade-offs among accuracy, speed, and numerical stability.

Abstract

Randomized linear algebra (RLA) algorithms are a modern class of numerical linear algebra techniques that play an essential role in scientific computing and machine learning, with broad and growing adoption. However, their discovery remains mostly a manual process that requires deep expert knowledge and inspiration. While Reinforcement Learning (RL) offers a pathway to automation, standard approaches struggle with sparse reward landscapes and vast search spaces inherent to high-performing RLA algorithms. In this paper, we present RL4RLA, a general RL framework that automates the discovery of interpretable, symbolic RLA algorithms. Unlike black-box approaches, our method builds explicit algorithms from basic linear algebra primitives, ensuring verifiable and implementable representations. To enable efficient discovery, we introduce: (1) a numerical curriculum that progressively increments problem difficulty to encode inductive bias specific to the RLA domain; (2) Monte Carlo Graph Search, which optimizes exploration by identifying and merging equivalent partial algorithms. We demonstrate that RL4RLA rediscovers state-of-the-art methods, including sketch-and-precondition solvers, Randomized Kaczmarz, and Newton Sketch, and can be targeted to produce algorithms optimized for specific trade-offs between accuracy, speed, and stability. Code is available at https://github.com/Tim-Xiong/RL4RLA.

Problem

• Randomized Linear Algebra
• Algorithm Discovery
• Reinforcement Learning
• Symbolic Algorithms
• Scientific Computing

Innovation

• Reinforcement Learning
• Randomized Linear Algebra
• Curriculum Learning
• Graph-Based Search
• Symbolic Algorithm Discovery