[論文レビュー] 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のネイティブな微分処理の制限を克服するため、ベクトル化された勾配操作と近似固有値分解を用いている。主な貢献は、オープンソースのpyhessianモジュールを介したスケーラブルな実装であり、2次的な空間計算量を削減することで、大規模モデルにおける実用的な曲率解析を可能にしている。

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()の非効率性と不正確さを解消すること。
• 大規模モデルにおいて、正確なヘッセ行列と外積勾配（OPG）近似の両方を実用的に計算できることを実現すること。
• パラメータ数Pが非常に多いモデルにおいて、ヘッセ行列計算の空間計算量をO(P²)からO(KP)に削減すること。ここでKは保持する主要固有値の数である。
• 一般のフィードフォワードおよび畳み込みニューラルネットワークアーキテクチャをサポートする、堅牢でオープンソースの実装（pyhessian）を提供すること。

提案手法

• 2階微分を計算するために、モデルパラメータに関して勾配ベクトルを微分するベクトル化された勾配操作を用いることで、個々の入力例ごとの勾配計算のボトル neck を回避する。
• 勾配ベクトルのヤコビ行列を用いてヘッセ行列の計算を行列-ベクトル積として再定式化し、TensorFlowの自動微分を介して効率的な計算を実現する。
• 外積勾配（OPG）行列は、個々の入力例ごとの勾配の外積として計算され、真のヘッセ行列の低ランク近似を提供する。
• ヘッセ行列に近似固有値分解を適用することで、空間計算量をO(P²)からO(KP)に削減する。ここでKは保持する主要固有値の数である。
• 全結合層、活性化関数、およびソフトマックス交差エントロピーなどの損失関数を適切に処理するため、モデルパラメータをベクトルω ∈ ℝᴾにフラット化する。
• 実装はGitHub上でオープンソースのpyhessianモジュールとして公開されており、TensorFlowとの統合を想定し、大規模モデルにスケーラブルである。

実験結果

リサーチクエスチョン

• RQ1 tf.hessians()の制限や個々の入力例ごとの勾配計算の制約を克服し、TensorFlowでヘッセ行列を効率的に計算する方法は何か？
• RQ2 パラメータ数Pが非常に多い大規模な深層学習モデルにおいて、ヘッセ行列を近似する最も効果的な方法は何か？
• RQ3 実際の応用において、ヘッセ行列計算の空間計算量をO(P²)からよりスケーラブルな形に削減するにはどうすればよいか？
• RQ4 外積勾配（OPG）行列は、深層学習における真のヘッセ行列の信頼できる低ランク近似として機能するか？
• RQ5 主要固有空間を用いて全ランク近似のヘッセ行列を構築し、残りの部分空間に対して定数固有値を割り当てるとすれば、その方法は何か？

主な発見

• 提案手法は、ベクトル化された勾配微分を用いることで、TensorFlowの制限を克服し、正確な2階微分の計算を可能にした。
• 外積勾配（OPG）行列は、真のヘッセ行列の有効で効率的な低ランク近似であることが示され、特に大規模モデルにおいて有用である。
• 近似固有値分解を適用することで、空間計算量をO(P²)からO(KP)に削減でき、大規模モデルにおけるヘッセ行列計算が実用可能になった。ここでKは保持する主要固有値の数である。
• 全ランクヘッセ行列近似は、H̃̃ = QₗΛₗQₗᵀ + λ̃(I - QₗQₗᵀ)として構築され、すべての固有値が正であり、行列が逆行列可能であることを保証する。
• 実装はオープンソースのpyhessianモジュールとして公開されており、ソフトマックス交差エントロピーなどの標準損失関数を用いた、一般のフィードフォワードおよび畳み込みニューラルネットワークアーキテクチャをサポートする。