[論文レビュー] Infinite attention: NNGP and NTK for deep attention networks

Jiri Hron, Yasaman Bahri+2 more|arXiv (Cornell University)|Jun 18, 2020

Gaussian Processes and Bayesian Inference被引用数 29

ひとこと要約

本稿では、マルチヘッドアテンションネットワークが無限幅極限においてガウス過程（NNGP）に収束し、ニューラルタングエントランスカーネル（NTK）行動を示すことを確立した。これは、単一ヘッド設定におけるアテンションの非ガウス出力が引き起こす理論的課題を解決するものである。本稿では、有限および無限幅モデルの両方で性能を向上させる、修正されたアテンションメカニズムと位置エンコーディング戦略を提案し、データ拡張なしでCIFAR-10で最先端の結果を達成した。

ABSTRACT

There is a growing amount of literature on the relationship between wide neural networks (NNs) and Gaussian processes (GPs), identifying an equivalence between the two for a variety of NN architectures. This equivalence enables, for instance, accurate approximation of the behaviour of wide Bayesian NNs without MCMC or variational approximations, or characterisation of the distribution of randomly initialised wide NNs optimised by gradient descent without ever running an optimiser. We provide a rigorous extension of these results to NNs involving attention layers, showing that unlike single-head attention, which induces non-Gaussian behaviour, multi-head attention architectures behave as GPs as the number of heads tends to infinity. We further discuss the effects of positional encodings and layer normalisation, and propose modifications of the attention mechanism which lead to improved results for both finite and infinitely wide NNs. We evaluate attention kernels empirically, leading to a moderate improvement upon the previous state-of-the-art on CIFAR-10 for GPs without trainable kernels and advanced data preprocessing. Finally, we introduce new features to the Neural Tangents library (Novak et al., 2020) allowing applications of NNGP/NTK models, with and without attention, to variable-length sequences, with an example on the IMDb reviews dataset.

研究の動機と目的

• マルチヘッドアテンション機構を備えた深層ニューラルネットワークの無限幅極限を厳密に特徴づけること。
• アテンションレイヤーへのガウス過程およびニューラルタングエントランスカーネルフレームワークの拡張に伴う理論的課題を解決すること。これは、標準的な交換可能性や条件付き分布の仮定を満たさないためである。
• 有限および無限幅ネットワークの両方で性能を向上させる、修正されたアテンションメカニズムおよび位置エンコーディング戦略を提案すること。
• 提案されたアテンションカーネルの実験的評価を行い、GP推論を用いてCIFAR-10などのベンチマークデータセットで性能向上を示すこと。
• Neural Tangentsライブラリに可変長シーケンスをサポートする機能を追加し、テキストなどの順序データにNNGP/NTK応用を可能にすること。

提案手法

• $d^{-1/2}$および$d^{-1}$スケーリングの下で、マルチヘッドアテンションネットワークのNNGPおよびNTKカーネルを導出。無限幅極限においてガウス過程に収束することを示した。
• $d^{-1}$スケーリングと補間ベースの共分散更新を備えた修正されたアテンションメカニズムを導入。これにより表現力が保持されるとともに、理論的解析が可能になった。
• 無限幅領域における一般化と性能向上に寄与する構造的ポジショナルエンコーディングを提案。NNGP/NTK推論と組み合わせても効果を示した。
• 残差接続およびレイヤーノーマライゼーションを扱える理論枠組みに適応。これにより、カーネル収束およびモデル行動への影響を明らかにした。
• 学習可能なポジショナル埋め込みからの勾配寄与を組み込んだ新しいカーネル更新ルールを導入。ハイブリッド補間メカニズムによりNTKが変更された。
• 可変長シーケンスをサポートするようにNeural Tangentsライブラリにコード拡張をリリース。これにより、IMDbなどのデータセットでエンドツーエンドのNNGP/NTK推論が可能になった。

実験結果

リサーチクエスチョン

• RQ1 広大なニューラルネットワークにおけるマルチヘッドアテンションは、無限幅極限においてガウス過程に収束するか？もしそうなら、どのような条件下で？
• RQ2 なぜ標準的なNNGPおよびNTK導出手法は単一ヘッドアテンションでは機能しないのか？また、マルチヘッドアーキテクチャに適応するにはどうすればよいか？
• RQ3 ポジショナルエンコーディングおよびレイヤーノーマライゼーションは、無限幅アテンションネットワークにおけるカーネル行動および一般化性能にどのように影響するか？
• RQ4 d^{-1}スケーリングを備えた修正されたアテンションメカニズムは、モデルの表現力を保持しつつ、厳密なNNGP/NTK解析を可能にするか？
• RQ5 CIFAR-10などの実世界ベンチマークに、導出されたNNGPおよびNTKカーネルを適用することで、データ拡張や高度な前処理なしでどの程度の性能向上が達成できるか？

主な発見

• マルチヘッドアテンションネットワークは無限幅極限においてガウス過程に収束するが、単一ヘッドアテンションは、共有されたアテンション重みが引き起こす統計的依存性のため収束しない。
• 本稿では、$d^{-1/2}$および$d^{-1}$スケーリングの両方で有効な、マルチヘッドアテンションのNNGPおよびNTKカーネルの厳密な導出を提供。これにより、アテンション機構への理論フレームワークの拡張が可能になった。
• d^{-1}スケーリングと補間ベースの共分散更新を備えた修正されたアテンションメカニズムは、モデルの表現力を保持しつつ、安定なカーネル収束を可能にし、従来のアプローチで見られるプーリングのような振る舞いを回避した。
• ポジショナルエンコーディングは、無限幅極限においても実験的性能を顕著に向上させた。これは、モデルがトレーニングされていなくても有効であり、一般化の向上に寄与する役割を示している。
• 本手法は、データ拡張や高度な前処理なしで、CIFAR-10において以前の最先端手法をわずかに上回る性能向上を達成した。
• 著者らは、可変長シーケンスをサポートするようにNeural Tangentsライブラリにコード拡張をリリース。これにより、テキストなどの順序データに対してNNGP/NTK推論が可能になり、IMDbデータセットを用いた実例が提示された。