MoH: Multi-Head Attention as Mixture-of-Head Attention

要約

この作業では、Transformer モデルの中核であるマルチヘッド アテンション メカニズムをアップグレードし、以前の精度レベルを維持または上回ると同時に効率を向上させます。
複数の頭の注意が合計形式で表現できることを示します。
すべてのアテンションヘッドが同等の重要性を持つわけではないという洞察に基づいて、我々は、アテンションヘッドを混合専門家 (MoE) メカニズムの専門家として扱う新しいアーキテクチャである混合ヘッド注意 (MoH) を提案します。
MoH には 2 つの大きな利点があります。 まず、MoH により、各トークンが適切なアテンション ヘッドを選択できるようになり、精度を犠牲にしたりパラメータの数を増やすことなく推論効率が向上します。
第 2 に、MoH はマルチヘッド アテンションの標準的な合計を重み付けされた合計に置き換え、アテンション メカニズムに柔軟性を導入し、さらなるパフォーマンスの可能性を解き放ちます。
ViT、DiT、および LLM に関する広範な実験により、MoH はアテンション ヘッドの 50 ～ 90% のみを使用することで、マルチヘッド アテンションよりも優れたパフォーマンスを発揮することが実証されました。
さらに、LLaMA3-8B などの事前トレーニング済みマルチヘッド アテンション モデルを MoH モデルにさらに継続的に調整できることを示します。
特に、MoH-LLaMA3-8B は 14 のベンチマーク全体で 64.0% の平均精度を達成し、アテンション ヘッドの 75% のみを利用することで LLaMA3-8B を 2.4% 上回っています。
私たちは、提案された MoH がマルチヘッド アテンションに代わる有望な代替品であり、高度で効率的なアテンション ベースのモデルを開発するための強力な基盤を提供すると信じています。

要約(オリジナル)

In this work, we upgrade the multi-head attention mechanism, the core of the Transformer model, to improve efficiency while maintaining or surpassing the previous accuracy level. We show that multi-head attention can be expressed in the summation form. Drawing on the insight that not all attention heads hold equal significance, we propose Mixture-of-Head attention (MoH), a new architecture that treats attention heads as experts in the Mixture-of-Experts (MoE) mechanism. MoH has two significant advantages: First, MoH enables each token to select the appropriate attention heads, enhancing inference efficiency without compromising accuracy or increasing the number of parameters. Second, MoH replaces the standard summation in multi-head attention with a weighted summation, introducing flexibility to the attention mechanism and unlocking extra performance potential. Extensive experiments on ViT, DiT, and LLMs demonstrate that MoH outperforms multi-head attention by using only 50%-90% of the attention heads. Moreover, we demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B, can be further continue-tuned into our MoH models. Notably, MoH-LLaMA3-8B achieves an average accuracy of 64.0% across 14 benchmarks, outperforming LLaMA3-8B by 2.4% by utilizing only 75% of the attention heads. We believe the proposed MoH is a promising alternative to multi-head attention and provides a strong foundation for developing advanced and efficient attention-based models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google