[论文解读] Efficient Computation of Hessian Matrices in TensorFlow

Geir Kjetil Nilsen, Antonella Z. Munthe-Kaas+2 more|arXiv (Cornell University)|May 14, 2019

Matrix Theory and Algorithms参考文献 2被引用 7

一句话总结

本文提出了一种在 TensorFlow 中高效计算精确与近似 Hessian 矩阵的方法，通过使用向量化梯度运算和近似特征分解，克服了 TensorFlow 原生微分的局限性。主要贡献是通过开源的 pyhessian 模块实现了可扩展的实现，将空间复杂度从二次方降低，从而在大型模型中实现实际的曲率分析。

ABSTRACT

The Hessian matrix has a number of important applications in a variety of different fields, such as optimzation, image processing and statistics. In this paper we focus on the practical aspects of efficiently computing Hessian matrices in the context of deep learning using the Python scripting language and the TensorFlow library. We define a general feed-forward neural network model and show how to efficiently compute two quantities: the cost function's exact Hessian matrix, and the cost function's approximate Hessian matrix, known as the Outer Product of Gradients (OPG) matrix. Furthermore, as the number of parameters (weights and biases) in deep learning usually is very large, we show how to reduce the quadratic space complexity by an efficient implementation based on approximate eigendecompositions.

研究动机与目标

• 解决 TensorFlow 内置 tf.hessians() 在深度学习中计算 Hessian 矩阵时效率低下和不准确的问题。
• 实现在大规模模型中对精确 Hessian 矩阵和梯度外积（OPG）近似进行实际计算。
• 通过近似特征分解，将具有大量参数（P）的模型中 Hessian 计算的空间复杂度从 O(P²) 降低至更可扩展的形式。
• 提供一个稳健的开源实现（pyhessian），支持通用的前馈神经网络和卷积神经网络架构。

提出的方法

• 该方法使用向量化梯度运算，通过对模型参数微分梯度向量来计算二阶导数，避免了逐样本梯度计算的性能瓶颈。
• 将 Hessian 计算重新表述为使用梯度向量的雅可比矩阵的矩阵-向量积，通过 TensorFlow 的自动微分实现高效计算。
• 梯度外积（OPG）矩阵通过逐样本梯度的外积计算，为真实 Hessian 提供低秩近似。
• 对 Hessian 应用近似特征分解，将空间复杂度从 O(P²) 降低至 O(KP)，其中 K 为保留的主要特征值数量。
• 该方法被推广以处理全连接层、激活函数以及如 Softmax 交叉熵等损失函数，通过将模型参数正确展平为向量 ω ∈ ℝᴾ。
• 该实现作为开源的 pyhessian 模块发布于 GitHub，专为与 TensorFlow 集成并可扩展至大型模型而设计。

实验结果

研究问题

• RQ1 如何在 tf.hessians() 和逐样本梯度计算存在局限性的前提下，高效地在 TensorFlow 中计算 Hessian 矩阵？
• RQ2 在参数数量 P 较高的大规模深度学习模型中，近似 Hessian 矩阵的最有效方法是什么？
• RQ3 如何在实践中将 Hessian 计算的空间复杂度从 O(P²) 降低至更可扩展的形式？
• RQ4 在深度学习中，梯度外积（OPG）矩阵能否作为真实 Hessian 的可靠低秩近似？
• RQ5 如何通过主要特征子空间和一个常数特征值来构建 Hessian 的满秩近似？

主要发现

• 所提出的方法通过使用向量化梯度微分，成功克服了 TensorFlow 在计算完整 Hessian 矩阵方面的局限，实现了准确的二阶导数计算。
• 结果表明，梯度外积（OPG）矩阵是 Hessian 的有效且高效的低秩近似，尤其适用于大型模型。
• 通过应用近似特征分解，该方法将空间复杂度从 O(P²) 降低至 O(KP)，其中 K 为保留的主要特征值数量，使大规模模型的 Hessian 计算成为可能。
• 满秩 Hessian 近似被构造为 H̃̃ = QₗΛₗQₗᵀ + λ̃(I - QₗQₗᵀ)，确保所有特征值为正且矩阵可逆。
• 该实现以开源的 pyhessian 模块形式发布，支持通用的前馈神经网络和卷积神经网络架构，以及如 Softmax 交叉熵等标准损失函数。