# Optimized Sparse Matrix Operations for Reverse Mode Automatic Differentiation

arXiv: https://arxiv.org/abs/2212.05159 · HTML: https://mlsanity.com/paper/2212.05159/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 61 references that are in the map.
- First posted on arXiv 10 Dec 2022.

## Builds on (most cited first, 15 of 15)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Spectral Networks and Locally Connected Networks on Graphs](https://mlsanity.com/paper/1312.6203/) (2013)
- [Revisiting Semi-Supervised Learning with Graph Embeddings](https://mlsanity.com/paper/1603.08861/) (2016)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [Learning to Optimize Multigrid PDE Solvers](https://mlsanity.com/paper/1902.10248/) (2019)
- [Learning Algebraic Multigrid Using Graph Neural Networks](https://mlsanity.com/paper/2003.05744/) (2020)
- [Learning Interface Conditions in Domain Decomposition Solvers](https://mlsanity.com/paper/2205.09833/) (2022)
- [Tuning Multigrid Methods with Robust Optimization](https://mlsanity.com/paper/2001.00887/) (2020)
- [GLU3.0: Fast GPU-based Parallel Sparse LU Factorization for Circuit Simulation](https://mlsanity.com/paper/1908.00204/) (2019)
- [GSoFa: Scalable Sparse Symbolic LU Factorization on GPUs](https://mlsanity.com/paper/2007.00840/) (2020)

## Built on by (most cited first, 5 of 5)

- [Neural incomplete factorization: learning preconditioners for the conjugate gradient method](https://mlsanity.com/paper/2305.16368/) (2023)
- [MG-GNN: Multigrid Graph Neural Networks for Learning Multilevel Domain Decomposition Methods](https://mlsanity.com/paper/2301.11378/) (2023)
- [Learning incomplete factorization preconditioners for GMRES](https://mlsanity.com/paper/2409.08262/) (2024)
- [Deep Learning-Enhanced Preconditioning for Efficient Conjugate Gradient Solvers in Large-Scale PDE Systems](https://mlsanity.com/paper/2412.07127/) (2024)
- [nablaSD: Differentiable Programming for Sparse Tensors](https://mlsanity.com/paper/2303.07030/) (2023)

## Abstract (arXiv)

Sparse matrix representations are ubiquitous in computational science and machine learning, leading to significant reductions in compute time, in comparison to dense representation, for problems that have local connectivity. The adoption of sparse representation in leading ML frameworks such as PyTorch is incomplete, however, with support for both automatic differentiation and GPU acceleration missing. In this work, we present an implementation of a CSR-based sparse matrix wrapper for PyTorch with CUDA acceleration for basic matrix operations, as well as automatic differentiability. We also present several applications of the resulting sparse kernels to optimization problems, demonstrating ease of implementation and performance measurements versus their dense counterparts.