Accelerating Matrix Diagonalization through Decision Transformers with Epsilon-Greedy Optimization

TL;DR

This work reimagines matrix diagonalization as a sequential decision problem by formulating it as an MDP and solving it with a Decision Transformer augmented by an epsilon-greedy exploration strategy. It introduces a Jacobi-rotation–based environment, constructs a trajectory dataset from Hamiltonian symmetric matrices, and trains a causal transformer to predict pivot actions and returns-to-go. The approach achieves a 49.23% reduction in required diagonalization steps and can transfer to smaller matrix sizes without retraining, offering significant speedups (up to ~45x) over the FastEigen baseline while maintaining robust performance. These results suggest that transformer-based sequence models can serve as effective, data-driven eigen-solvers with practical implications for high-performance computing and scientific workflows.

Abstract

This paper introduces a novel framework for matrix diagonalization, recasting it as a sequential decision-making problem and applying the power of Decision Transformers (DTs). Our approach determines optimal pivot selection during diagonalization with the Jacobi algorithm, leading to significant speedups compared to the traditional max-element Jacobi method. To bolster robustness, we integrate an epsilon-greedy strategy, enabling success in scenarios where deterministic approaches fail. This work demonstrates the effectiveness of DTs in complex computational tasks and highlights the potential of reimagining mathematical operations through a machine learning lens. Furthermore, we establish the generalizability of our method by using transfer learning to diagonalize matrices of smaller sizes than those trained.

Figures (5)

• Figure 1: Training architecture of decision transformer.
• Figure 2: Overall training and inference pipeline. Blue outline represents the training pipeline and red outline represents the evaluation pipeline.
• Figure 3: Evaluation Results. Larger episode length implies it takes longer to diagonalize and vice-versa. Note: The results from FastEigen are derived from the model trained on 500K temperature matrices.
• Figure 4: Probability of pivot points chosen by FastEigen and DT method for diagonalization. Each matrix represents a step in the diagonalization process, with the numbers within each matrix element indicating the probability of selecting that specific element as the pivot for that step.
• Figure 5: a) Decision Transformer Action transistion probability. The figure is organized from top to bottom, with each level corresponding to the $i^{th}$ step in the sequence. The numbers enclosed in red circles indicate the potential pivot locations available for use in the $i^{th}$ step. For a detailed explanation of the numbering scheme, please refer to action definition under section \ref{['J-game']}. Additionally, the numbers displayed on the edges represent the probabilities associated with selecting the pivot indicated by the respective edge in the subsequent step. b) Transfer Learning performance on a 3x3 matrix test set for a DT trained on 5x5 matrices