Flexible and Efficient Surrogate Gradient Modeling with Forward Gradient Injection

TL;DR

This paper introduces forward gradient injection (FGI), a novel method for surrogate gradient modeling that enhances flexibility and efficiency in deep learning frameworks, especially for non-differentiable operations like in spiking neural networks.

Contribution

The paper proposes FGI as an alternative to traditional custom backward methods, enabling direct gradient injection during the forward pass for improved performance and usability.

Findings

FGI significantly improves model performance in SNNs compared to traditional methods.

FGI enables training speedup of over 7x and inference speedup of over 16x with TorchScript.

FGI offers a straightforward and flexible approach for surrogate gradient modeling.

Abstract

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.