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.

本文介绍了一种替代方法，即前向梯度注入（FGI），用于在前向传递过程中将任意梯度形状直接注入计算图，证明了使用 FGI 相对于自定义向后传递方法在 SNNs 中使用 TorchScript 可以显著提高模型性能，并通过在 TorchScript 和 torch.compile 中对循环 SNNs 进行的一般性能研究结果补充，揭示了与纯 PyTorch 相比，训练速度可提高 7 倍以上，推理速度可提高 16 倍以上的潜力。