Towards Faster Matrix Diagonalization with Graph Isomorphism Networks and the AlphaZero Framework

TL;DR

The paper reframes matrix diagonalization via the Jacobi method as reinforcement learning problems, modeling large-matrix diagonalization as a Semi-Markov Decision Process and small-matrix diagonalization as a Markov Decision Process, with state transitions driven by $M^{i+1} = J(p,q,\theta)^{\top} M^{i} J(p,q,\theta)$. It couples Monte Carlo Tree Search with Graph Isomorphic Networks (GIN) to learn rotation heuristics, achieving reduced rotation counts and efficient inference on small matrices, and explores scalability to larger matrices through transfer learning and state abstractions. Empirically, AlphaZero-based strategies produce superior heuristics over traditional MaxElem baselines, reveal structured transition patterns via learned sweep distributions, and demonstrate potential cost savings in rotations, though scaling to very large matrices remains challenged without enhanced abstractions. The work advances practical, scalable approaches for accelerating matrix diagonalization in scientific and engineering computations by integrating RL planning, graph representations, and structured policy learning.

Abstract

In this paper, we introduce innovative approaches for accelerating the Jacobi method for matrix diagonalization, specifically through the formulation of large matrix diagonalization as a Semi-Markov Decision Process and small matrix diagonalization as a Markov Decision Process. Furthermore, we examine the potential of utilizing scalable architecture between different-sized matrices. During a short training period, our method discovered a significant reduction in the number of steps required for diagonalization and exhibited efficient inference capabilities. Importantly, this approach demonstrated possible scalability to large-sized matrices, indicating its potential for wide-ranging applicability. Upon training completion, we obtain action-state probabilities and transition graphs, which depict transitions between different states. These outputs not only provide insights into the diagonalization process but also pave the way for cost savings pertinent to large-scale matrices. The advancements made in this research enhance the efficacy and scalability of matrix diagonalization, pushing for new possibilities for deployment in practical applications in scientific and engineering domains.

Figures (8)

• Figure 1: Primitive pivot orderings for all 8 baselines/options. During the Jacobi cyclic game, the goal is to construct an ordering such that the primitive rotations are reduced. Darker pivots $(i,j)$ are rotated first when the option is selected. During the classic Jacobi rotation game, the agent selects each pivot individually as an action.
• Figure 2: Comparative analysis of FastGIN performance under different conditions.
• Figure 3: Sweep transition probabilities for a 15x15 matrix.
• Figure 6: Comparative analysis of sweep transition probabilities and graph representations of sweep transitions across matrices of sizes 15x15, 30x30, and 50x50.
• Figure 7: Progressive comparison of rotation counts between baseline and AlphaZero strategies across various matrix sizes. The black points denote the baseline approach, which iteratively applies sweeps in one of the 8 specific directions/options until successful diagonalization is achieved. The red dots represent the AlphaZero method, demonstrating that AlphaZero outperforms the baseline more significantly as the dimensions of the matrix increase.