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

If you like our project, please give us a star ⭐ on GitHub for the latest update.

💡 I also have other projects that may interest you ✨.

Chat-UniVi: Unified Visual Representation Empowers Large Language Models with Image and Video Understanding
Peng Jin, Ryuichi Takanobu, Wancai Zhang, Xiaochun Cao, Li Yuan

MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts
Peng Jin, Bo Zhu, Li Yuan, Shuicheng Yan

-->

📣 News

• [2024/10/23] We updated tokenizer config for MoH-LLaMA3-8B.
• [2024/10/22] Now MoH-LLaMA3-8B is available.
• [2024/10/10] MoH-LLaMA3-8B weights are being approved and will be available for download after approval.
• [2024/10/09] Model weight and inference code are available now! Welcome to watch 👀 this repository for the latest updates.

⚡ Overview

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.

😮 Highlights

💡 General Framework

We evaluate our proposed MoH across various popular model frameworks, including Vision Transformers (ViT) for image classification, Diffusion models with Transformers (DiT) for class-conditional image generation, and Large Language Models (LLMs) for language tasks.

🔥 High Performance

Extensive experiments on ViT, DiT, and LLMs demonstrate that MoH outperforms multi-head attention by using only 50%~90% of the attention heads.

🤗 Support Continue-Tuning Starting from the Multi-Head Attention Models

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.

The MoH model quickly recovers to over 95% of the performance of the original model within a training budget of 10B tokens. Then, the performance gradually improves with the increase of the training tokens.

🚀 Main Results

ViT for ImageNet-1K Classification

DiT for Class-Conditional Image Generation (ImageNet-1K)

Training LLMs from Scratch

Continue-Tuning LLaMA3-8B

😍 Why is MoH better than Multi-Head Attention?

Flexible Head Assignment Patterns

We observe significant variation in attention head assignments across different categories and task topics, indicating that the MoH model adapts to diverse tasks by employing distinct head assignment patterns. This characteristic of MoH allows different attention heads to focus on different types of tasks, making parameter utilization more efficient than multi-head attention.

Weighted Summation of Heads

By replacing the standard summation in multi-head attention with a weighted summation, MoH enhances the flexibility of the attention mechanism and increases the performance potential.

🗝️ Training & Validating

• The training instruction of MoH-ViT is in MoH-ViT.
• The training instruction of MoH-DiT is in MoH-DiT.
• The instruction of MoH-LLaMA-8B is in MoH-LLaMA-8B.

👍 Acknowledgement

• Skywork-MoE It is an advanced MoE language model.
• LLaMA3 It is an advanced open-source language model.
• TransNeXt (CVPR 2024) It is an advanced ViT model.
• DiT (ICCV 2023) It is an advanced diffusion model with the Transformer.

🤝 Related Projects

• MoE++ MoE++ achieves better performance while delivering 1.1~2.1x expert forward throughput compared to a vanilla MoE model of the same size, which lays a solid foundation for developing advanced and efficient MoE-related models.
• Chat-UniVi (CVPR 2024 Highlight) The model is an efficient large language and video assistant. This framework exhibits remarkable interactive capabilities between images and videos.

🔒 License

• The majority of this project is released under the Apache 2.0 license as found in the LICENSE file.
• The service is a research preview intended for non-commercial use only, subject to the model License of LLaMA, Terms of Use of the data generated by OpenAI, and Privacy Practices of ShareGPT. Please contact us if you find any potential violations.

✏️ Citation

If you find this paper useful, please consider staring 🌟 this repo and citing 📑 our paper:

@article{jin2024moh,
  title={MoH: Multi-Head Attention as Mixture-of-Head Attention},
  author={Jin, Peng and Zhu, Bo and Yuan, Li and Yan, Shuicheng},
  journal={arXiv preprint arXiv:2410.11842},
  year={2024}
}

✨ Contributors