MoH: Attenzione Multi-Head come Attenzione a Mischia di Teste

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

October 15, 2024

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

cs.AI

Abstract

In questo lavoro, aggiorniamo il meccanismo di attenzione multi-testa, il nucleo del modello Transformer, per migliorare l'efficienza mantenendo o superando il livello di precisione precedente. Mostriamo che l'attenzione multi-testa può essere espressa nella forma di sommatoria. Basandoci sull'idea che non tutte le teste di attenzione hanno la stessa importanza, proponiamo l'attenzione Mixture-of-Head (MoH), una nuova architettura che tratta le teste di attenzione come esperti nel meccanismo Mixture-of-Experts (MoE). MoH ha due significativi vantaggi: in primo luogo, MoH consente a ciascun token di selezionare le teste di attenzione appropriate, migliorando l'efficienza inferenziale senza compromettere la precisione o aumentare il numero di parametri. In secondo luogo, MoH sostituisce la sommatoria standard nell'attenzione multi-testa con una sommatoria pesata, introducendo flessibilità nel meccanismo di attenzione e sbloccando un potenziale di prestazioni aggiuntivo. Esperimenti approfonditi su ViT, DiT e LLMs dimostrano che MoH supera l'attenzione multi-testa utilizzando solo il 50%-90% delle teste di attenzione. Inoltre, dimostriamo che i modelli di attenzione multi-testa preaddestrati, come LLaMA3-8B, possono essere ulteriormente adattati ai nostri modelli MoH. In particolare, MoH-LLaMA3-8B raggiunge una precisione media del 64,0% su 14 benchmark, superando LLaMA3-8B del 2,4% utilizzando solo il 75% delle teste di attenzione. Riteniamo che il MoH proposto sia una promettente alternativa all'attenzione multi-testa e fornisca una solida base per lo sviluppo di modelli avanzati ed efficienti basati sull'attenzione.

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.