[論文レビュー] Low-Rank Bottleneck in Multi-head Attention Models

Srinadh Bhojanapalli, Chulhee Yun+3 more|arXiv (Cornell University)|Feb 17, 2020

Topic Modeling参考文献 29被引用数 14

ひとこと要約

この論文は、ヘッド数を増加させるに従ってヘッドサイズが縮小することにより生じる低ランクのボトルネックを特定し、これにより表現力が制限され、大規模な埋め込み次元に依存するようになる。著者らは、ヘッドサイズをヘッド数とは独立して入力シーケンス長に固定することで、モデルの表現力を向上させ、より小さな埋め込み次元でも優れた性能を発揮できるようにするとともに、精度を落とさずにモデルパラメータを削減できると提言する。

ABSTRACT

Attention based Transformer architecture has enabled significant advances in the field of natural language processing. In addition to new pre-training techniques, recent improvements crucially rely on working with a relatively larger embedding dimension for tokens. Unfortunately, this leads to models that are prohibitively large to be employed in the downstream tasks. In this paper we identify one of the important factors contributing to the large embedding size requirement. In particular, our analysis highlights that the scaling between the number of heads and the size of each head in the current architecture gives rise to a low-rank bottleneck in attention heads, causing this limitation. We further validate this in our experiments. As a solution we propose to set the head size of an attention unit to input sequence length, and independent of the number of heads, resulting in multi-head attention layers with provably more expressive power. We empirically show that this allows us to train models with a relatively smaller embedding dimension and with better performance scaling.

研究の動機と目的

• 現代のトランスフォーマー・モデルにおける大規模な埋め込み次元の要請の根本的原因を調査すること。
• ヘッド数に反比例してヘッドサイズをスケーリングすることで、表現力が制限される低ランクボトルネックが生じることを特定すること。
• ヘッドサイズを入力シーケンス長に固定し、ヘッド数や埋め込み次元から分離した新しいアーキテクチャ設計を提言すること。
• この変更により、より小さなパラメータ数で効率的かつ性能が向上するモデルの訓練が可能になることを示すこと。
• ヘッドサイズが、ヘッド数や埋め込み次元よりも重要な容量制御パラメータであることを検証すること。

提案手法

• 各アテンションヘッドが入力シーケンス長に等しい次元の部分空間で動作する固定ヘッドサイズ設定を提唱する。
• ヘッドサイズをヘッド数および埋め込み次元から分離し、ランク制約なしに任意のヘッド数を許容できるようにする。
• 理論的に（定理1）ヘッドサイズがシーケンス長未満の場合に低ランクボトルネックが生じ、表現力が低下することを証明する。
• ヘッドサイズを増加させると性能が単調に向上するが、ヘッド数を増やすだけではそうならないことの実証を行う。
• LM1B、SQuAD、MNLIベンチマークで、より小さな埋め込み次元を使用して固定ヘッドサイズのモデルを訓練・評価する。
• 従来のヘッドサイズヒューリスティック（ヘッド数に比例してスケーリング）を用いた標準モデルと性能を比較する。

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャの進歩にもかかわらず、なぜ現代のトランスフォーマー・モデルは大規模な埋め込み次元を必要としているのか？
• RQ2 ヘッド数を増やすに従ってヘッドサイズを小さくすると、マルチヘッドアテンションの表現能力にどのような影響を与えるか？
• RQ3 ヘッドサイズを入力シーケンス長に固定することで、アテンションヘッドにおける低ランクボトルネックを解消できるか？
• RQ4 固定ヘッドサイズにより、パラメータ数を減らしても性能スケーリングが向上するか？
• RQ5 ヘッドサイズは、ヘッド数や埋め込み次元よりも重要な容量制御パラメータであるか？

主な発見

• 従来のヘッドサイズヒューリスティック（ヘッド数に反比例してスケーリング）は、定理1で示されるように低ランクボトルネックを生じさせ、表現力を制限する。
• 固定ヘッドサイズ（128）かつ小規模な埋め込み次元（256）のモデルが、大規模な埋め込み次元（512）を用いた標準モデルを、言語モデリングタスクで上回る性能を発揮する。
• SQuADおよびMNLIのタスクでは、512の埋め込み次元と8〜32のヘッド数を用いた固定ヘッドサイズモデルが、標準のBERT-Large（1024の埋め込み次元）を上回るF1スコアと正答率を達成する。
• ヘッドサイズを固定した場合、ヘッド数を増加させると性能が単調に向上するが、標準モデルでは閾値を超えると性能が低下する。
• ヘッドサイズを32から128に増加させると性能が顕著に向上し、ヘッドサイズが主要な容量パラメータであることが確認される。
• 固定ヘッドサイズアプローチにより、パラメータ数を減らしたにもかかわらず性能を維持または向上できる小型モデルの訓練が可能になり、推論コストが低減される。