# Flexible and Efficient Surrogate Gradient Modeling with Forward Gradient Injection

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 7 of its 11 references that are in the map.
- First posted on arXiv 31 May 2024.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Surrogate Gradient Learning in Spiking Neural Networks](https://mlsanity.com/paper/1901.09948/) (2019)
- [Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural Networks](https://mlsanity.com/paper/2007.05785/) (2020)
- [Long short-term memory and learning-to-learn in networks of spiking neurons](https://mlsanity.com/paper/1803.09574/) (2018)
- [SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence](https://mlsanity.com/paper/2310.16620/) (2023)
- [Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks](https://mlsanity.com/paper/2103.12593/) (2021)
- [Meta-learning Spiking Neural Networks with Surrogate Gradient Descent](https://mlsanity.com/paper/2201.10777/) (2022)

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

- [The Resurrection of the ReLU](https://mlsanity.com/paper/2505.22074/) (2025)
- [Pulling Back the Curtain on Deep Networks](https://mlsanity.com/paper/2507.22832/) (2025)

## Abstract (arXiv)

Automatic differentiation is a key feature of present deep learning frameworks. Moreover, they typically provide various ways to specify custom gradients within the computation graph, which is of particular importance for defining surrogate gradients in the realms of non-differentiable operations such as the Heaviside function in spiking neural networks (SNNs). PyTorch, for example, allows the custom specification of the backward pass of an operation by overriding its backward method. Other frameworks provide comparable options. While these methods are common practice and usually work well, they also have several disadvantages such as limited flexibility, additional source code overhead, poor usability, or a potentially strong negative impact on the effectiveness of automatic model optimization procedures. In this paper, an alternative way to formulate surrogate gradients is presented, namely, forward gradient injection (FGI). FGI applies a simple but effective combination of basic standard operations to inject an arbitrary gradient shape into the computational graph directly within the forward pass. It is demonstrated that using FGI is straightforward and convenient. Moreover, it is shown that FGI can significantly increase the model performance in comparison to custom backward methods in SNNs when using TorchScript. These results are complemented with a general performance study on recurrent SNNs with TorchScript and torch.compile, revealing the potential for a training speedup of more than 7x and an inference speedup of more than 16x in comparison with pure PyTorch.