MoH : Attention Multi-Tête en tant qu'Attention Mélange-de-Têtes

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

October 15, 2024

Auteurs: Peng Jin, Bo Zhu, Li Yuan, Shuicheng Yan

cs.AI

Résumé

Dans ce travail, nous améliorons le mécanisme d'attention multi-têtes, au cœur du modèle Transformer, afin d'améliorer l'efficacité tout en maintenant ou dépassant le niveau de précision précédent. Nous montrons que l'attention multi-têtes peut être exprimée sous forme de sommation. En nous appuyant sur l'idée que toutes les têtes d'attention ne sont pas également significatives, nous proposons l'attention Mélange-de-Têtes (MoH), une nouvelle architecture qui considère les têtes d'attention comme des experts dans le mécanisme de Mélange-d'Experts (MoE). MoH présente deux avantages significatifs : Premièrement, MoH permet à chaque jeton de sélectionner les têtes d'attention appropriées, améliorant l'efficacité de l'inférence sans compromettre la précision ou augmenter le nombre de paramètres. Deuxièmement, MoH remplace la sommation standard dans l'attention multi-têtes par une sommation pondérée, introduisant de la flexibilité dans le mécanisme d'attention et débloquant un potentiel de performance supplémentaire. Des expériences approfondies sur ViT, DiT et LLMs démontrent que MoH surpasse l'attention multi-têtes en n'utilisant que 50 % à 90 % des têtes d'attention. De plus, nous montrons que des modèles d'attention multi-têtes pré-entraînés, tels que LLaMA3-8B, peuvent être encore affinés pour devenir nos modèles MoH. Notamment, MoH-LLaMA3-8B atteint une précision moyenne de 64,0 % sur 14 benchmarks, surpassant LLaMA3-8B de 2,4 % en utilisant seulement 75 % des têtes d'attention. Nous croyons que le MoH proposé est une alternative prometteuse à l'attention multi-têtes et fournit une base solide pour le développement de modèles d'attention avancés et efficaces.

English

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.