Mixture of Depths (MoD)

深度混合（MoD）是一种Transformer技术，由学习得到的路由器逐模块决定哪些token要参与计算、哪些 token 通过残差连接跳过该模块，因此不同 token 获得的计算量各不相同。Google DeepMind 的 David Raposo、Adam Santoro 及其同事于 2024 年 4 月在论文《深度混合：动态分配基于 Transformer 的语言模型中的计算量》中介绍了该技术。固定容量 k 决定每个路由模块处理多少个 token，因此，即使具体选择哪些 token 取决于输入，计算预算和张量形状也能预先确定。MoD 大型语言模型（LLM）在相同训练FLOPs和实际运行时间下，达到了与基线模型相当的质量；每次前向传递只需使用基线模型一小部分的 FLOPs，采样速度则提高了 50% 以上。在固定训练预算下，表现最好的 MoD 模型比最佳基线模型更大，损失也更低。

MoD 路由机制如何工作#

在路由模块之前，线性路由器会为每个 token 生成一个标量分数。得分最高的 k 个 token 会经过模块中的自注意力和 MLP，其输出会乘以对应的路由器分数，再加回残差流。乘以分数后，路由器也处于梯度路径上，因此反向传播可以教会路由器判断哪些 token 能从计算中受益。其余 token 则原样通过，不参与该模块的注意力计算，因此后续 token 在该模块中也无法关注它们；注意力成本会随容量的平方缩减：处理一半的 token，查询—键乘积的数量就只有四分之一。论文中效果最佳的设置是每隔一个模块进行路由，容量设为序列长度的 12.5%。以下是在 PyTorch 中实现一个路由模块的精简版本，其中仅包含 MLP：

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

它会打印 Block ran on 16 of 128 tokens; 112 skipped it unchanged：该模块中的 MLP 只执行了通常八分之一的计算量，而输出张量仍保留完整的序列形状，供下一个模块使用。

生成过程中的路由#

Top-k 选择不具备因果性：一个 token 是否属于其序列中得分最高的 k 个 token，取决于它后面 token 的分数；而模型逐个 token 生成文本时，后续 token 尚不存在。作者通过两种方式使路由具备因果性：一种是辅助二元交叉熵损失，用于训练路由器预测 token 是否属于 top-k，代价是语言建模损失仅增加 0.2% 至 0.3%；另一种是在停止梯度的输入上使用单独的小型 MLP 预测器，准确率很快就达到 97% 至 99%。无论采用哪种方法，自回归采样的质量损失都很小。训练好的路由器还呈现出一种规律：有些 token 会在模型深度上的每个路由模块中参与计算，而另一些 token 则会尽可能跳过模块。

MoD 与相关技术的对比#

• 专家混合（MoE）：在每一层中，将每个 token 路由到一个或少数几个专家 MLP，从而增加参数数量，同时保持每个 token 的计算量不变。论文还将 MoD 与 MoE 结合成 MoDE；在这种方法中，token 可以被路由到专家，也可以选择不执行操作，两者的增益会叠加。
• 提前退出：token 在某一层后停止处理，跳过所有后续层。MoD 可以跳过一个模块，但仍继续使用后续模块。
• Token 合并（ToMe）：在视觉 Transformer 中合并相似的 token，使后续层处理更少的 token，且不需要学习路由器。
• 模型剪枝：永久移除权重，因此所有输入都会使用相同的较小网络。
• 模型路由：将整个请求发送到多个模型中的一个，这是在完整模型层面实现的相同理念。

大多数网络在构建时就固定了深度。Ultralytics YOLO26会根据模型的规模设置深度：模型 YAML中的深度乘数在 n、s 和 m 规模下为 0.50，在 l 和 x 规模下为 1.00。因此，配置为重复 2 次的 C3k2 阶段，在 YOLO26n 中运行 1 个模块，在 YOLO26x 中运行 2 个模块；每张图像都会经过其中的所有模块。