MoH: Многоголовое внимание как смесь многоголового внимания

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

October 15, 2024

Авторы: Peng Jin, Bo Zhu, Li Yuan, Shuicheng Yan

cs.AI

Аннотация

В данной работе мы улучшаем механизм внимания с несколькими головами, являющийся основой модели Transformer, с целью повышения эффективности при сохранении или превосходстве предыдущего уровня точности. Мы показываем, что механизм внимания с несколькими головами можно выразить в форме суммирования. Основываясь на понимании того, что не все головы внимания имеют одинаковое значение, мы предлагаем механизм внимания Mixture-of-Head (MoH), новую архитектуру, которая рассматривает головы внимания как экспертов в механизме Mixture-of-Experts (MoE). MoH имеет два значительных преимущества: Во-первых, MoH позволяет каждому токену выбирать соответствующие головы внимания, улучшая эффективность вывода без ущерба точности или увеличения количества параметров. Во-вторых, MoH заменяет стандартное суммирование в механизме внимания с несколькими головами на взвешенное суммирование, вводя гибкость в механизм внимания и раскрывая дополнительный потенциал производительности. Обширные эксперименты на моделях ViT, DiT и LLMs показывают, что MoH превосходит механизм внимания с несколькими головами, используя всего 50%-90% голов внимания. Более того, мы демонстрируем, что предварительно обученные модели механизма внимания с несколькими головами, такие как LLaMA3-8B, могут быть дополнительно донастроены наши MoH модели. Заметно, что MoH-LLaMA3-8B достигает средней точности 64,0% по 14 показателям, превосходя LLaMA3-8B на 2,4% при использовании всего 75% голов внимания. Мы считаем, что предложенный MoH является многообещающей альтернативой механизму внимания с несколькими головами и обеспечивает прочное основание для разработки передовых и эффективных моделей на основе внимания.

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.