Flexible and Efficient Surrogate Gradient Modeling with Forward Gradient Injection

自动微分是当前深度学习框架的一个关键特性。此外，它们通常提供了各种方式来指定计算图中的自定义梯度，这在非可微分操作领域，例如脉冲神经网络（SNNs）中定义替代梯度非常重要，比如海维赛德函数。例如，PyTorch允许通过覆盖其反向方法来自定义操作的反向传递。其他框架提供类似的选项。虽然这些方法是常见的实践并且通常效果良好，但它们也有一些缺点，例如灵活性有限、额外的源代码开销、较差的可用性或对自动模型优化过程的有效性可能产生强烈的负面影响。本文提出了一种替代的替代梯度形式的方法，即前向梯度注入（FGI）。FGI将基本标准操作的简单但有效的组合应用于直接在前向传递中向计算图注入任意梯度形状。实验证明，使用FGI是简单和方便的。此外，还展示了在使用TorchScript时，与SNNs中的自定义反向方法相比，FGI可以显著提高模型性能。这些结果还通过使用TorchScript和torch.compile对循环SNNs进行了一般性能研究，揭示了与纯PyTorch相比训练加速超过7倍，推理加速超过16倍的潜力。

本论文探讨了在深度学习中使用前向梯度注入（FGI）来定义替代梯度的方法，以避免使用自定义反向传播方法的缺点。

FGI是一种直接在前向传递中注入任意梯度形状的简单而有效的组合基本标准操作的方法。

FGI使用简单方便，可以显著提高SNNs模型的性能。实验使用TorchScript和torch.compile进行了性能研究，结果显示与纯PyTorch相比，训练速度提高了7倍以上，推理速度提高了16倍以上。

最近在这个领域中，也有其他相关的研究，如使用自定义反向传播方法来定义替代梯度的方法。