AutoSpec: A Neural Network Framework for Discovering Iterative Spectral Algorithms

Installation

To start using AutoSpec training and the auxiliary Python evaluation pipeline, first create conda environment and install packages from requirements.txt:

conda create -n autospec python=3.10
conda activate autospec
pip install -r requirements.txt

Repository Layout

AutoSpec/
├── README.md
├── requirements.txt
├── checkpoints/
└── src/
    ├── evaluation/
    │   ├── README.md
    │   ├── evaluation_matlab/
    │   └── evaluation_python/
    ├── model/
    ├── utils/
    ├── scripts/
    ├── train_linsolve_backbone.py
    ├── train_linsolve_embedding.py
    ├── train_eigs_backbone.py
    └── train_eigs_embedding.py

• checkpoints/: packaged trained checkpoint used by the example evaluation workflow, containing a trained 10-layer AutoSpec model ./checkpoints/linsolve_model_10_layer.pth.
• src/: main Python implementation for training and evaluation.
• src/evaluation/: evaluation implementations grouped under one source-tree location. Contains MATLAB and Python evaluation implementations.
• src/model/: AutoSpec Neural Network Engine definitions and reusable model modules.
• src/scripts/: shell entry points for training and auxiliary Python evaluation.
• src/train_linsolve_backbone.py: training entry point for the linear-solve backbone model.
• src/train_linsolve_embedding.py: training entry point for the embedding layer on top of the backbone.
• src/train_eigs_backbone.py: training entry point for the eigs backbone model.
• src/train_eigs_embedding.py: training entry point for the embedding layer on top of the eigs backbone.

Training Neural Network Engine for Linear Systems

Step 1: Training Backbone Model (Pre-training)

To run the training script for the backbone model, simply run the following command:

# To train a 10-layer backbone
export BACKBONE_LAYERS="10" BACKBONE_LR="0.0005" BACKBONE_EPOCHS="1300" BACKBONE_LOSS="log_ratio_exp_cond_8" \
CUDA_VISIBLE_DEVICES=0 bash ./src/scripts/train_linsolve_backbone.sh

# To train a 15-layer backbone
export BACKBONE_LAYERS="15" BACKBONE_LR="0.0002" BACKBONE_EPOCHS="1500" BACKBONE_LOSS="log_ratio_exp_cond_16" \
CUDA_VISIBLE_DEVICES=0 bash ./src/scripts/train_linsolve_backbone.sh

You can also specify hyperparameters in train_linsolve_backbone.sh, such as lr, 'epochs', loss, input_dim, num_coeffs, etc. The current configurations are set to the hyperparameters for training a good backbone model, which can be further tuned for better performance.

Step 2: Training Embedding Layer (Post-training)

To run the training script for the embedding layer, simply run the following command:

# Train the embedding layer for the 10-layer backbone
export BACKBONE_LAYERS="10" BACKBONE_LR="0.0005" BACKBONE_EPOCHS="1300" BACKBONE_LOSS="log_ratio_exp_cond_8" HEAD_LR="0.001" HEAD_EPOCHS="500" HEAD_LOSS="log_ratio_exp_cond_8" HEAD_INPUT_DIM="20" \
CUDA_VISIBLE_DEVICES=0 bash ./src/scripts/train_linsolve_embedding.sh

# Train the embedding layer for the 15-layer backbone
export BACKBONE_LAYERS="15" BACKBONE_LR="0.0002" BACKBONE_EPOCHS="1500" BACKBONE_LOSS="log_ratio_exp_cond_16" HEAD_LR="0.001" HEAD_EPOCHS="500" HEAD_LOSS="log_ratio_exp_cond_2" HEAD_INPUT_DIM="20" \
CUDA_VISIBLE_DEVICES=0 bash ./src/scripts/train_linsolve_embedding.sh

Same as the backbone training script, you can also specify hyperparameters in train_linsolve_embedding.sh. The current configurations are set to the hyperparameters for training a good embedding layer, which can be further tuned for better performance.

Evaluation

Primary MATLAB Evaluation

The primary hardware-efficient evaluation setup is the MATLAB transfer package:

cd src/evaluation/evaluation_matlab
experimentDirs = main();

The MATLAB package is self-contained and preserves its original structure: experiments/, matlab_helpers/, pytorch_model/, and test_matrices/. It uses MATLAB for the solver experiments and calls a small PyTorch inference script for AutoSpec coefficients.

The Python executable used by MATLAB is detected in this order:

1. AUTOSPEC_PYTHON environment variable
2. paths.pythonExe in src/evaluation/evaluation_matlab/autospec_transfer_paths.m
3. python3 on PATH, if it can import Torch
4. python on PATH, if it can import Torch

For details, see src/evaluation/evaluation_matlab/README.md.

Auxiliary Python Evaluation

The Python evaluation setup is kept as an auxiliary implementation for users who want a pure-Python SuiteSparse workflow. To apply the model from AutoSpec for PCG on real-world systems, run the evaluation script as follows:

First, use the packaged checkpoint or override it with the desired checkpoint:

export AUTOSPEC_INFERENCE_CHECKPOINT_PATH=/path/to/trained/model.pth

An already-trained checkpoint is available at ./checkpoints/linsolve_model_10_layer.pth.

Second, run the evaluation script as follows:

bash ./src/scripts/eval_autospec.sh <matrix_name> <first_level_preconditioner> <autospec_probe_iters> <max_iters> <seed>

where

• <matrix_name> is the name of the matrix in the matrix name from SuiteSparse (e.g., Schmid/thermal2).
• <first_level_preconditioner> is the first level preconditioner used for PCG, e.g., jacobi, amg, or simply identity if no preconditioning is desired.
• <autospec_probe_iters> is the number of iterations for the AutoSpec spectral probe (50 or 200).
• <max_iters> is the maximum number of iterations for the CG solver (e.g., 500).

As an example,

bash ./src/scripts/eval_autospec.sh Schmid/thermal2 jacobi 200 2000 42

Acknowledgement

We'd like to thank the creators of the GNP codebase, which we followed for our python implementation of the evaluation pipeline on SuiteSparse matrix collection.