matrixpfn-base

GNN-based learned preconditioner for sparse linear systems.

Version: 20260309-0350

Usage

import numpy as np
from scipy.io import mmread
from matrixpfn import MatrixPFN

pfn = MatrixPFN.from_pretrained("Csed-dev/matrixpfn-base")

A = mmread("matrix.mtx")  # any scipy sparse matrix
b = A @ np.random.randn(A.shape[0])

result = pfn.solve(A, b)  # accepts scipy sparse directly
print(f"Converged: {result.converged} in {result.iterations} iterations")

Architecture

Parameter	Value
Network	ContextResGCN
Layers	24
Embed	128
Hidden	512
Context pairs	10
Parameters	2,056,706
dtype	float32

Training

• epochs: 5000
• best_loss: 0.006055
• loss_function: l1_direct
• batch_size: 256
• matrices_per_epoch: 4
• domains: diffusion, diffusion_advection, graph_laplacian, elasticity, stokes, sbm, spectral_stress, variable_diffusion, variable_advection, enhanced_diffusion, enhanced_advection, random_sparse, directed_power_law
• grid_sizes: [16, 24, 32, 48]

Benchmark

domain	grid	converged	avg_iters	avg_residual
diffusion	16x16	0/20	300.0	3.96e-02
diffusion	24x24	0/20	300.0	3.40e-02
diffusion	32x32	0/20	300.0	2.43e-02
diffusion	48x48	0/20	300.0	3.41e-02
diffusion	64x64	0/20	300.0	4.11e-02
diffusion_advection	16x16	0/20	300.0	3.69e-02
diffusion_advection	24x24	0/20	300.0	2.91e-02
diffusion_advection	32x32	0/20	300.0	2.59e-02
diffusion_advection	48x48	0/20	300.0	2.95e-02