Precision autotuning for linear solvers via contextual bandit-based RL

TL;DR

This paper introduces a reinforcement learning framework using a contextual bandit approach to adaptively tune precision in linear solvers, improving efficiency while maintaining accuracy.

Contribution

It presents the first RL-based precision autotuning method for linear solvers, demonstrating effectiveness and generalization to unseen data.

Findings

Reduces computational cost compared to double-precision baselines.

Maintains accuracy and convergence in iterative refinement.

Generalizes well to out-of-sample data.

Abstract

We propose a reinforcement learning (RL) framework for adaptive precision tuning for linear solvers, which can be extended to general algorithms. The framework is formulated as a contextual bandit problem and solved using incremental action-value estimation with a discretized state space to select optimal precision configurations for computational steps, balancing precision and computational efficiency. To verify its effectiveness, we apply the framework to iterative refinement for solving linear systems $Ax = b$. In this application, our approach dynamically chooses precisions based on calculated features from the system while maintaining acceptable accuracy and convergence. In detail, an action-value estimator takes discretized features (e.g., approximate condition number and matrix norm) as input and outputs estimated action values, from which a policy selects the actions (chosen precision configurations for specific steps), optimized via an $\epsilon$-greedy strategy to maximize a multi-objective reward to balance accuracy and computational cost. Empirical results demonstrate effective precision selection, reducing computational cost while maintaining accuracy comparable to double-precision baselines. The framework generalizes to diverse out-of-sample data and provides insights into applying RL precision selection to other numerical algorithms, advancing mixed-precision numerical methods in scientific computing. To the best of our knowledge, this is the first work on precision autotuning with RL with verification on unseen datasets.