让注意力头跨层看信息，解决深层模型信号衰减问题
Mixture-of-Depths Attention
- 每个注意力头可访问当前层及前序层的键值对
- 1.5B模型在10个基准上平均困惑度降0.2，下游任务提升2.11%
- 仅增3.7%计算开销，适配大模型深度扩展需求
扩大模型深度是大语言模型发展的关键。然而，随着模型加深，浅层形成的有用特征会因反复残差更新而逐渐稀释，难以在深层恢复。本文提出混合深度注意力（MoDA），使每个注意力头不仅能关注当前层的键值对，还能访问前序层的键值对。我们进一步设计了一种硬件高效的算法，解决了非连续内存访问问题，在64K序列长度下达到FlashAttention-2 97.3%的效率。在15亿参数模型上的实验表明，MoDA持续优于强基线：在10个验证基准上平均困惑度降低0.2，在10个下游任务上平均性能提升2.11%，计算开销仅增加3.7% FLOPs。此外，将MoDA与后归一化结合比与前归一化效果更好。结果表明，MoDA是一种有前景的深度扩展基础组件。代码已开源。
原文摘要 · Abstract (English)
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 .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。