一个 32 层 Transformer 在处理一段文本时,最整齐的执行方式是让每个 Token 依次经过全部 32 个 Block。无论当前位置只是标点、常见连接词,还是承载了代码变量、数学关系或长距离依赖,它获得的网络深度都一样。这样的计算图对 GPU 很友好:张量形状稳定、每层工作量规则、批处理容易组织。但它也隐含了一个强假设——每个 Token 在每一层都值得支付同样的 Attention 和 MLP 成本。
Mixture-of-Depths,简称 MoD,尝试放松这个假设。它不先判断“某个 Token 总共应该走几层”,而是在若干 Transformer Block 前加入轻量路由,让当前层只处理容量范围内的一部分 Token。没有被选中的 Token 不执行该 Block 的主要计算,而是沿残差路径继续向后流动。于是模型总深度保持不变,但不同 Token 在不同层实际消耗的计算预算可以不同。
固定深度 Transformer 把计算预算平均分给了所有 Token
标准 Transformer Block 通常包含 Attention、前馈网络以及残差连接等组件。对于长度为 N 的序列,常规执行会让 N 个 Token 的隐藏状态一起进入这一层。下一层仍然接收 N 个 Token,只是隐藏表示已经更新。这个模式从第一层持续到最后一层,因此单个 Token 的“计算深度”基本由模型层数决定,而不是由当前输入状态决定。
这种统一路径有明显工程优势。训练时可以使用规则的矩阵计算;推理服务可以按照固定模型结构规划显存、kernel 和批处理;编译器也更容易优化静态计算图。但从计算分配角度看,它没有机制区分“当前层继续变换这个 Token 是否值得”。即使某个位置的表示在这一层变化很小,也会照常进入 Attention 和 MLP。
MoD 的目标不是判断哪些词在人类语义上“简单”或“复杂”。Router 学到的是一种与训练目标相关的计算分配策略。某些标点可能被频繁跳过,也可能在特定上下文里仍被选中;某些看起来复杂的词也不保证每层都获得计算。把 Router 解释成“理解了哪些 Token 更难”会过度拟人化。更准确的说法是:模型学习了哪些隐藏状态在当前层更值得占用有限计算容量。
MoD 的核心不是提前退出,而是逐层选择计算对象
MoD Block 为 Token 计算一个路由分数,然后在固定容量内选择一部分 Token 执行该层变换。假设当前序列有 4096 个 Token,某层容量比例设为 50%,那么这层最多让 2048 个位置进入主要计算。Router 决定的是“哪 2048 个”,而不是临时把容量从 2048 改成 137 或 3900。
没有被选中的 Token 会绕过当前 Block 的主要计算,并通过残差路径保留已有表示。到了下一层,它仍然会重新参与路由,因此“这一层跳过”不意味着“从此退出模型”。一个 Token 可以在第 10 层跳过,在第 11 层被选中,第 12 层再次跳过。这与传统 Early Exit 的执行边界不同。
flowchart TD
A[当前层全部 Token] --> B[Router 计算分数]
B --> C[按固定 Capacity 选 Top-K]
C --> D[Gather 被选 Token]
D --> E[执行 Attention 和 MLP]
E --> F[Scatter 回原位置]
B --> G[未选 Token 走旁路]
F --> H[合并隐藏状态]
G --> H
H --> I[进入下一层重新路由]
图里的 Gather 和 Scatter 很重要。算法描述可以只说“Top-K Token 进入 Block”,GPU 却需要把分散在序列不同位置的隐藏状态重新组织成适合矩阵计算的连续批次,计算完成后再写回原位置。MoD 减少的是 Block 内部的大计算,但 Router、排序、Gather、Scatter 和状态合并都会新增成本。
固定 Capacity 是动态计算能够落到硬件上的关键约束
完全自由的动态深度听起来更灵活:某层可以根据输入决定处理 10%、40% 或 95% 的 Token。但这种设计会让每层实际张量尺寸持续变化,训练中的工作量难以预测,批处理里的不同样本也可能产生完全不同的计算形状。GPU 依赖大规模、规则的并行工作,这种不确定性很容易抵消“少算”的理论收益。
MoD 采用固定 Capacity,把“计算预算”和“计算对象”拆开。系统提前知道每个 MoD Block 最多处理多少 Token,因此可以预估 FLOPs 和中间张量规模。Router 只负责在这个预算内选择位置。这种设计让动态路由没有彻底破坏规则计算图,也使训练成本能够被比较稳定地控制。
容量本身就是一个质量与成本旋钮。容量高时,更多 Token 进入 Block,行为更接近 Dense Transformer;容量低时,Block FLOPs 下降更明显,但 Router 必须在更激烈的竞争中决定哪些位置保留计算。容量不能脱离模型训练过程随意缩小。一个按特定 Capacity 训练出来的模型,不应假设推理时把比例从 50% 改成 10% 仍能维持同样行为。
MoD 与 MoE、Early Exit 分别稀疏了什么
MoD 容易和 MoE 混淆,因为两者都有 Router,也都会说“只计算一部分”。差别在于稀疏维度不同。MoE 中,一个 Token 到达 MoE 层后通常仍需要计算,只是 Router 决定它进入哪些 Expert;MoD 则先决定这个 Token 是否执行当前 Block 的主要计算。
Early Exit 又是另一种边界。它通常在某个位置判断当前表示已经足够用于输出,从而跳过后续一整段网络。MoD 的跳过是局部的:Token 可以绕过一层,再在后面的层重新进入计算。LayerSkip 等方法则把 Early Exit 与自推测解码结合,通过让较浅层具备可用预测能力来减少完整模型的重复执行,优化对象和 MoD 仍然不同。
| 技术 | 动态选择对象 | 被跳过后是否还能在后续层重新计算 | 主要节省位置 | 主要工程代价 |
|---|---|---|---|---|
| Dense Transformer | 不动态选择 | 不适用 | 无动态节省 | 计算规则、实现简单 |
| MoD | 当前层哪些 Token 执行 Block | 可以 | Attention / MLP 的 Token 计算量 | Router、Top-K、Gather/Scatter、不规则数据流 |
| MoE | 当前 Token 使用哪些 Expert | 通常仍经过该层 | 参数激活量 | Expert Parallelism、All-to-All、负载均衡 |
| Early Exit | 是否结束后续层计算 | 通常不会重新进入 | 后续整段网络 | 退出判定、质量校准、批内不同深度 |
| LayerSkip 类路线 | 较浅层预测并用后续层验证 | 按具体算法处理 | 完整解码轮次或后续层计算 | 训练目标、验证路径和缓存协调 |
MoD 与 MoE 在概念上可以组合:先决定“这一层是否处理当前 Token”,再决定“如果处理,使用哪些专家”。但两套 Router 叠加后,训练稳定性、Token 分桶、负载均衡和运行时调度都会更复杂。算法上的稀疏可以相乘,系统开销也可能相乘,因此不能仅凭 FLOPs 推断端到端收益。
训练时 Top-K 很自然,逐 Token 推理却有因果约束
训练阶段通常拿到完整序列,因此一层可以先计算整段 Token 的路由分数,再选择固定数量的 Top-K。这种选择天然利用了“同一段序列里谁的分数更高”这一全局比较。但在线自回归生成每次只新增一个 Token,未来 Token 还不存在。如果推理逻辑要求知道整段未来序列的分数才能决定当前 Token 是否进入某层,就违反了生成过程的因果约束。
因此,不能简单认为训练时的 Sequence Top-K 路由可以原封不动地用于逐 Token decode。MoD 原始工作专门讨论了采样阶段的路由处理,核心问题就是把训练阶段可用的固定容量决策转换成在线生成能够执行的选择机制。工程实现需要验证实际推理路径如何产生路由决定,而不能只看到论文里的 Top-K 公式就假设 decode 端也拥有完整序列。
这也会影响性能解释。训练或长 Prompt prefill 一次处理很多 Token,Router 可以形成较大的候选集合,Gather 后仍有机会得到适合 GPU 的矩阵规模。逐 Token decode 的每个请求一次只有很少新位置,高并发服务虽然可以跨请求组成 batch,但不同请求的层路由选择会产生额外不规则性。MoD 在训练 FLOPs 上成立的节省,不会自动以相同比例转换成在线每 Token 延迟。
少算 FLOPs 为什么不一定让 GPU 同比例变快
MoD 论文的核心价值之一,是展示了在固定计算预算下动态分配深度的可行性,并报告了部分配置在采样速度上的明显收益。但“少执行多少理论 FLOPs”和“服务快多少”是两个不同指标。Router 本身要做打分和选择;Top-K 需要排序或选择操作;Token 要重新 Gather 成紧凑张量;执行后还要 Scatter;跳过路径和计算路径最终还要恢复原序列顺序。
如果 Capacity 很高,节省的 Block 计算不多,路由开销可能占比上升。若 Capacity 很低,理论 FLOPs 很漂亮,但 Gather 后的矩阵过小可能导致 GPU 利用率下降。对于 GPU 来说,大而规整的矩阵计算往往比大量小而动态的工作更容易跑满执行单元。因此,实际性能取决于“节省的大矩阵成本”是否足以覆盖“动态组织数据的成本”。
高并发在线服务还存在批处理问题。传统 Dense decode 可以让一批请求在同一层执行相同 kernel。加入 MoD 后,同一批请求的 Token 可能在不同层走计算或旁路。Runtime 若无法把需要计算的 Token 高效重组,可能出现 kernel 数量增加、batch 变小、同步变多等问题。此时单模型 FLOPs 下降,整机吞吐却未必按比例提高。
KV Cache 让“跳过 Attention”比表面看起来更复杂
自回归 Transformer 的 Attention 会为历史 Token 保存 K/V 状态。普通模型每层都有连续的历史缓存;当前 Token 到达这一层时,会生成新的 K/V 并参与后续 Attention。MoD 若允许某个 Token 跳过带 Attention 的 Block,就必须明确这个位置在该层缓存里如何表示,以及后续 Token 查询这一层历史时能看到什么。
缓存语义必须与模型训练方式保持一致。不能在 Runtime 中自行决定“跳过计算,所以完全不记录这个位置”,除非模型结构和 Attention 实现本来就是这样定义的;也不能为了保持缓存形状而偷偷补算完整 Attention,否则节省会被吃掉。不同 MoD 设计可能对哪些子层被稀疏化、旁路状态如何参与后续计算有具体约束,部署系统需要跟随模型实现,而不是把“Skip Layer”理解成普通条件分支。
对于 PagedAttention 一类缓存管理器,还多出一个布局问题。缓存系统擅长按请求、层和 Token 位置管理连续增长的 K/V 块,而动态深度可能让不同层的有效计算位置呈现不同模式。工程上需要确认缓存索引、batch 调度和模型 kernel 都理解相同的路由语义。否则最危险的结果不是性能下降,而是模型在错误缓存历史上继续生成。
Router 学到的是计算分配策略,不是人类定义的“难度标签”
观察路由结果时,很容易讲成“标点简单所以跳过,数学 Token 困难所以深入计算”。这种叙述直观,却不应当当作机制定义。Router 接收到的是隐藏状态,优化目标来自模型训练损失和路由设计,它没有显式的“语法词”“困难题”标签。某些类型 Token 在统计上可能更常获得计算,这只能说明模型在当前训练分布下形成了这样的资源分配模式。
更有价值的分析是测量路由行为与模型结果之间的关系。例如,可以统计不同层的实际选择比例、不同 Token 类别的入选频率、同一位置跨层被选择的次数,以及路由模式在领域变化后是否发生漂移。如果代码流量进入服务后某些模式长期改变,Router 可能仍能正确工作,也可能暴露训练分布之外的行为,需要用任务质量而不是“看起来合理”来判断。
路由稳定性同样值得关注。若微小输入变化导致大量 Token 在 Top-K 边界附近来回切换,可能造成质量抖动或 kernel 工作模式变化。固定 Capacity 让总预算稳定,却不保证选择集合稳定。线上监控不能只记录平均 Capacity,因为它按设计本来就固定;更有信息的是不同层的分数分布、边界 Token、路由熵以及实际 Gather 后的有效 batch 形状。
生产环境应该同时测质量、计算和数据重排成本
评估 MoD 是否值得部署,第一组指标应是任务质量。在相同模型规模和训练条件下比较 Dense 与 MoD,观察困惑度或目标任务指标;对于生成服务,还需要覆盖短回答、长回答、代码、结构化输出和长上下文等不同流量。特别要检查连续跳过若干层后,模型是否在某些长尾输入上出现明显退化。
第二组指标是计算路径。记录每个 MoD Block 的 Capacity、被选 Token 数、Router 耗时、Top-K 耗时、Gather/Scatter 耗时以及 Block 本身耗时。只有把这些阶段拆开,才能知道 FLOPs 节省是否真正落到了 kernel 时间。如果 Block 时间下降 40%,但重排和同步新增 30%,端到端收益就远低于理论值。
第三组是服务级指标,包括 TTFT、每 Token 延迟、吞吐、GPU 利用率、HBM 带宽、kernel 数量、平均 batch 大小和不同并发度下的尾延迟。MoD 可能在单请求上节省计算,却因为动态分支破坏 continuous batching;也可能在大 batch 下收益明显,而低并发时重排开销占比过高。真正有意义的是沿真实请求分布做 A/B 测试。
还需要保留回退能力。动态路由属于模型执行语义的一部分,不能像普通性能开关一样随意关闭;但服务可以选择不部署 MoD 模型,或在灰度阶段让同类 Dense 基线承担对照流量。若 Runtime 对某个 shape、硬件或缓存模式支持不完整,应优先回退到经过验证的模型路径,而不是通过修改 Capacity 或跳过缓存检查来勉强运行。
自适应计算的目标不是让每个 Token 都“更聪明地思考”
MoD 放在更大的 Adaptive Computation 方向里更容易理解。AdaTape、Adaptive Computation Time、Early Exit 等工作都在尝试让网络根据输入改变计算预算,但“动态”可以发生在不同维度:有的改变层数,有的增加或减少外部计算单元,有的决定是否提前输出,有的像 MoE 一样改变激活参数集合。
MoD 的特点是把动态预算嵌入现代 Transformer 的深度方向,同时用固定 Capacity 保留可预测的总计算规模。它没有让每个 Token 自由选择任意深度,也没有让模型在运行时无限追加“思考步骤”。路由只是把有限 Block 计算重新分配给不同位置。
这条路线是否会成为和 MoE、GQA、FlashAttention 一样常见的基础组件,仍取决于系统实现。算法层已经说明“所有 Token × 所有层”不是唯一可行的计算组织方式;真正困难的是让 Router、Top-K、Gather/Scatter、KV Cache 和批处理协同工作,使减少的理论计算能够转化为稳定的吞吐和延迟收益。
因此,判断 Mixture-of-Depths 的价值不能只问“跳过了多少层”。更关键的问题是:在保持任务质量的前提下,一次层路由实际省下了多少高成本矩阵计算,又为 GPU Runtime 增加了多少数据重排和不规则调度。如果前者长期大于后者,动态深度才真正从一种计算稀疏思想变成可用的推理优化。