Mixture-of-Depths Attention

📄 arXiv: 2603.15619v1 📥 PDF

作者: Lianghui Zhu, Yuxin Fang, Bencheng Liao, Shijie Wang, Tianheng Cheng, Zilong Huang, Chen Chen, Lai Wei, Yutao Zeng, Ya Wang, Yi Lin, Yu Li, Xinggang Wang

分类: cs.CL, cs.AI

发布日期: 2026-03-16

备注: Code is released at https://github.com/hustvl/MoDA

🔗 代码/项目: GITHUB

💡 一句话要点

提出混合深度注意力机制(MoDA)，解决深度语言模型中的信号衰减问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 深度学习 大型语言模型 注意力机制 信号衰减 深度扩展

📋 核心要点

  1. 深度扩展是提升大型语言模型性能的关键，但深层模型易受信号衰减影响，浅层信息难以在深层恢复。
  2. MoDA允许注意力头关注当前层和先前层的KV对，从而保留并利用浅层特征，缓解信号衰减。
  3. 实验表明，MoDA在困惑度和下游任务性能上均优于基线，且计算开销小，是深度扩展的有效方法。

📝 摘要（中文）

本文提出了一种混合深度注意力（MoDA）机制，旨在解决大型语言模型（LLMs）深度扩展时出现的信号衰减问题。随着LLMs的加深，浅层形成的信息特征会被重复的残差更新逐渐稀释，导致深层难以恢复。MoDA允许每个注意力头关注当前层的序列KV对以及来自先前层的深度KV对。此外，本文还提出了一种硬件高效的MoDA算法，解决了非连续内存访问模式，在64K序列长度下实现了FlashAttention-2 97.3%的效率。在15亿参数模型上的实验表明，MoDA始终优于强大的基线模型。值得注意的是，MoDA在10个验证基准测试中平均困惑度降低了0.2，在10个下游任务中平均性能提高了2.11%，而计算开销仅增加了3.7%。研究还发现，将MoDA与后归一化（post-norm）结合使用比与预归一化（pre-norm）结合使用效果更好。这些结果表明，MoDA是一种有前景的深度扩展原语。

🔬 方法详解

问题定义：大型语言模型（LLMs）通过增加深度来提升性能，但随着模型深度的增加，浅层网络层学习到的信息特征会逐渐被重复的残差连接所稀释，导致深层网络难以有效利用这些信息，从而限制了模型的整体性能。现有方法难以有效解决这种信号衰减问题。

核心思路：MoDA的核心思路是让每个注意力头不仅关注当前层的键值对（KV pairs），还能选择性地关注来自前面若干层的KV对。通过这种方式，模型可以在深层网络中重新利用浅层网络学习到的信息，从而缓解信号衰减问题，提升模型性能。这种设计借鉴了混合专家模型（MoE）的思想，但应用于深度维度。

技术框架：MoDA可以集成到Transformer架构中，替代或补充原有的注意力机制。具体来说，对于每一层，MoDA首先计算当前层的注意力权重，然后根据这些权重，选择性地聚合当前层和前面若干层的KV对。最终的输出是这些聚合后的KV对的加权和。整个过程可以看作是在深度维度上进行注意力加权。

关键创新：MoDA的关键创新在于引入了深度维度的注意力机制。与传统的注意力机制只关注序列维度不同，MoDA同时关注序列维度和深度维度，从而能够更好地利用模型不同层的信息。此外，论文还提出了一种硬件高效的MoDA算法，解决了非连续内存访问问题，提高了计算效率。

关键设计：MoDA的关键设计包括：1) 如何选择需要关注的先前层。可以选择固定数量的先前层，也可以使用注意力机制动态选择。2) 如何融合来自不同层的KV对。可以使用简单的加权平均，也可以使用更复杂的融合方法。3) 硬件加速算法的设计，以解决非连续内存访问问题。论文中提到在64K序列长度下实现了FlashAttention-2 97.3%的效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在15亿参数模型上，MoDA在10个验证基准测试中平均困惑度降低了0.2，在10个下游任务中平均性能提高了2.11%，而计算开销仅增加了3.7%。此外，MoDA与后归一化结合使用比与预归一化结合使用效果更好。这些结果证明了MoDA的有效性和实用性。

🎯 应用场景

MoDA作为一种深度扩展技术，可以广泛应用于各种需要大型语言模型的场景，例如机器翻译、文本生成、对话系统、代码生成等。通过缓解深度模型中的信号衰减问题，MoDA可以提升这些应用的性能和效果，尤其是在需要处理长序列的任务中，其优势更为明显。未来，MoDA有望成为构建更大规模、更深层次语言模型的重要组成部分。

📄 摘要（原文）

Scaling depth is a key driver for large language models (LLMs). Yet, as LLMs become deeper, they often suffer from signal degradation: informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers. We introduce mixture-of-depths attention (MoDA), a mechanism that allows each attention head to attend to sequence KV pairs at the current layer and depth KV pairs from preceding layers. We further describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2's efficiency at a sequence length of 64K. Experiments on 1.5B-parameter models demonstrate that MoDA consistently outperforms strong baselines. Notably, it improves average perplexity by 0.2 across 10 validation benchmarks and increases average performance by 2.11% on 10 downstream tasks, with a negligible 3.7% FLOPs computational overhead. We also find that combining MoDA with post-norm yields better performance than using it with pre-norm. These results suggest that MoDA is a promising primitive for depth scaling. Code is released at https://github.com/hustvl/MoDA .