MoH: Multi-Head Aandacht als Mixture-of-Head Aandacht

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

October 15, 2024

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

cs.AI

Samenvatting

In dit werk upgraden we het multi-head aandachtsmechanisme, de kern van het Transformer model, om de efficiëntie te verbeteren terwijl we het vorige nauwkeurigheidsniveau behouden of overtreffen. We tonen aan dat multi-head aandacht kan worden uitgedrukt in de vorm van een sommatie. Puttend uit het inzicht dat niet alle aandachtskoppen gelijkwaardig zijn, stellen we Mixture-of-Head aandacht (MoH) voor, een nieuwe architectuur die aandachtskoppen behandelt als experts in het Mixture-of-Experts (MoE) mechanisme. MoH heeft twee significante voordelen: Ten eerste stelt MoH elk token in staat om de juiste aandachtskoppen te selecteren, wat de inferentie-efficiëntie verbetert zonder de nauwkeurigheid in gevaar te brengen of het aantal parameters te verhogen. Ten tweede vervangt MoH de standaard sommatie in multi-head aandacht door een gewogen sommatie, wat flexibiliteit aan het aandachtsmechanisme toevoegt en extra prestatiepotentieel ontsluit. Uitgebreide experimenten op ViT, DiT en LLMs tonen aan dat MoH multi-head aandacht overtreft door slechts 50%-90% van de aandachtskoppen te gebruiken. Bovendien tonen we aan dat vooraf getrainde multi-head aandachtmodellen, zoals LLaMA3-8B, verder kunnen worden afgestemd op onze MoH-modellen. Opmerkelijk genoeg behaalt MoH-LLaMA3-8B een gemiddelde nauwkeurigheid van 64,0% over 14 benchmarks, waarbij LLaMA3-8B met 2,4% wordt overtroffen door slechts 75% van de aandachtskoppen te gebruiken. We geloven dat het voorgestelde MoH een veelbelovend alternatief is voor multi-head aandacht en een sterke basis biedt voor de ontwikkeling van geavanceerde en efficiënte op aandacht gebaseerde modellen.

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.