用于极小内存下神经 ODE 精确梯度的辛伴随方法
机器学习
2021-10-20 v2
摘要
微分方程神经网络模型，即神经 ODE，已能够实现高精度的连续时间动力系统与概率分布的学习。神经 ODE 在数值积分过程中重复使用同一网络。反向传播算法的内存消耗与重复使用次数乘以网络规模成正比。即便采用检查点方案将计算图划分为子图，这一结论依然成立。另一方面，伴随方法通过时间反向的数值积分获得梯度。尽管该方法仅消耗单次网络使用的内存，但需要较高的计算代价以抑制数值误差。本研究提出辛伴随方法，即由辛积分器求解的伴随方法。辛伴随方法以与重复使用次数加网络规模成正比的内存获得精确梯度（直至舍入误差）。实验结果表明，辛伴随方法比朴素反向传播算法和检查点方案消耗更少的内存，比伴随方法运行更快，并且对舍入误差更具鲁棒性。
引用
@article{arxiv.2102.09750,
title = {Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory},
author = {Takashi Matsubara and Yuto Miyatake and Takaharu Yaguchi},
journal= {arXiv preprint arXiv:2102.09750},
year = {2021}
}
备注
19 pages