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

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.

Many scientific computing and machine learning systems rely on specialized algorithms that use randomness to efficiently solve large mathematical problems — for example, finding solutions to systems with millions of equations. Designing these algorithms has traditionally required deep expertise and years of manual effort, as researchers must creatively combine mathematical building blocks in the right order. We built a system called RL4RLA that uses reinforcement learning — the same family of AI techniques behind game-playing agents like AlphaGo — to automatically discover these algorithms. Instead of producing opaque "black box" solutions, our system constructs step-by-step symbolic programs from basic mathematical operations, so humans can read, verify, and implement what it finds. Two key ideas make this possible. First, we designed a training curriculum that gradually increases problem difficulty, much like how a student progresses from arithmetic to calculus. Each stage introduces one new challenge that forces the system to learn a specific new technique. Second, we use a graph-based search method that recognizes when different sequences of steps produce the same partial algorithm, avoiding redundant exploration. Our system successfully rediscovered several well-known algorithms that took human experts years to develop, and generalizes to new problem types with minimal modification. This approach could accelerate the design of efficient numerical methods across scientific computing.