基于前向梯度注入的灵活高效替代梯度建模
机器学习
2024-06-04 v1 神经与进化计算
摘要
自动微分是当前深度学习框架的一个关键特性。此外，它们通常提供各种方法来在计算图中指定自定义梯度，这对于在不可微操作（如脉冲神经网络中的Heaviside函数）领域定义替代梯度尤为重要。例如，PyTorch允许通过重写其backward方法来自定义指定操作的反向传播。其他框架也提供类似选项。虽然这些方法是常见实践且通常效果良好，但它们也存在一些缺点，例如灵活性有限、额外的源代码开销、可用性差，或对自动模型优化过程的有效性产生潜在的强烈负面影响。本文提出了一种替代梯度的新表述方法，即前向梯度注入（FGI）。FGI应用了一种简单但有效的基本标准操作组合，直接在前向传播中将任意梯度形状注入计算图。研究表明，使用FGI简单方便。此外，结果表明，在使用TorchScript时，与脉冲神经网络中的自定义反向方法相比，FGI可以显著提高模型性能。这些结果通过对使用TorchScript和torch.compile的循环脉冲神经网络的通用性能研究得到补充，揭示了与纯PyTorch相比，训练速度提升超过7倍、推理速度提升超过16倍的潜力。
引用
@article{arxiv.2406.00177,
title = {Flexible and Efficient Surrogate Gradient Modeling with Forward Gradient Injection},
author = {Sebastian Otte},
journal= {arXiv preprint arXiv:2406.00177},
year = {2024}
}