鹏城实验室提出MoH | ViT、DiT和LLM上MoH只使用50%左右Head即可完成超越

文章预览

关注 「 AIWalker 」 并 星标 从此AI不迷路 来源 | AI落地之芯     链接 | AI芯落 在这项工作中，作者将Transformer模型的核心多头自注意力机制升级，以提高效率，同时保持或超过先前的准确度水平。作者证明了多头自注意力可以表示为求和形式。借鉴到注意力头的重要性并不相同这一洞察，作者提出了一种新的架构Mixture-of-Head attention (MoH)，将注意力头视为Mixture-of-Experts (MoE)机制中的专家。 MoH具有两个显著的优势：首先，MoH使每个 Token 可以选择合适的注意力头，在不降低准确率或增加参数数量的情况下，提高推理效率。其次，MoH用加权求和替换多头自注意力中的标准求和，为注意力机制引入灵活性，并解锁额外的性能潜力。在ViT、DiT和LLMs上的广泛实验表明，MoH通过只使用50%~90%的注意力头，优于多头自注意力。此外，作者证明了预训练的多头自注意力 ………………………………