Mixture-of-Depths Attention

扩展模型深度是大语言模型（LLM）性能提升的关键驱动力。然而，随着LLM不断加深，往往会出现信号退化问题：浅层所提取的富含信息的特征，在经过多轮残差更新后逐渐被稀释，导致其在深层中难以有效恢复。为此，我们提出“混合深度注意力机制”（Mixture-of-Depths Attention, MoDA），该机制允许每个注意力头在当前层对序列的键值（KV）对进行关注的同时，还能跨层访问前序各层所生成的深度KV对。我们进一步设计了一种面向硬件优化的MoDA实现算法，专门解决非连续内存访问模式带来的开销问题，在序列长度达64K时，其运行效率可达FlashAttention-2的97.3%。在15亿参数规模模型上的实验表明，MoDA始终优于多个强基线方法：在10个验证基准数据集上，其平均困惑度（perplexity）降低0.2；在10项下游任务上，平均性能提升2.11%，而计算开销仅增加微乎其微的3.7% FLOPs。此外，我们还发现，将MoDA与后归一化（post-norm）结构结合，效果优于其与前归一化（pre-norm）的组合。上述结果表明，MoDA是一种极具潜力的、支撑模型深度持续扩展的基础性机制。代码已开源：https://github.com/hustvl/MoDA。

随着LLM深度增加，浅层提取的有用特征在深层中因重复残差更新而逐渐退化（signal degradation），导致深度扩展遭遇性能瓶颈。这是一个在大模型缩放实践中日益突出但尚未被系统性解决的新问题。

提出混合深度注意力（MoDA）：每个注意力头可动态融合当前层的KV与来自先前若干层的‘深度KV’（depth KV），显式保留跨层语义信息；不同于传统单层KV访问或参数共享机制，MoDA通过可学习门控实现层间KV选择，且设计了硬件友好的非连续内存访问优化算法（基于FlashAttention-2改进），首次在长序列下高效实现跨层KV复用。

在1.5B模型上验证：平均困惑度下降0.2（10个验证集）、下游任务平均提升2.11%（10个基准），仅增3.7% FLOPs；发现MoDA与post-norm组合优于pre-norm；提供完整开源代码（https://github.com/hustvl/MoDA）；实验覆盖64K长序列，验证硬件效率达FlashAttention-2的97.3%；值得深入的方向包括MoDA的动态深度选择机制、与稀疏注意力/专家混合架构的协同设计。

DeepNet: Scaling Transformers to 1,000 Layers (ICLR 2022); Stable Diffusion with Depth-Aware Attention (NeurIPS 2023); LayerSkip: Learning to Skip Layers in Deep Networks (ICML 2023); FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023); RetNet: Multi-Scale Retention for Large Language Models (ICLR 2024)