混合深度(MoD)：动态计算深度
计算量的另一维度的稀疏化 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是正交的两种稀疏化策略： 维度 MoE MoD 稀疏方向 宽度（选择不同专家） 深度（选择是否计算） 路由目标 哪个专家 是否参与 计算节省 减少FFN计算量 减少整层计算量 可组合 — 可以与MoE叠加 两者可以组合使用：在MoE层的FFN部分使用专家路由，同时在层级别使用MoD路由。这种"双重稀疏"可以进一步降低计算量。 ...
