[論文レビュー] Eigen Analysis of Self-Attention and its Reconstruction from Partial Computation

Srinadh Bhojanapalli, Ayan Chakrabarti+4 more|arXiv (Cornell University)|Jun 16, 2021

Advanced Neural Network Applications参考文献 22被引用数 5

ひとこと要約

この論文では、自己注意行列の完全な再構築手法を提案する。この手法は、ドット積注意スコアの部分集合のみを計算することで、注意スコアが低次元固有空間に存在することに着目している。グリーディアルゴリズムにより重要なクエリ-キー対を選択し、残りのスコアを学習された線形変換で再構築することで、FLOPsを25–45%削減しつつ、MLMおよびNLIタスクにおけるモデル精度の低下を最小限に抑えることができる。

ABSTRACT

State-of-the-art transformer models use pairwise dot-product based self-attention, which comes at a computational cost quadratic in the input sequence length. In this paper, we investigate the global structure of attention scores computed using this dot product mechanism on a typical distribution of inputs, and study the principal components of their variation. Through eigen analysis of full attention score matrices, as well as of their individual rows, we find that most of the variation among attention scores lie in a low-dimensional eigenspace. Moreover, we find significant overlap between these eigenspaces for different layers and even different transformer models. Based on this, we propose to compute scores only for a partial subset of token pairs, and use them to estimate scores for the remaining pairs. Beyond investigating the accuracy of reconstructing attention scores themselves, we investigate training transformer models that employ these approximations, and analyze the effect on overall accuracy. Our analysis and the proposed method provide insights into how to balance the benefits of exact pair-wise attention and its significant computational expense.

研究の動機と目的

• トランスフォーマー・モデルにおける自然言語入力の自己注意スコアの低次元構造を調査すること。
• 異なる層やモデル間で、注意スコアの変動が少数の主成分で捉えられるかどうかを特定すること。
• 正確なドット積計算の部分集合のみを用いて、完全な注意行列を近似する手法を開発すること。
• エンドツーエンド学習中の部分的注意計算が、下流モデル性能に与える影響を評価すること。
• 異なるモデルや層に共通する固有空間が、一般化可能な再構築手法の実現を可能にするかどうかを検討すること。

提案手法

• 完全な注意スコア行列および個々のクエリ行行列に対して固有値分析を実施し、主要な主成分を同定する。
• 情報量の増加を最大化する基準に基づき、グリーディアルゴリズムを用いて正確な注意計算を実施するクエリ-キー対のサブセットを選択する。
• 最小二乗最適化により、選択されたスコアを完全な注意ベクトルにマッピングする再構築行列 R を学習する。
• 部分的計算と再構築をトランスフォーマーの学習パイプラインに統合し、選択および再構築プロセスをバックプロパゲーション可能にする。
• モデル重みと同時に再構築行列 R を学習し、層ごとに共有するか、層ごとに別々に学習することで、タスク固有の分布に適応する。
• BERT-baseを用いて、マスク言語モデル学習とMNLIベンチマークにおける微調整の評価を実施する。

実験結果

リサーチクエスチョン

• RQ1 自然言語入力で計算された自己注意スコアは、その変動に低次元構造を示すか？
• RQ2 異なる層、ヘッド、トランスフォーマー・モデル間で、注意スコアの主成分がどの程度共有されているか？
• RQ3 少数の正確なスコアを用いて線形再構築により、完全な行列を低誤差で再構築できるか？
• RQ4 部分的注意計算と再構築を用いたトランスフォーマーの学習は、下流NLPタスクで競争力ある性能を維持できるか？
• RQ5 再構築行列の選択（固定 vs. 学習可能、共有 vs. 層ごと）が、モデル精度と効率性に与える影響はいかほどか？

主な発見

• 異なる層やモデル間で、注意スコアの主固有空間に顕著な重なりが見られ、共通の低次元構造が存在することが示された。
• 1クエリあたり16–32個の正確なスコアのみを用いても、完全な行列の再構築誤差は低く、平均二乗誤差が学習された再構築行列によって最小化された。
• バックプロパゲーション中に再構築行列 R を学習することで、固定で事前に計算された R を用いる場合よりも性能が向上し、特に k が小さい場合に顕著であった。
• k が小さい（例：k=16）場合、層間で共有された R を使用すると性能が向上するが、k が大きい（例：k=32）場合は、層ごとの R がより優れた性能を示した。
• この手法により、注意のFLOPsを最大45%削減しつつ、MNLIの下流タスクでベースライン精度の97%を維持できた。
• k=32の場合、部分的注意を用いたモデルはMNLIで79.7%の精度を達成したのに対し、完全なモデルでは81.6%であった。これは、精度と効率の優れたトレードオフを示している。