[论文解读] A Geometric Theory of Higher-Order Automatic Differentiation

Michael Betancourt|arXiv (Cornell University)|Dec 30, 2018

Numerical methods for differential equations参考文献 10被引用 14

一句话总结

本文提出了一种基于喷射丛的微分几何框架，用于高阶自动微分，能够系统地识别有效的高阶微分算子，并通过计算图实现显式、高效的传播规则。核心贡献在于建立了一个统一的理论，将高阶偏导数的计算与传播更新解耦，揭示了软件实现中的性能优化机会，并凸显了混合模式算法中内存使用与计算冗余之间的权衡。

ABSTRACT

First-order automatic differentiation is a ubiquitous tool across statistics, machine learning, and computer science. Higher-order implementations of automatic differentiation, however, have yet to realize the same utility. In this paper I derive a comprehensive, differential geometric treatment of automatic differentiation that naturally identifies the higher-order differential operators amenable to automatic differentiation as well as explicit procedures that provide a scaffolding for high-performance implementations.

研究动机与目标

• 解决计算统计学与机器学习中缺乏系统性、高性能的高阶自动微分框架的问题。
• 识别并形式化可应用于自动微分的高阶微分算子，超越一阶梯度。
• 解决现有高阶AD工具中递归一阶方法引发的混淆与低效问题。
• 提供一个几何基础，通过将导数计算与传播过程分离，实现高阶导数的显式、优化实现。
• 揭示在实现混合模式高阶AD时，内存使用与计算冗余之间的关键性能权衡。

提出的方法

• 使用喷射丛形式化高阶自动微分，其将切向量推广以编码高阶导数。
• 引入速度空间与余速度空间，用于建模高阶导数在计算图中的前向与后向传播。
• 推导出沿复合函数向前传播速度与向后拉回余速度的显式变换规则。
• 开发局部混合模式算法，通过在每个计算节点追踪多个导数分量，计算高阶导数（如海森矩阵、三阶张量）。
• 提出一种基于栈的数据结构，用于存储值、速度与高阶余速度（如 a, b, g, e），以实现高效的反向传播扫描。
• 证明显式、非递归的高阶AD算法可通过避免冗余计算并支持更优缓存策略，优于递归一阶AD。

实验结果

研究问题

• RQ1 哪些高阶微分算子天然适用于自动微分？如何系统地识别它们？
• RQ2 如何利用喷射丛的几何结构推导出高阶导数传播的正确且高效的规则？
• RQ3 在内存与计算方面，递归一阶AD与显式高阶AD算法之间存在哪些性能权衡？
• RQ4 将高阶偏导数计算与传播逻辑分离，如何改善软件实现？
• RQ5 在复杂函数中，中间计算的共享对高阶AD中的缓存与性能有何影响？

主要发现

• 喷射丛形式化提供了一个严格的几何基础，能自然识别有效的高阶微分算子，如海森矩阵与三阶导数。
• 通过将偏导数计算与传播过程分离，可推导出显式、非递归的高阶AD算法，从而实现更高效的实现。
• 混合模式高阶AD方法需谨慎管理中间值，因为共享的高成本计算可在不同导数分量间缓存，以减少冗余。
• 使用局部混合模式算法计算高阶导数的成本与输入变量数量呈线性关系，使其在中等维度下具有实用性。
• 对于高维目标函数，纯反向模式海森矩阵计算可能优于多轮扫描的混合模式方法，因其具有更好的内存局部性并减少冗余计算。
• 该理论表明，递归应用一阶AD会掩盖高阶算子的真实结构，阻碍优化。