# Optimal Gradient Checkpointing for Sparse and Recurrent Architectures using Off-Chip Memory

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 19 references that are in the map.
- First posted on arXiv 16 Dec 2024.

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

- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Surrogate Gradient Learning in Spiking Neural Networks](https://mlsanity.com/paper/1901.09948/) (2019)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Converting Static Image Datasets to Spiking Neuromorphic Datasets Using Saccades](https://mlsanity.com/paper/1507.07629/) (2015)
- [SuperSpike: Supervised learning in multi-layer spiking neural networks](https://mlsanity.com/paper/1705.11146/) (2017)
- [Long short-term memory and learning-to-learn in networks of spiking neurons](https://mlsanity.com/paper/1803.09574/) (2018)
- [The Heidelberg spiking datasets for the systematic evaluation of spiking neural networks](https://mlsanity.com/paper/1910.07407/) (2019)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [Recurrent Neural Networks (RNNs): A gentle Introduction and Overview](https://mlsanity.com/paper/1912.05911/) (2019)
- [Dissecting the Graphcore IPU Architecture via Microbenchmarking](https://mlsanity.com/paper/1912.03413/) (2019)
- [Block-Sparse Recurrent Neural Networks](https://mlsanity.com/paper/1711.02782/) (2017)
- [Spiking Neural Networks and Bio-Inspired Supervised Deep Learning: A Survey](https://mlsanity.com/paper/2307.16235/) (2023)
- [FPGA Resource-aware Structured Pruning for Real-Time Neural Networks](https://mlsanity.com/paper/2308.05170/) (2023)
- [Efficient Sparse-Dense Matrix-Matrix Multiplication on GPUs Using the Customized Sparse Storage Format](https://mlsanity.com/paper/2005.14469/) (2020)
- [Harnessing Manycore Processors with Distributed Memory for Accelerated Training of Sparse and Recurrent Models](https://mlsanity.com/paper/2311.04386/) (2023)
- [On Performance Analysis of Graphcore IPUs: Analyzing Squared and Skewed Matrix Multiplication](https://mlsanity.com/paper/2310.00256/) (2023)

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

- [Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts](https://mlsanity.com/paper/2602.02108/) (2026)
- [Bruno: Backpropagation Running Undersampled for Novel device Optimization](https://mlsanity.com/paper/2505.17791/) (2025)
- [DongYuan: An LLM-Based Framework for Integrative Chinese and Western Medicine Spleen-Stomach Disorders Diagnosis](https://mlsanity.com/paper/2603.28191/) (2026)

## Abstract (arXiv)

Recurrent neural networks (RNNs) are valued for their computational efficiency and reduced memory requirements on tasks involving long sequence lengths but require high memory-processor bandwidth to train. Checkpointing techniques can reduce the memory requirements by only storing a subset of intermediate states, the checkpoints, but are still rarely used due to the computational overhead of the additional recomputation phase. This work addresses these challenges by introducing memory-efficient gradient checkpointing strategies tailored for the general class of sparse RNNs and Spiking Neural Networks (SNNs). SNNs are energy efficient alternatives to RNNs thanks to their local, event-driven operation and potential neuromorphic implementation. We use the Intelligence Processing Unit (IPU) as an exemplary platform for architectures with distributed local memory. We exploit its suitability for sparse and irregular workloads to scale SNN training on long sequence lengths. We find that Double Checkpointing emerges as the most effective method, optimizing the use of local memory resources while minimizing recomputation overhead. This approach reduces dependency on slower large-scale memory access, enabling training on sequences over 10 times longer or 4 times larger networks than previously feasible, with only marginal time overhead. The presented techniques demonstrate significant potential to enhance scalability and efficiency in training sparse and recurrent networks across diverse hardware platforms, and highlights the benefits of sparse activations for scalable recurrent neural network training.