arXiv:2410.11842cs.CVcs.AI2024-10ICML被引 59

用专家混合机制优化注意力头,提速降耗还提效。

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

  • 将注意力头视为专家,按需动态选择激活
  • 仅用50%-90%头实现更高精度,75%头时提升2.4%
  • 可直接继续训练已有模型,适合高效部署

本文升级了Transformer核心的多头注意力机制,在保持或超越原有精度的前提下提升效率。我们发现多头注意力可表示为求和形式,基于‘并非所有注意力头同等重要’的洞察,提出混合头注意力(MoH):将注意力头类比为混合专家(MoE)中的专家。MoH具有两大优势:一是每个令牌可自主选择适配的注意力头,提升推理效率而不损失精度且不增加参数量;二是用加权求和替代标准求和,增强机制灵活性并释放性能潜力。在ViT、DiT和大语言模型上的广泛实验表明,MoH仅使用50%-90%的注意力头即可超越传统多头注意力。此外,预训练的多头模型如LLaMA3-8B可进一步持续微调为MoH模型。值得注意的是,MoH-LLaMA3-8B在14个基准上平均准确率达64.0%,较原版提升2.4%,仅使用75%的注意力头。我们认为MoH是多头注意力的有力替代方案,为高效注意力模型发展奠定基础。

原文摘要 · Abstract (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.

注意力机制模型压缩MoE高效模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。