基于高层伴随微分的记忆高效神经ODE框架
机器学习
2023-06-12 v3
摘要
神经常微分方程（neural ODE）作为一种连接动力系统与深度学习的新颖网络架构已崭露头角。然而，原始神经ODE中采用连续伴随方法得到的梯度并非逆向精确的。其他方法要么因深层计算图导致内存需求过高，要么因时间积分方案选择受限，阻碍了其面向大规模复杂动力系统的应用。为在不牺牲内存效率与灵活性的前提下获得精确梯度，我们提出一种基于高层离散伴随算法微分的新型神经ODE框架PNODE。通过利用离散伴随时间积分器以及为这些积分器量身定制的先进检查点策略，PNODE能够在内存与计算代价之间取得平衡，同时一致且精确地计算梯度。我们提供了基于PyTorch与PETSc（最常用可移植、可扩展科学计算库之一）的开源实现。我们通过图像分类与连续归一化流问题上的大量数值实验展示了性能。结果表明，与其他逆向精确方法相比，PNODE实现了最高的内存效率。在图像分类问题上，PNODE比原始神经ODE快至多2倍，比现有最佳逆向精确方法快至多2.3倍。我们还展示了PNODE支持刚体动力系统所需的隐式时间积分方法的使用。
引用
@article{arxiv.2206.01298,
title = {A memory-efficient neural ODE framework based on high-level adjoint differentiation},
author = {Hong Zhang and Wenjun Zhao},
journal= {arXiv preprint arXiv:2206.01298},
year = {2023}
}