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

Randomized linear algebra (RLA) algorithms are essential for scaling scientific computing and machine learning, yet 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. 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 domain-based inductive bias; (2) Monte Carlo Graph Search (MCGS), 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.