MoH：将多头注意力视为混合注意力头

在本工作中，我们对Transformer模型的核心机制——多头注意力机制进行升级，以在保持或超越先前准确性水平的同时提高效率。我们表明，多头注意力可以表示为求和形式。基于并非所有注意力头都等重要这一洞察，我们提出了混合注意力头(Mixture-of-Head attention, MoH)，一种新的架构将注意力头视为混合专家(MoE)机制中的专家。MoH具有两个显著优势：首先，MoH使每个token能够选择合适的注意力头，在不损害准确性或增加参数数量的前提下提高推理效率；其次，MoH将多头注意力中的标准求和替换为加权求和，为注意力机制引入灵活性，释放额外的性能潜力。在ViT、DiT和LLM上进行的大规模实验表明，MoH通过仅使用50%~90%的注意力头即可超越多头注意力。此外，我们展示了预训练的多头注意力模型（如LLaMA3-8B）可以进一步微调为我们的MoH模型。值得注意的是，MoH-LLaMA3-8B在14项基准测试中实现平均64.0%的准确率，凭借仅使用75%的注意力头超越LLaMA3-8B 2.4%。我们认为，提出的MoH是多头注意力的有前景的替代方案，为开发高级且高效的基于注意力的模型提供了坚实基础。