Technical Summary: Learning to Discover Iterative Spectral Algorithms (AutoSpec)

Problem Statement

Numerical Linear Algebra (NLA) and numerical optimization rely heavily on iterative spectral methods to solve large-scale problems, such as linear systems, eigenvalue computations, and matrix function approximations. Traditional approaches often use fixed algorithmic structures (e.g., Chebyshev polynomials, Neumann series) or discrete search spaces for algorithm discovery (e.g., symbolic regression, reinforcement learning over programs). These methods face limitations:

1. Lack of Adaptability: Fixed methods do not adapt well to the specific spectral profiles of input operators.
2. Discrete Search Mismatch: Many ML-driven discovery frameworks search over discrete hypothesis classes (programs), which does not transfer cleanly to NLA where effective methods are continuous objects requiring adaptation to operator-specific structures.
3. Spectral Access: Existing neural preconditioners often train separate models for each matrix or require full spectral knowledge, limiting scalability and transferability.

The paper proposes AutoSpec, a neural network framework designed to discover continuous, iterative spectral algorithms that adapt to input operators using coarse spectral information.

Methodology: The AutoSpec Framework

AutoSpec is a self-supervised, end-to-end differentiable framework that learns mappings from spectral probes to recurrence coefficients defining an operator polynomial.

1. Core Architecture

The framework consists of a Neural Network Engine paired with an Executable Recurrence.

• Input (Spectral Probe): Instead of full eigendecompositions, the model accepts a "spectral probe" $(\hat{\lambda}, \hat{r})$, consisting of a small set of coarse eigenvalue estimates and associated residual norms. These are obtained via a cheap, short warm-start eigensolver run (e.g., 20–50 Lanczos or Arnoldi steps).
• Neural Engine: A neural network maps the spectral probe to a set of scalar coefficients $\{\rho_k, \gamma_k, \eta_k, \alpha_k, \beta_k\}$.
• Executable Recurrence: These coefficients parameterize a state transition operator $M_k(X)$ in a short recurrence relation:
  $$V_{k+1} = M_k(X) V_k$$
  Starting from $V_0 = z$, the process iterates $d$ times to produce $V_{d+1} = P(X)z$, where $P(\cdot)$ is the learned matrix polynomial. This is implemented in a matrix-free manner, requiring only matrix-vector multiplications (matvecs).

2. Hypothesis Class and Parameterization

The framework focuses on linear transitions where the state $V_k$ is composed of blocks (e.g., $A_k, B_k, C_k$). The authors identify two primary regimes subsumed by their parameterization:

• Affine Three-term Recurrence: Covers classical orthogonal polynomial recurrences (e.g., Chebyshev) and matrix-free filters.
• Basis Generation + Learned Expansion: Separates the generation of a stable polynomial basis (via homogeneous recurrence) from the learned expansion weights, allowing optimization beyond simple function fitting.

3. Training Strategy

• Self-Supervised Learning: The model is trained end-to-end by minimizing task-specific NLA objectives (e.g., residual decay, condition number reduction) without imitating existing algorithms.
• Synthetic-to-Real Transfer: Training occurs on small, synthetic diagonal matrices designed to exhibit diverse spectral properties (e.g., heavy-tailed decay, varying gaps). The model learns a size-independent mapping, enabling extrapolation to real-world matrices with dimensions up to $O(10^6)$.
• Efficiency: Training is restricted to diagonal matrices to avoid full matrix operations.
• Stability: To prevent numerical overflow during high-degree iterations, the state vectors are normalized by their norm after each update.
• Loss Function: A unified loss combines:
  1. Task Objective ($L_{obj}$): Inverse of the log-ratio of convergence rates between the learned model and a baseline.
  2. Structural Objective ($L_{struct}$): Constraints on the preconditioner's condition number to ensure robustness.
  3. Regularization ($L_{reg}$): Penalizes outliers to promote generalizability.

Key Contributions

1. General Framework for Continuous Discovery: AutoSpec introduces a neural framework for discovering iterative spectral algorithms that learn continuous updates suited for NLA, contrasting with discrete program search methods.
2. Discovery of Nontrivial Structures: The framework successfully discovers complex recurrence structures using a simple neural architecture and no supervision. The authors identify that these learned recurrences often exhibit behavior consistent with the Chebyshev equioscillation theorem, demonstrating a connection to classical approximation theory.
3. Task-Specific Optimization: Unlike methods that approximate a prescribed function, AutoSpec optimizes directly for downstream task objectives (e.g., residual reduction, eigengap maximization, preconditioner conditioning).
4. Performance Gains: Applied to symmetric positive-definite (spd) matrices, the learned algorithms achieve up to order-of-magnitude improvements in accuracy and/or reductions in iteration counts compared to spectrum-agnostic baselines (e.g., Neumann series, standard Chebyshev).

Empirical Results

The authors evaluated AutoSpec on three representative tasks using real-world sparse matrices from the SuiteSparse collection:

1. Accelerating Linear Solvers:

   • Task: Preconditioning Conjugate Gradient (CG) solvers.
   • Result: Learned degree-11 polynomial preconditioners significantly outperformed spectrum-agnostic Neumann preconditioners, reducing iteration counts by orders of magnitude across various matrices (e.g., G2_circuit, parabolic_fem).

2. Eigenvalue Computations:

   • Task: Computing the $k$ smallest-magnitude eigenpairs using the eigs algorithm.
   • Result: A degree-21 polynomial preconditioner enhanced spectral separation near the target boundary. AutoSpec consistently outperformed a power preconditioner baseline ($P(X)=X^d$) across different subspace dimensions and target eigenvalue counts.

3. Matrix Function Approximation:

   • Task: Approximating $X^{-1/2}$ for covariance whitening of a DNA similarity Gram matrix.
   • Result: The learned degree-21 polynomial achieved order-of-magnitude lower operator-norm error compared to a truncated Neumann-series baseline.

Significance and Claims

The paper claims that AutoSpec demonstrates the viability of using machine learning to automate the design of robust, deployable numerical algorithms.

• Connection to Theory: The learned algorithms are not "black boxes"; they recover classical behaviors (like Chebyshev equioscillation) while offering the flexibility to optimize for specific structural constraints (e.g., avoiding small eigenvalue clusters that hinder solver convergence).
• Scalability: By relying only on coarse spectral probes and training on small synthetic instances, the framework bridges the gap between small-scale learning and large-scale real-world application.
• Beyond Classical Limits: The framework can discover algorithms that outperform classical minimax constructions by trading off slight improvements in operator conditioning for better preconditioner conditioning, a flexibility not easily achieved with fixed polynomial families.

Limitations Acknowledged:
The current evaluation is limited to symmetric positive-definite (spd) matrices. The authors note that extending the framework to indefinite and unsymmetric problems is a necessary future direction. Additionally, performance depends on the training set containing spectral features representative of the test-time operators.