MoH: Atención Multi-Cabeza como Atención Mezcla-de-Cabezas

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

October 15, 2024

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

cs.AI

Resumen

En este trabajo, mejoramos el mecanismo de atención multi-cabeza, el núcleo del modelo Transformer, para aumentar la eficiencia manteniendo o superando el nivel de precisión previo. Mostramos que la atención multi-cabeza puede ser expresada en forma de suma. Basándonos en la idea de que no todas las cabezas de atención tienen la misma importancia, proponemos la atención de Mezcla de Cabezas (MoH), una nueva arquitectura que trata las cabezas de atención como expertos en el mecanismo de Mezcla de Expertos (MoE). MoH tiene dos ventajas significativas: Primero, MoH permite que cada token seleccione las cabezas de atención apropiadas, mejorando la eficiencia de inferencia sin comprometer la precisión o aumentar el número de parámetros. Segundo, MoH reemplaza la suma estándar en la atención multi-cabeza con una suma ponderada, introduciendo flexibilidad al mecanismo de atención y desbloqueando un potencial de rendimiento adicional. Experimentos exhaustivos en ViT, DiT y LLMs demuestran que MoH supera a la atención multi-cabeza utilizando solo el 50%-90% de las cabezas de atención. Además, demostramos que los modelos pre-entrenados de atención multi-cabeza, como LLaMA3-8B, pueden ser afinados aún más en nuestros modelos MoH. Notablemente, MoH-LLaMA3-8B logra una precisión promedio del 64.0% en 14 pruebas, superando a LLaMA3-8B en un 2.4% al utilizar solo el 75% de las cabezas de atención. Creemos que el MoH propuesto es una alternativa prometedora a la atención multi-cabeza y proporciona una base sólida para el desarrollo de modelos avanzados y eficientes basados en atención.

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.