Mixture-of-Depths (MoD) 与动态深度分配深度解析：从等计算范式到按 Token 自适应路由的下一代 Transformer 架构演进

Mixture-of-Depths (MoD) 与动态深度分配深度解析：从等计算范式到按 Token 自适应路由的下一代 Transformer 架构演进核心痛点：现行 Transformer 把 FLOPs 在序列维度与深度维度上"均匀涂抹"——每条序列的每个位置都必须经过全部e l l ellell层 Self-Attention 与 MLP，结果是"算力不足"与"算力浪费"同时存在：高频出现的停用词（标点、虚词、连接词）消耗掉与命名实体相同的算力，而真正需要深层推理的 token（数学符号、代码标识符、长程依赖锚点）反而被淹没在平均算力预算之下；Mixture-of-Depths (MoD) 通过"按 token 自适应路由到残差直通或层计算"在总计算预算固定的前提下，让每个 token 自学"该算多少层"，在 isoFLOP 下与 vanilla Transformer 性能持平甚至略优 1.5 个百分点，同时把单次前向 FLOPs 削减 50% 以上、推理采样 step 时延最高降低 50%——但它带来的副作用是非因果路由、自回归采样困难、batch 内负载对齐、Causal Mask 与动态深度耦合、设计空间远大于 MoE，这些工程难题正是 2024-2026 年"动态计算"研究最热的战场适配人群：负责 LLM 推理优化（KV Cache/Continuous Batching/Speculative Decoding）的平台工程师、研究 Conditional Computation 与 Adaptive Computation 的科研人员、阅读 Transformer 原始论文并尝试改造骨干架构的算法工程师、关注 Mixture-of-Experts 系列与 MoD/MoE/SA 三大动态架构差异化定位的架构师、对 LayerSkip/MoDification/MoSA 等 MoD 后续工作形成完整谱系认知的技术负责人收获能力：建立"等计算 vs 动态计算 vs 条件计算 vs 早退计算"四象限认知地图，深入理解 MoD 的 expert-choice 路由（C 个 token 参与层计算、B-C 个 token 经残差直通）、top-k 路由器数学（W r i n m a t h b b R d × 1 W_r in mathbb{R}^{d × 1}Wr​inmathbbRd×1投影到标量分数、按层预算k kk取 top）、对偶形式（token-choice 与 expert-choice 互为对偶）及其训练稳定性（aux-loss 与容量因子）、自回归采样的非因果性挑战（采样时只看历史 token）与三种解决方案（aux-loss、aux-MLP 预测器、Sequence-Level 路由），掌握 MoD vs MoE vs MoSA vs LayerSkip vs Skypick 的差异化定位，能在 PyTorch 中亲手实现一个完整的 MoD Transformer Block 并理解 Capacity Factor、Batch Augmentation、Sequence Packing 三大工程适配技巧技术背景与演进逻辑背景一：等计算（Uniform Compute）范式的算力浪费现象：vanilla Transformer 在每条序列每个位置都执行e l l c d o t ( 2 m a t h r m S e l f A t t n + 2 m a t h r m M L P ) ell cdot (2 mathrm{SelfAttn} + 2 mathrm{MLP})ellcdot(2mathrmSelfAttn+2mathrmMLP)次相同的浮点运算，不管 token 是高频停用词还是稀有关键字推演：当 batch 内包含大量"易于预测"的 token 时，深度模型在它们身上浪费的 FLOPs 与"难以预测"的关键 token 完全相等——这是结构性浪费而非调度低效子点一 - 算力分布实证：DeepMind MoD 论文实测显示，序列中约 12-25% 的 token 是"易预测停用词"（标点/虚词/低信息词），它们在浅层就足以建模，约 30-40% 是"中等难度"（实词/常见短语），约 35-55% 是"高难度"（命名实体/数字/代码标识符/逻辑锚点）需要深层特征子点二 - 算力预算守恒：训练总 FLOPs 约等于6 c d o t P c d o t D 6 cdot P cdot D6cdotPcdotD（P PP参数数、D DD数据 token 数）；MoD 的核心命题是"在6 P D 6PD6PD守恒的前提下，把算力从易 token 转移到难 token 即可让等计算模型变成精准计算模型"子点三 - 推理浪费：推理时 step 时延p r o p t o proptoproptoFLOPs，如果能用 50% 的 FLOPs 达到 100% 的 FLOPs 性能，每 token 时延就降一半——这是 MoD 推理加速的核心动机推演 - 等计算是 ML 早期的"省心设计"——一个统一计算图方便 CUDA kernel 优化、方便 batch、方便 distributed；但它不是物理或信息论的要求，"动态分配"是结构性更优的解法背景二：Conditional Computation 与 Mixture-of-Experts（MoE）现象：稀疏激活的 MoE（Switch Transformer、GShard、Mixtral）已经在"按 token 分配算力"上迈出第一步——通过 router 把 token 分发到 top-k 个专家，每个 token 走"小部分参数"推演 - MoE 解决"宽度方向的稀疏"，但深度方向仍是均匀的；如果在 MoE 之上再加一层"按 token 跳过部分层"的控制，就得到 MoD——两者是"维度正交、可叠加"的稀疏化子点一 - MoE 数学：router 输出g i n m a t h b b R E g in mathbb{R}^{E}ginmathbbRE（E EE专家数）、选 top-k 专家、y = s u m i i n m a t h r m t o p k ( g ) g i c d o t m a t h r m E x p e r t i ( x ) y = sum_{i in mathrm{topk}(g)} g_i cdot mathrm{Expert}_i(x)y=sumiinmathrmtopk(g)​gi​cdotmathrmExperti​(x)；激活参数占总参数的k / E k/Ek/E比例子点二 - MoE 限制：所有被选中的专家都要"等深度计算"（激活后还要走完整套 Self-Attn + MLP 链路），不能跳过整层 Transformer 块子点三 - 算力维度差异：MoE 削减"参数维度"的算力，MoD 削减"深度维度"的算力；二者在理论上完全正交、可叠加成 MoDE / MoMa 等混合架构推演 - 理解 MoD 必须建立在已理解 MoE 上；后续 MoMa（Mamba+MoE）、MoDE（MoE+MoD）等都是这条主线上的变体背景三：MoD 与 MoE 的本质区别——算力维度不同现象：MoD 不改变"激活 token 走的是哪一套参数"，而是改变"激活 token 是否需要走过这一层"——这是深度方向的选择，是 MoE 不具备的维度推演 - MoD 与 MoE 是"在算力分配的不同维度上的稀疏化"；理解这一点就能解释"为什么 MoD 的工程难度更高"——因为跳过整层意味着残差连接要承担"信息流"的全部职责，路由器决定"哪些 token 不进入这层"子点一 - 路由对象：MoE 把 token 路由到 top-k 专家（横切：宽度）；MoD 把 token 路由到 top-k 容量（纵切：深度）子点二 - 容量属性：MoE 的容量 = “每个专家最多处理多少 token”（token-level capacity）；MoD 的容量 = “每层最多处理多少 token”（layer-level capacity）子点三 - 决策粒度：MoE 决策粒度 = token × 专家（每个 token 选 k 个专家）；MoD 决策粒度 = token × 层（每层每个 token 是否参与）子点四 - 残差连接：MoE 路由后所有 token 仍然经过完整的层（attn + MLP）；MoD 的"未被选中的 token" 100% 跳过该层、仅走残差直通推演 - 在 token 维度上做"算力配额"的思想 MoD 与 MoE 完全一致（expert-choice），但在"跳过整层"这个动作上 MoD 是独一无二的——这也是后续 LayerSkip、Skypick 等"层级 early exit"工作的源头背景四：自回归采样的非因果挑战现象：MoD 的 expert-choice top-k 路由器在训练时看完整序列，但推理时只能看到历史 token——这与 Causal Mask 直接冲突推演 - MoD 的"非因果性"是工程难题的核心来源；三种解决方案（aux-loss 软约束、aux-MLP 预测器、Sequence-Level 静态路由）各有优缺点子点一 - 非因果来源：top-k 选择需要看完整序列的 router 分数才能"按分数排队选 top"，但推理时未来 token 不存在子点二 - 方案 1（Auxiliary Loss）：加一个"前文预测"的 aux loss 让路由器学会"基于前文选择"，但训练/推理分布有 gap子点三 - 方案 2（Aux-MLP Predictor）：训练一个小 MLP 预测"该 token 是否会被 top-k 选中"，推理时用预测代替路由子点四 - 方案 3（Sequence-Level Routing）：让整条序列的预算固定为C CC，所有 token 共享预算，推理时直接按"已知预算"分配，无需 top-k推演 - MoDification（NAACL 2025）的核心贡献就是"如何让 MoD 不依赖繁重的从头训练就能应用"——给出 aux-loss + 选择性 fine-tune 的简化路径背景五：动态深度与 Early Exit 的关系现象：LayerSkip（Meta 2024 ACL）、CALM（Google 2022）、Skypick 等"早退"工作与 MoD 在某些方面看起来相似——都是"不是每个 token 必须走完整深度"推演 - Early Exit 是"序列级"动态深度（一条序列决定退出层），MoD 是"token 级"动态深度（每个 token 走不同深度）——细粒度差异决定了 MoD 的灵活性子点一 - Early Exit 决策粒度：整条序列共享一个出口层；MoD 决策粒度：每个 token 可独立选子点二 - 信息流：Early Exit 后 token 间 attention 仍按"已退出的"层对齐；MoD 是同一层内"直通残差 vs 进层计算"两种状态并存子点三 - 与 Speculative Decoding 结合：LayerSkip 与 MoD 都可与 Self-Speculative Decoding（自推测解码）结合——早退层提供 draft、最终层 verify推演 - MoD 与 Early Exit 不是替代关系而是互补关系：MoD 控制"逐层哪些 token 进入"，Early Exit 控制"整条序列在哪层退出"，二者组合可获得"逐 token 逐层"双重动态深度总结：MoD 是 DeepMind 2024 年提出的"按 token 动态分配 Transformer 层深度"方法，它在算力预算固定（isoFLOP）前提下让重要 token 走完整深度、易 token 走残差直通，实现"等算力但更准"或"等准但省 50%+ FLOPs"的双向优化——它的工程挑战集中在 expert-choice 路由的因果性、容量因子的负载对齐、batch 内的 token 重分布这三个层面，是动态深度领域第一个落地完整的方案演进时间线（text 树）：动态深度与 MoD 演进时间线 ├── 1991 - Schmidhuber 早期条件计算：Adaptive Computation Time ├── 2017 - Transformer 确立等计算范式，每条序列每个位置都走 ℓ 层 ├── 2018 - Conditional Computation（Early Exit, PonderNet）萌芽 ├── 2022 - Google CALM: Confident Adaptive Language Modeling，序列级 early exit ├── 2023 - Switch Transformer / Mixtral: MoE 宽度稀疏，深度仍等计算 ├── 2024.04 - DeepMind Mixture-of-Depths 论文 (arXiv 2404.02258)，token 级动态深度 ├── 2024.07 - Meta LayerSkip (ACL 2024)，层 dropout 训练 + early exit 推理 + 自推测解码 ├── 2024.10 - MoDification (NAACL 2025)，aux-loss 简化 MoD，1.2× 延迟 1.8× 显存 ├── 2025.02 - Skypick: 随机层跳过 + 推理时按层预算分配 ├── 2025.04 - MoE + MoD 混合 (MoDE)，按 token 同时路由专家与深度 ├── 2025.06 - MoSA (arXiv 2505.00315)，Sparse Attention 用 expert-choice 路由 ├── 2025.09 - DeepSeek-V3.5 实验 MoD 模式长上下文推理加速 ├── 2026.01 - CALM-MoD：MoD 与 confident exit 结合，token + 序列双重动态 ├── 2026.04 - 生产级 MoD 在 Llama-3 变体实验：长上下文推理 1.4× 加速 ├── 2026.07 - NVIDIA FasterTransformer 接入 MoD 模式作为可选 kernel 调度策略核心原理深度解析MoD 的数学建模：Layer-Level Capacity 与 Top-K Routing本质：每层 Transformer Block 在处理 token 时有一个"层预算"C CC（Layer Capacity），由路由器选出 top-k =C CC个 token 进入完整计算（Attn + MLP），其它L − C L - CL−C个 token 仅走残差直通（Skip Connection）设模型e l l ellell层、batchB BB条序列、每序列长度L LL