Weighted Grouped Query Attention in Transformers

要約

アテンション メカニズムは、トランスフォーマー言語モデルの基本ブロックを形成します。
最近のアプローチでは、モデルをスケーリングすることで人間レベルのパフォーマンスが達成されることが示されています。
ただし、スケーリングに対する要求の増大とハードウェア メモリの制約により、これらのモデルの推論コストは依然として高くなっています。
推論時間を短縮するために、マルチクエリ アテンション (MQA) とグループ化クエリ アテンション (GQA) がそれぞれ (Shazeer、2019) と (Ainslieet al.、2023) で提案されました。
この論文では、重み付けグループ化クエリ アテンション (WGQA) と呼ばれる、グループ化クエリ アテンションのバリエーションを提案します。
T5 デコーダー アテンション ブロックの各キーと値のヘッドに新しい学習可能なパラメーターを導入し、モデルが微調整中に加重平均を取れるようにしました。
私たちのモデルは GQA と比較して平均 0.53% の改善を達成し、推論中に追加のオーバーヘッドを発生させることなく、パフォーマンスは従来のマルチヘッド アテンション (MHA) に収束します。
これらのパラメーターの導入を評価し、その後の微調整によってトレーニング中にグループ化メカニズムについてモデルに情報を提供することで、パフォーマンスが向上しました。
さらに、T5 スモール アーキテクチャと T5 ベース アーキテクチャの結果を比較することで、分析におけるスケーリングの法則を実証します。

要約(オリジナル)

The attention mechanism forms the foundational blocks for transformer language models. Recent approaches show that scaling the model achieves human-level performance. However, with increasing demands for scaling and constraints on hardware memory, the inference costs of these models remain high. To reduce the inference time, Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) were proposed in (Shazeer, 2019) and (Ainslieet al., 2023) respectively. In this paper, we propose a variation of Grouped-Query Attention, termed Weighted Grouped-Query Attention (WGQA). We introduced new learnable parameters for each key and value head in the T5 decoder attention blocks, enabling the model to take a weighted average during finetuning. Our model achieves an average of 0.53% improvement over GQA, and the performance converges to traditional Multi-head attention (MHA) with no additional overhead during inference. We evaluated the introduction of these parameters and subsequent finetuning informs the model about the grouping mechanism during training, thereby enhancing performance. Additionally, we demonstrate the scaling laws in our analysis by comparing the results between T5-small and T5-base architecture.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google