Mixture of Depths (MoD)

Google DeepMind 于 2024 年提出的动态计算架构，让 Transformer 模型学会为不同 token 分配不同的计算深度，实现与 Mixture of Experts 正交的稀疏化维度。

概述

标准 Transformer 对所有输入 token 均匀地分配计算资源：每个 token 都必须经过全部 $L$ 层的 Self-Attention 和 MLP 计算。Mixture of Depths (MoD) 挑战了这一均匀假设——并非所有 token 都需要同等深度的处理。MoD 通过 top-k 路由机制，让模型动态决定每个 token 在每一层是否参与计算，从而在保持性能的同时显著减少前向传播的计算量。

核心机制

问题设定

传统 Transformer 的计算分布是均匀的：

Token A (简单词)  → Layer 1 → Layer 2 → ... → Layer L
Token B (复杂概念) → Layer 1 → Layer 2 → ... → Layer L

MoD 的目标是让计算分布动态化：

Token A (简单词)  → Layer 1 → [跳过] → [跳过] → Layer L
Token B (复杂概念) → Layer 1 → Layer 2 → Layer 3 → ... → Layer L

top-k 路由

MoD 的核心是一个每层独立的 top-k 路由机制：

路由决策: 每层接收全部 token，但只选择 $k$ 个 token 进入该层的 Self-Attention 和 MLP 计算
选择标准: 由网络自身学习的路由分数决定哪些 token 最需要该层的处理
静态图优势: $k$ 是预先定义的常数，因此计算图是静态的（张量大小已知），与 MoE 的动态路由不同
流体身份: 虽然 $k$ 固定，但具体哪些 token 被选中是上下文相关的

关键特性

总计算量可预测: 每层固定处理 $k$ 个 token，总 FLOPs 是确定的
token 级动态: 单个 token 的计算路径是上下文敏感的
与 MoE 正交: MoD 在深度维度稀疏化，MoE 在宽度维度稀疏化，两者可结合为 MoDE

性能与效率

实验结果

根据 DeepMind 的原始论文：

等 FLOPs 对比: MoD 模型与基线性能相当
推理加速: 前向传播 FLOPs 减少约 50%，采样速度提升可达 50%+
训练效率: wall-clock 训练时间与基线相当

与相关技术的对比

技术	稀疏维度	机制	代表工作
MoD	深度 (层)	top-k 路由选择 token	DeepMind, 2024
MoE	宽度 (专家)	路由到不同 FFN 专家	Shazeer et al., 2017
Early Exit	深度	基于置信度提前退出	DeeBERT, 2020
CALM	深度	自适应计算预算分配	Schuster et al., 2022

MoD 与 Early Exit / CALM 的关键区别：

Early Exit 是确定性的（置信度阈值），MoD 是学习出来的路由策略
MoD 的路由决策是每层独立的，token 可以在不同层被选中或跳过

与 MoE 的结合: MoDE

MoD 和 MoE 是正交的稀疏化维度：

MoD: 决定哪些 token 进入当前层
MoE: 决定 token 进入当前层的哪个专家

两者结合形成 MoDE (Mixture of Depths and Experts)，同时在深度和宽度两个维度实现稀疏计算，理论上可以获得乘积级的效率提升。

局限与挑战

自回归复杂性: 在解码阶段，动态路由与 KV Cache 的交互需要精心设计
内存访问模式: 虽然总计算量减少，但不规则的内存访问可能影响实际加速比
负载均衡: 需要辅助损失函数确保各层的计算负载均衡
长序列扩展: 极长序列下 top-k 选择的语义意义需要验证

关键论文

论文	作者	年份	链接
Mixture-of-Depths: Dynamically allocating compute in transformer-based language models	Raposo et al. (Google DeepMind)	2024	arXiv

Mixture of Depths (MoD) ​

概述 ​

核心机制 ​

问题设定 ​

top-k 路由 ​

关键特性 ​

性能与效率 ​

实验结果 ​

与相关技术的对比 ​

与 MoE 的结合: MoDE ​

局限与挑战 ​

关键论文 ​

相关概念 ​