[2026 技术前沿] MoDA:打破深度瓶颈,让 LLM 能够“翻看”每一层的记忆
Mixture-of-Depths Attention
本文提出了 Mixture-of-Depths Attention (MoDA),一种旨在解决深度扩展(Depth Scaling)中信息稀释问题的统一注意力机制。通过让每个注意力头同时连接序列 KV 与历史层(深度)的 KV,MoDA 在不增加 DenseNet 式计算负担的前提下实现了跨层信息检索。实验显示,在 1.5B 参数模型上,MoDA 仅以 3.7% 的额外 FLOPs 开销,将 10 项下游任务的平均性能提升了 2.11%。
TL;DR
深度本应是 LLM 表达能力的核心,但“信号稀释”让加深模型变得收效甚微。来自字节跳动和华中科技大学的研究者提出了 Mixture-of-Depths Attention (MoDA)。它打破了传统 ResNet 式的层级传递逻辑,允许注意力头直接检索之前所有层的隐藏状态。该方案不仅在 1.5B 模型上刷榜,更通过巧妙的算子融合达到了 FlashAttention-2 级别的极限硬件效率。
痛点深挖:深度扩展的“天花板”
目前 LLM 的 Scaling 往往侧重于数据量、宽度和上下文长度。虽然理论上更深的模型能学到更复杂的逻辑,但现实中我们面临两大阻碍:
- 信号稀释 (Information Dilution):在数百层的残差网络中,浅层捕捉到的关键特征会被不断的 Residual Updates 冲淡,深层想找回这些“初心”变得异常困难。
- 效率陷阱:如果模仿 DenseNet 让所有层全连接,计算量会随层数 的平方级增长 (),这对于大模型是不可接受的。
核心机制:MoDA 的“深度检索”直觉
作者将模型每一层的操作拆解为 Read - Operate - Write 流程。
- Read (读):不再仅仅读取上一层的输出(Identity Read),而是发起一个针对历史层 KV 缓存的查询。
- Operate (算):由一个统一的 Softmax 处理“序列 KV”和“深度 KV”。这意味着模型可以动态决定:此时此刻是该关注序列中的某个 token(空间信息),还是该关注此 token 在前几个层级中的表达特征(深度历史)。
- Write (写):将当前层的计算结果追加到深度 KV 流中,供后续层使用。
图注:(a) 残差连接, (b) 密集连接, (c) 深度注意力, (d) 本文提出的 MoDA 将两者融合。
硬件加速:97.3% 的 FlashAttention 效率是如何炼成的?
MoDA 最具工业价值的地方在于其硬件感知实现(Hardware-aware implementation)。原生 PyTorch 实现因为需要频繁处理非连续的内存块,速度极慢。
作者通过以下三招实现了突围:
- Chunk-aware Layout:将查询划分为块(Chunk),让每一块仅访问对应的深度 KV 区间,将原本 的深度利用率极大地拉升。
- Group-aware Indexing:利用 GQA(分组查询注意力)的特性,让 G 个 Query 头复用同一组 Depth KV 的索引计算,大幅减少内存寻址。
- 算子融合:将序列注意力和深度注意力整合在同一个 Online-Softmax 循环中,减少了 HBM(访存)压力。
图注:随着序列长度增加,MoDA 带来的额外耗时被摊薄,在大规模推理中极具竞争力。
实验结果:全线 SOTA 提升
在 1.5B 模型、400B Token 的严苛测试下,MoDA 的表现令人印象深刻:
- 常识推理:WinoGrande 提升 2.37%,COPA 提升 4.00%。
- 多任务能力:MMLU (大规模多任务语言理解) 提升 1.86%。
- 可视化洞察:通过注意力热图发现,MoDA 会自发地将分配给“Attention Sink”(无意义的填充占位)的概率权重转移到有意义的深度层级中。
图注:观察红虚线右侧,深层模型明显在积极检索来自历史层的“深度信息”。
总结与未来:超深模型的基石
MoDA 为模型架构设计提供了一个重要启示:深度维度的信息不应是被动叠加的,而应该是主动检索的。
局限性与挑战: 虽然计算开销很低,但缓存所有层的 KV 会占用额外的内存(KV Cache)。作者在讨论中也诚实地指出,未来可以通过**固定插槽缓存(Fixed-size Depth-KV Slot Caching)**来解决这个 Limit,即只保留最有价值的 个历史层级。
对于追求极致 Scaling 的开发者,MoDA 提供了一个无需大幅改动算子框架就能榨取模型深度潜力的优雅路径。
