[论文解读] Neural Networks with Cheap Differential Operators

Ricky T. Q. Chen, David Duvenaud|arXiv (Cornell University)|Dec 8, 2019

Model Reduction and Neural Networks被引用 12

一句话总结

该论文提出HollowNet，一种神经网络架构，通过仅使用k次反向模式自动微分（reverse-mode automatic differentiation）即可高效计算逐维微分算子（如散度和雅可比对角线），无论输入维度如何，均可实现k阶导数的高效计算。通过将网络结构划分为调节器（conditioner）和变换器（transformer）两部分，该方法分离了各维之间的依赖关系，从而实现对这些算子的精确、低成本访问，显著提升了求解隐式常微分方程（ODEs）、连续归一化流（continuous normalizing flows）以及随机微分方程（SDEs）的福克-普朗克方程匹配的效率。

ABSTRACT

Gradients of neural networks can be computed efficiently for any architecture, but some applications require differential operators with higher time complexity. We describe a family of restricted neural network architectures that allow efficient computation of a family of differential operators involving dimension-wise derivatives, used in cases such as computing the divergence. Our proposed architecture has a Jacobian matrix composed of diagonal and hollow (non-diagonal) components. We can then modify the backward computation graph to extract dimension-wise derivatives efficiently with automatic differentiation. We demonstrate these cheap differential operators for solving root-finding subproblems in implicit ODE solvers, exact density evaluation for continuous normalizing flows, and evaluating the Fokker--Planck equation for training stochastic differential equation models.

研究动机与目标

• 解决标准神经网络中计算散度和雅可比对角线等微分算子时计算成本过高的问题。
• 设计一种神经网络架构，实现对逐维导数的高效、可扩展访问，且计算复杂度不随输入维度增加而上升。
• 通过匹配福克-普朗克方程，实现对随机微分方程（SDEs）的高效参数估计。
• 通过支持精确的密度评估，提升连续归一化流的效率，避免使用随机近似方法。
• 通过高效计算根查找子问题，支持隐式ODE求解器的运行。

提出的方法

• 网络被划分为一个调节器网络，其从除x_i外的所有输入中计算隐藏表示h_i，确保h_i与x_i无关。
• 随后，变换器网络计算f_i(x) = τ_i(x_i, h_i)，其中τ_i是作用于x_i与h_i拼接后的输入的神经网络。
• 该架构确保f的雅可比矩阵被分解为对角分量（来自x_i到f_i）与空心分量（来自h_i到f_i），从而实现导数的分离。
• 通过修改反向计算图，排除对调节器网络的梯度传播，仅保留逐维导数，从而在单次反向传播中即可高效提取∂f_i/∂x_i。
• 该方法利用反向模式自动微分，但仅限制梯度流经对角分量，将k阶导数的计算复杂度从O(d)降低至O(k)。
• 该方法与标准深度学习操作兼容，并可通过掩码权重矩阵支持全连接和卷积架构。

实验结果

研究问题

• RQ1 我们能否设计一种神经网络架构，实现对散度和雅可比对角线等逐维导数的高效计算？
• RQ2 我们能否将计算这些算子的计算成本从O(d)降低至k阶导数仅需O(k)次反向模式AD评估？
• RQ3 该架构能否通过支持快速根查找子问题，提升求解刚性隐式ODE的效率？
• RQ4 它能否在连续归一化流中实现无需随机近似的精确密度评估？
• RQ5 它能否在观测稀疏的情况下，通过福克-普朗克方程匹配，支持SDE的功能性参数估计方法？

主要发现

• HollowNet无论输入维度d如何，仅需k次反向模式AD即可计算k阶逐维导数，将复杂度从O(kd)降低至O(k)。
• 该方法可在单次反向传播中实现精确的散度与雅可比对角线计算，对高维问题具有极高效率。
• 在连续归一化流中，HollowNet可实现无需随机似然估计的精确密度评估，显著提升了训练稳定性和准确性。
• 在SDE参数估计中，HollowNet结合福克-普朗克方程匹配在稀疏数据条件下优于伪最大似然方法，且在传统基于离散化的方法失效时仍表现良好。
• 该方法成功恢复了随机摆动SDE中的多模态动力学，包括正确的边缘分布与正弦行为，即使在观测稀疏时亦能保持准确。
• 该方法对非平稳数据具有鲁棒性，在传统基于离散化方法因采样稀疏而失效的场景下仍能保持精度。