[論文レビュー] Improving Transformers with Probabilistic Attention Keys

Tam Nguyen, Tan M. Nguyen+6 more|arXiv (Cornell University)|Oct 16, 2021

Topic Modeling被引用数 6

ひとこと要約

本稿では、標準的なアテンションキーをガウス分布の混合モデルに置き換えることで、アテンションヘッド間の冗長性を低減する、新しいトランスフォーマー・アーキテクチャ「Transformer-MGK」を提案する。確率的に入力キーをモデル化することにより、学習および推論の高速化、パラメータとFLOPsの削減が可能となり、少ないヘッド数で標準的なトランスフォーマーと同等またはそれ以上の性能を達成する。これは、Wikitext-103およびLong Range Arenaベンチマークで実証された。

ABSTRACT

Multi-head attention is a driving force behind state-of-the-art transformers, which achieve remarkable performance across a variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with a Mixture of Gaussian Keys (Transformer-MGK), a novel transformer architecture that replaces redundant heads in transformers with a mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer-MGK accelerates training and inference, has fewer parameters, and requires fewer FLOPs to compute while achieving comparable or better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attention. We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications, including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.

研究の動機と目的

• 複数ヘッドアテンション機構において観察される冗長性（多くのヘッドが類似したアテンションパターンを学習すること）を解消すること。
• モデルのパフォーマンスを損なわせることなく、トランスフォーマーにおける計算およびパラメータのオーバーヘッドを低減すること。
• 各ヘッドの表現能力を向上させるための確率的フレームワークを構築し、アテンションキーをモデル化すること。
• 複数の冗長なヘッドを1つのより表現力の高いガウス混合キー機構に置き換えることで、効率的な推論および学習を可能にすること。
• 長序列列処理に適した線形アテンションに対してもこのフレームワークを拡張すること。

提案手法

• アテンションスコアをガウス・ミックスチャネル・モデル（GMM）における事後確率にリンクさせることで、自己アテンションを確率的解釈する。
• 各マルチヘッドアテンションヘッドの標準的なキー・ベクトルを、ガウス成分の混合に置き換える。これにより、1つのヘッドが同時に複数の異なるシーケンス位置に注目できるようになる。
• アテンション計算を、ガウス成分ごとの重み付き和として導出する。ここで、アテンション重みはGMM下での事後確率に対応する。
• ガウス成分の平均、分散、混合重みを最適化する微分可能学習手順を導入する。
• 標準的なトランスフォーマー・ブロックにガウス混合キーを統合し、既存のアーキテクチャおよび最適化手法と互換性を保つ。
• 長序列列処理に適した線形アテンションへ拡張するため、確率的キー表現を用いてアテンション機構を近似する。

実験結果

リサーチクエスチョン

• RQ1 アテンションキーの確率的モデル化により、トランスフォーマーにおけるマルチヘッドアテンションヘッド間の冗長性を低減できるか？
• RQ2 標準的なキーをガウス分布の混合に置き換えることで、アテンションヘッド1つあたりの表現能力を向上させつつ、モデルの複雑さを低減できるか？
• RQ3 標準的なトランスフォーマーと比較して、少ないアテンションヘッド数でTransformer-MGKがパフォーマンスを維持または向上できる程度はどの程度か？
• RQ4 ガウス混合キー機構は、長序列列タスクおよびリソースが限られた環境でも効果的に機能するか？
• RQ5 提案手法を線形アテンション機構に効率的に拡張できるか？これにより、スケーラブルな長文脈モデリングが可能になるか？

主な発見

• Wikitext-103言語モデリングベンチマークにおいて、4ヘッドのTransformer-MGKは、8ヘッドの標準的トランスフォーマーと同等またはそれ以上の性能を達成した。
• Long Range Arenaベンチマークにおいて、4ヘッドのTransformer-MGKは、8ヘッドの標準的トランスフォーマーと同等またはそれ以上の性能を示し、特に長文脈タスクで顕著な優位性を示した。
• ヘッドの冗長性が低減されたことにより、FLOPsとパラメータ数が削減され、学習および推論が高速化された。
• ガウス混合キー機構により、各アテンションヘッドが入力シーケンスの複数の異なる部分に同時に注目できるようになり、表現の多様性が向上した。
• この手法は線形アテンションとも互換性があり、非常に長いシーケンスの効率的処理を可能にした。
• 実験的結果から、確率的キーのモデリングにより、さまざまなNLPタスクにおいてより安定的かつ汎用性の高いアテンションパターンが得られることを示した。