计算量的另一维度的稀疏化

MoE在"宽度"维度上实现了稀疏激活——不同token使用不同的专家。而Mixture-of-Depth（MoD）在"深度"维度上实现稀疏化——不同token经过不同数量的Transformer层。

这个想法直指一个核心观察：不是所有token都需要同等深度的计算。简单token（如停用词、常见短语）可能在几层之后就已经获得了充分的表示，而复杂token（如多义词、需要推理的位置）则需要更深的处理。

MoD的工作原理

基本架构

在MoD架构中，每一层都有一个路由器，决定哪些token需要继续经过当前层的计算，哪些可以"跳过"当前层直接传递到下一层：

class MoDLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, capacity_ratio=0.5):
        super().__init__()
        self.router = nn.Linear(d_model, 1)  # 输出一个标量路由分数
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = FFN(d_model, d_ff)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)
        self.capacity_ratio = capacity_ratio  # 参与计算的token比例
    
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 路由决策
        router_scores = self.router(x).squeeze(-1)  # [batch, seq_len]
        n_process = int(seq_len * self.capacity_ratio)
        
        # 选择top-k个token参与计算
        _, process_indices = torch.topk(router_scores, n_process, dim=-1)
        process_mask = torch.zeros_like(router_scores, dtype=torch.bool)
        process_mask.scatter_(1, process_indices, True)
        
        # 只有被选中的token经过注意力计算
        processed = x.clone()
        if process_mask.any():
            selected = x[process_mask].unsqueeze(0)
            attn_out = self.self_attn(self.norm1(selected))
            ffn_out = self.ffn(self.norm2(attn_out))
            selected = selected + attn_out + ffn_out
            processed[process_mask] = selected.squeeze(0)
        
        # 未选中的token直接传递（残差连接的等价效果）
        return processed

路由器的设计

路由器极其简单——只是一个单层线性映射加sigmoid，输出一个标量分数。这个分数表示"这个token需要当前层处理的程度"。

选择策略有两种：

  • Top-k选择：固定比例的token参与计算，保证计算量可预测
  • 阈值选择：分数超过阈值的token参与计算，计算量动态变化

实践中，Top-k选择更受青睐，因为它保证了计算量的可预测性，便于批处理和资源调度。

与MoE的关系与区别

MoD和MoE是正交的两种稀疏化策略：

维度MoEMoD
稀疏方向宽度（选择不同专家）深度（选择是否计算）
路由目标哪个专家是否参与
计算节省减少FFN计算量减少整层计算量
可组合—可以与MoE叠加

两者可以组合使用：在MoE层的FFN部分使用专家路由，同时在层级别使用MoD路由。这种"双重稀疏"可以进一步降低计算量。

MoD的训练

端到端训练

MoD可以端到端训练，路由器通过反向传播自动学习哪些token需要更深的计算。训练损失中通常加入一个小的负载均衡损失，确保每层处理的token数量大致均匀。

路由损失

为了鼓励路由器做出有意义的决策，可以加入一个辅助损失：

def mod_aux_loss(router_scores, capacity_ratio):
    """鼓励路由分数分布均匀的辅助损失"""
    probs = torch.sigmoid(router_scores)
    # 理想的平均概率应接近capacity_ratio
    mean_prob = probs.mean()
    return (mean_prob - capacity_ratio) ** 2

知识蒸馏初始化

从一个预训练的密集模型初始化MoD模型时，可以先让所有token都参与所有层的计算（capacity_ratio=1.0），然后逐渐降低到目标比例。这种"退火"策略比直接从头训练效果好得多。

性能与效率分析

计算量减少

以12层Transformer、capacity_ratio=0.5为例，平均每个token经过6层计算。与标准12层模型相比，FLOPs减少约50%。

但MoD的效率提升不是线性的——因为即使某些token跳过当前层，它们仍然需要在内存中保存和传递。实际加速比通常在1.5-2倍之间（相比理论2倍）。

性能保持

DeepMind的研究表明，在相同计算预算下，MoD模型优于更浅的密集模型。一个12层MoD（50%容量）的性能接近9-10层密集模型，而计算量相当于6层。

关键发现是：不同token确实需要不同深度的处理。路由器学会了将简单token（如"的"、“是"等）路由为跳过，而将复杂token（如多义词、推理关键点）路由为参与计算。

宨态计算深度的早期探索

MoD并非第一个"动态深度"的想法。此前的探索包括：

通用Transformer

在标准Transformer的基础上添加一个” halting mechanism"——每个位置在每层后决定是否停止计算。已停止的位置将隐藏状态复制到所有后续层。这种方法理论上优雅但训练困难。

早退机制

在每个层后添加一个分类器，如果分类器的置信度足够高，就直接输出结果而不经过后续层。这种方法在分类任务上有效，但不适合生成任务（因为不同位置的退出深度不同，破坏了批处理的效率）。

MoD相对于这些早期方法的关键改进在于Top-k选择保证了计算量的可预测性，使得高效的批处理成为可能。

2026年的MoD发展

层级MoD

不同层使用不同的capacity_ratio——浅层capacity较高（大部分token参与），深层capacity较低（只有需要深度推理的token参与）。这符合"浅层通用、深层专用"的直觉。

MoD + MoE联合训练

将MoD的深度稀疏与MoE的宽度稀疏联合优化。研究表明，这种组合可以在相同计算预算下提升3-5%的性能。

自适应推理深度

在推理时根据输入复杂度动态调整capacity_ratio——简单输入使用更低的capacity（更快），复杂输入使用更高的capacity（更准）：

def adaptive_inference(model, input_ids, complexity_estimator):
    """根据输入复杂度动态调整MoD参数"""
    complexity = complexity_estimator(input_ids)
    
    if complexity < 0.3:
        # 简单输入：激进的层跳过
        for layer in model.layers:
            layer.capacity_ratio = 0.3
    elif complexity < 0.7:
        # 中等输入：适度跳过
        for layer in model.layers:
            layer.capacity_ratio = 0.5
    else:
        # 复杂输入：几乎不跳过
        for layer in model.layers:
            layer.capacity_ratio = 0.8
    
    return model(input_ids)

工程实现挑战

批处理中的不规则性

同一batch中不同token的路由决策不同，导致计算模式不规则。高效实现需要使用掩码张量操作或scatter/gather操作，增加了kernel的复杂性。

与KV Cache的交互

当token跳过某些层时，其KV Cache在这些层中不存在。注意力计算需要处理这种"稀疏KV Cache"——某些token在特定层的K/V缺失。实践中通常用零向量填充，或直接在路由中确保KV相关层不被跳过。

分布式训练

MoD的层跳过在流水线并行中引入了复杂性——不同GPU上处理的token数量不均匀，导致负载不均衡。需要动态调度策略来缓解。

结语

MoD为模型效率优化开辟了新的维度。与MoE在宽度上的稀疏化互补，MoD在深度上的稀疏化让我们能在保持模型表达能力的同时减少冗余计算。随着"动态计算"理念从研究走向工程成熟，MoD有望成为下一代高效大模型的标配组件。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。