[論文レビュー] Tensor Programs II: Neural Tangent Kernel for Any Architecture

Greg Yang|arXiv (Cornell University)|Jun 25, 2020

Stochastic Gradient Optimization Techniques参考文献 59被引用数 53

ひとこと要約

本研究は、 neural tangent kernel (NTK) が無限幅極限において任意のニューラルネットワークアーキテクチャに対して決定論的に収束することを証明し、勾配独立性仮定が有効となる条件を明らかにする。実用的なテンソルプログラムベースの手法と参照実装を提供する。

ABSTRACT

We prove that a randomly initialized neural network of *any architecture* has its Tangent Kernel (NTK) converge to a deterministic limit, as the network widths tend to infinity. We demonstrate how to calculate this limit. In prior literature, the heuristic study of neural network gradients often assumes every weight matrix used in forward propagation is independent from its transpose used in backpropagation (Schoenholz et al. 2017). This is known as the *gradient independence assumption (GIA)*. We identify a commonly satisfied condition, which we call *Simple GIA Check*, such that the NTK limit calculation based on GIA is correct. Conversely, when Simple GIA Check fails, we show GIA can result in wrong answers. Our material here presents the NTK results of Yang (2019a) in a friendly manner and showcases the *tensor programs* technique for understanding wide neural networks. We provide reference implementations of infinite-width NTKs of recurrent neural network, transformer, and batch normalization at https://github.com/thegregyang/NTK4A.

研究の動機と目的

• 幅が無限大に成長するにつれて任意のアーキテクチャをもつランダムに初期化されたネットワークは決定論的なNTK極限を持つことを示す。
• 勾配独立性仮定がNTK計算に対して正しいNTK極限を導く条件を特定する（Simple GIA Check）。
• Simple GIA Check が成り立たない場合には勾配独立性の失敗がNTK計算に与える影響を示し、正しいNTK計算の指針を提供する。
• テンソルプログラムを用いたNTK極限の友好的な解説を広いネットワークに対して提示する。
• RNN、Transformer、およびバッチ正規化の無限幅NTKの参照実装を提供する。

提案手法

• テンソルプログラムのフレームワークを適用して広いニューラルネットワークを分析しNTK極限を導出する。
• 勾配独立性仮定（GIA）を形式化し、Simple GIA Checkを正しいNTK計算の条件として確立する。
• 幅が無限大に向かうとき任意のアーキテクチャに対してNTKが決定論的な極限に収束することを証明する。
• リカレントニューラルネットワーク、Transformer、バッチ正規化のNTK極限を導出・提示する。
• 無限幅NTKを計算するための実用的な指針と参照実装を提供する。

実験結果

リサーチクエスチョン

• RQ1 無限幅極限においてNTKは任意のニューラルネットワークアーキテクチャに対して決定論的な極限へ収束するのか？
• RQ2 NTK計算において勾配独立性仮定が有効である条件は何か（Simple GIA Check）？
• RQ3 RNN、Transformer、バッチ正規化などの特定のアーキテクチャに対してNTK極限をどのように計算できるか？
• RQ4 GIAの失敗モードは何か、そしてそれらはNTKの結果にどのように影響するか？
• RQ5 テンソルプログラムは広いネットワークにおけるNTK極限を理解するための友好的で実用的なフレームワークを提供できるのか？

主な発見

• NTKはネットワーク幅が無限大に向かうとき、任意のアーキテクチャに対して決定論的な極限へ収束する。
• Simple GIA CheckはGIAベースのNTK計算が正しいという条件を特定する。
• Simple GIA Checkが成り立たない場合、GIAはNTKの結果を誤らせる可能性がある。
• 本論文はテンソルプログラムを用いたNTK結果の友好的な解説を提示する。
• RNN、Transformer、バッチ正規化の無限幅NTKの参照実装が提供される。