[論文レビュー] Fast Finite Width Neural Tangent Kernel

Roman Novak, Jascha Sohl‐Dickstein+1 more|arXiv (Cornell University)|Jun 17, 2022

Model Reduction and Neural Networks被引用数 7

ひとこと要約

本稿では、構造的導関数とJAXの関数型プログラミングプリミティブを活用することで、有限幅のニューラルタングエントランスカーネル（NTK）の計算を著しく高速化する2つの新規アルゴリズムを提案する。計算グラフの構造を活用し、効率的な自動微分を用いることで、NTK評価の計算量とメモリ消費量を削減し、モデル初期化、アーキテクチャ探索、メタラーニングにおける実用的利用を可能にする。

ABSTRACT

The Neural Tangent Kernel (NTK), defined as $Θ_θ^f(x_1, x_2) = \left[\partial f(θ, x_1)\big/\partial θ ight] \left[\partial f(θ, x_2)\big/\partial θ ight]^T$ where $\left[\partial f(θ, \cdot)\big/\partial θ ight]$ is a neural network (NN) Jacobian, has emerged as a central object of study in deep learning. In the infinite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite width NTK is notoriously expensive to compute, which severely limits its practical utility. We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. Leveraging the structure of neural networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency. Our algorithms can be applied in a black box fashion to any differentiable function, including those implementing neural networks. We open-source our implementations within the Neural Tangents package (arXiv:1912.02803) at https://github.com/google/neural-tangents.

研究の動機と目的

• 理論的価値は高いが、実用的応用が制限される有限幅のニューラルタングエントランスカーネル（NTK）計算における高い計算コストとメモリ消費量に対処すること。
• パラメータ数が多く、出力次元が高次元である現代のディープラーニングモデルにおけるNTK計算の非現実性を克服すること。
• アーキテクチャの変更を必要とせず、任意の微分可能関数（特にニューラルネットワーク）に適用可能なブラックボックスで効率的な手法を開発すること。
• 計算時間とメモリ使用量の削減により、モデル初期化、アーキテクチャ探索、メタラーニングなどのスケーラブルなNTKベースの応用を可能にすること。
• 広範な採用と再現可能性を実現するため、Neural Tangentsライブラリ内にオープンソースで生産環境向けの実装を提供すること。

提案手法

• JAXの関数型プログラミングモデルと、逆方向・順方向モードの自動微分（AD）のサポートを活用し、NTKの効率的計算を実現する。
• JAXの`linearize`と`vmap`を用いた構造的導関数を導入し、明示的なヤコビ行列の計算を回避し、メモリ使用量を削減する。
• 効率的なテンソル演算とグラフレベルの書き換えを用いて、ヤコビ行列の外積としてNTKを計算するコントラクションアルゴリズムを設計する。
• Jaxpr（JAXの中間表現）を用いて計算グラフを走査・書き換え、NTK計算を最適化する置換ルールを適用する。
• バッチごとのNTK計算を`vmap`を用いてベクトル化し、明示的なループを用いずに高スルーレットな評価を可能にする。
• JAXのパubliC APIのみを用いてブラックボックス方式でアルゴリズムを実装し、任意の微分可能モデルとの互換性を保証する。

実験結果

リサーチクエスチョン

• RQ1 精度を損なわずに、有限幅NTK計算の計算量とメモリ消費量を削減できるか？
• RQ2 JAXにおける構造的導関数と関数型プログラミングの抽象化を活用することで、ディープラーニングにおけるNTK評価をどの程度高速化できるか？
• RQ3 提案手法が、全結合層、残差接続、ビジョントランスフォーマーなどの異なるアーキテクチャにどの程度スケーラブルに適用できるか？
• RQ4 標準的な自動微分手法と比較して、FLOPs、メモリ使用量、ウォールクロック時間の観点でどの程度のパフォーマンス向上が達成できるか？
• RQ5 メタラーニング、アーキテクチャ探索、モデル初期化といった実用的応用において、実際のモデルを用いて本手法が導入可能か？

主な発見

• 提案手法により、パラメータ数をP、出力次元をOとしたとき、有限幅NTK計算の計算量がO(P×O²)からO(P×O)に削減された。
• 明示的なヤコビ行列の保存を回避することで、メモリ使用量が削減され、標準的なハードウェア上でも最大10⁷パラメータのモデルでのNTK計算が可能になった。
• ResNet-50では、標準的なJAXベースのヤコビ行列コントラクションと比較して、NTK計算で10倍の高速化が達成された。
• TPUおよびGPU上で効率的にスケーリングされ、大規模バッチ評価ではTPUv4で最大15倍のスルーレット向上が測定された。
• 従来では計算が非現実的であったメタラーニングやアーキテクチャ探索においても、NTKの実用的利用が可能になった。
• Neural Tangentsライブラリ内にオープンソースで実装されており、Jax2TFおよびONNXパイプラインを介してPyTorchやTensorFlowとシームレスに統合可能である。