[論文レビュー] Multi-Head Attention: Collaborate Instead of Concatenate

Jean-Baptiste Cordonnier, Andreas Loukas+1 more|arXiv (Cornell University)|Jun 29, 2020

Advanced Neural Network Applications参考文献 27被引用数 76

ひとこと要約

この論文はコラボラティブ・マルチヘッド・アテンション（collab MHA）を提案し、ヘッド間でキー/クエリの射影を共有することでパラメータを削減し、NLP翻訳、NLP理解、ビジョンタスクで性能を維持または向上させる。標準MHAからのテンソル分解による事後再パラメータ化も示す。

ABSTRACT

Attention layers are widely used in natural language processing (NLP) and are beginning to influence computer vision architectures. Training very large transformer models allowed significant improvement in both fields, but once trained, these networks show symptoms of over-parameterization. For instance, it is known that many attention heads can be pruned without impacting accuracy. This work aims to enhance current understanding on how multiple heads interact. Motivated by the observation that attention heads learn redundant key/query projections, we propose a collaborative multi-head attention layer that enables heads to learn shared projections. Our scheme decreases the number of parameters in an attention layer and can be used as a drop-in replacement in any transformer architecture. Our experiments confirm that sharing key/query dimensions can be exploited in language understanding, machine translation and vision. We also show that it is possible to re-parametrize a pre-trained multi-head attention layer into our collaborative attention layer. Collaborative multi-head attention reduces the size of the key and query projections by 4 for same accuracy and speed. Our code is public.

研究の動機と目的

• MHAヘッド間のキー/クエリ射影の冗長性を特徴づける。
• 射影をヘッド間で共有するCollaborative Head Attentionを提案する。
• コラボラティブ注意機構を使用するように事前訓練済みトランスフォーマーを再パラメータライズする方法を示す。
• NLPおよび視覚タスクでのパラメータと計算効率を実証する。
• トランスフォーマー構造に対するドロップイン置換の実用的ガイダンスを提供する。

提案手法

• W_Q W_KをヘッドごとにPCAして冗長性を分析し、結合ヘッドの低ランク挙動を示す。
• Collab MHAを定義し、すべてのヘッドがW_QとW_Kを共有し、共有次元˜Dkに射影し、混成ベクトル m_iを使用する。
• collabヘッドは標準MHAの特別なケースとして関係づける（˜Dk = Nh dk）。
• Tucker/CP（カノニカル）テンソル分解を用いて、既存の訓練済みアテンション層をcollaborative attentionに再パラメータライズする（W̃_Q, W̃_K, および混合行列Mを導出）。
• コンテンツ/文脈分解とアテンションスコアの非寄与バイアス項を含むバイアス処理について議論。
• 複雑さ/パラメータ分析を提供し、プルーニングと比較。
• NMT、NLU（GLUE）、Vision（ImageNet）タスクで検証し、ゼロからの訓練と事後再パラメータライズを含む。

実験結果

リサーチクエスチョン

• RQ1 アテンションヘッドは共通のキー/クエリ射影を共有できる冗長性を学習するのか？
• RQ2 パラメータを削減しつつcollab MHAは標準MHAの性能に匹敵または上回るのか？
• RQ3 事前訓練済みトランスフォーマーを完全再訓練なしでcollaborative MHAに変換することは可能か（事後再パラメータライズ）？
• RQ4 NLPとビジョンタスクにおけるパラメータと速度の実践的な効率向上はどの程度か？
• RQ5 GLUEやImageNetなどのベンチマークでcollaborative MHAはモデル圧縮と微調整の必要性にどう影響するのか？

主な発見

• ヘッドはサブスペース間で射影を共有する傾向があり、結合されたW_Q W_Kは低ランクになるが、個別にはそうでない。
• collaborative MHAはキー/クエリ射影（˜Dk）と混合行列（M）をNhヘッド全体で共有することでパラメータ数を削減する。
• 標準MHAをcollab MHAに置換しても性能が維持または向上し、NMTでは精度低下なしでキー/クエリ次元を4倍削減可能。
• collab MHAをゼロから訓練すると、Attentionパラメータを最大35%削減できBLEUの低下なし、場合によりWMT14 EN-DEでBLEUが改善。
• 事後再パラメータライズによる事前訓練済みトランスフォーマーをコラボラティブ注意に変換して、精度を維持するかわずかに低下させつつ、GLUEで˜Dkの2×〜3×削減など大幅な圧縮が可能。
• ビジョンでは、collab MHAは結合ベースより低い˜Dk値で同等またはそれを上回る性能を発揮し、ImageNetで最小限の精度損失で圧縮を可能にする。