[論文レビュー] Neural Tangents: Fast and Easy Infinite Neural Networks in Python

Roman Novak, Lechao Xiao+5 more|arXiv (Cornell University)|Dec 5, 2019

Model Reduction and Neural Networks参考文献 28被引用数 57

ひとこと要約

Neural Tangents は解析的 NNGP/NTK カーネルまたは有限幅近似を通じて、無限幅ニューラルネットワークを定義・訓練・分析するための JAX ベースの API を提供し、CPU/GPU/TPU のスケーラブルなサポートを実現します。

ABSTRACT

Neural Tangents is a library designed to enable research into infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finite-width as usual or in their infinite-width limit. Infinite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks in either function space or weight space. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. Neural Tangents is available at www.github.com/google/neural-tangents. We also provide an accompanying interactive Colab notebook.

研究の動機と目的

• 無限幅ニューラルネットワークを研究するためのオープンソースライブラリである Neural Tangents を紹介する。
• 複雑なアーキテクチャを指定し、解析的な無限幅カーネル（NNGP/NTK）を取得するためのハイレベル API を提供する。
• 無限幅極限における厳密なベイズ推論と勾配降下に着想を得た訓練の両方を可能にする。
• 解析形式を持たないカーネルに対してモンテカルロ近似をサポートする。
• アーキテクチャとデータセット全体でのスケーラビリティと実用性を示す。

提案手法

• 標準のテンソル演算（dense、conv、非線形性）を NNPG/NTK に対応するカーネル演算へ翻訳する定義。
• 解析形式が利用できないときには、無限幅ネットワークの解析的カーネル計算とMCベースのカーネル近似を提供する。
• 厳密なカーネルを介して無限幅極限でベイズ推論または勾配降下に似た訓練を実行する API を提供する。
• 連続時間トレーニングダイナミクスとテイラー展開近似による有限幅重み空間の探索を組み込む。
• 自動デバイス分散とバッチ処理を用いて CPU/GPU/TPU 実行のために XLA/JAX を活用する。

実験結果

リサーチクエスチョン

• RQ1 無限幅極限で広範なニューラルアーキテクチャに対して、NNGPおよびNTKカーネルを効率的に計算・利用する方法はどのようなものか。
• RQ2 有限幅ネットワークは解析的またはモンテカルロカーネル推定を用いて、無限幅の対応物を通じて正確に分析できるか。
• RQ3 デバイスとアーキテクチャを跨いだ解析的カーネル計算の性能とスケーラビリティはどの程度か。
• RQ4 異なるアーキテクチャ（例：全結合、CNN、WideResNet）は、無限幅推論とカーネルベースの訓練の下でどのように比較されるか。
• RQ5 無限幅ニューラルネットワークの研究を加速する実用的なツールとワークフローは何か。

主な発見

• Neural Tangents は、多くのアーキテクチャに対して解析的 NNGP/NTK カーネルによる無限幅ネットワークの厳密推論を可能にする。
• モンテカルロ近似は解析的カーネルが利用できない場合に柔軟な近似を提供し、幅やサンプル数が増えると解析的値に収束する。
• 実験では、無限幅推論が合成データ上で有限幅のアンサンブル訓練とほぼ一致し、実用的な忠実性を示す。
• CIFAR-10 に評価した場合、アーキテクチャ間の性能収束と階層を示し、WideResNet様の構造がテストされた変種の中で最も良い性能を示した。
• 追加のアクセラレータでほぼ線形スケーリングを達成し、バッチ処理、オペレーション融合、効率のための2D畳込みベース共分散計算をサポートする。