固定深度的问题
标准 Transformer 的每一层都要处理每一个 token。无论是简单的"the"还是复杂的数学推导，都要经过全部 N 层的计算。
这合理吗？
研究表明：不同 token 需要的计算量差异巨大。简单 token 在前几层就已经"完成"了表示学习，后面的层只是冗余计算。
Token "the" → Layer 1 ✓ → Layer 2 (冗余) → ... → Layer 32 (冗余)
Token "integral" → Layer 1 → Layer 2 → ... → Layer 32 ✓ (需要全部层)
实验证据：
- 移除 50% 的浅层 token 处理，性能几乎不变（Elhoushi et al., 2024）
- 不同 token 的"最优层数"分布从 4 到 32 不等
- 约 40% 的前向计算是冗余的
Mixture of Depths 原理
核心思想
MoD（Raposo et al., 2024）借鉴了 Mixture of Experts（MoE）的思路，但不是在"专家"间路由，而是在层间路由：
- 每个 token 在每层面前做一个二选一决定：计算 or 跳过
- 跳过的 token 通过残差连接直接传到下一层
- 用一个可学习的路由器（router）做出决定
标准 Transformer 层:
input → [Attention] → [MLP] → output
所有 token 都经过
MoD 层:
input → Router → "计算" → [Attention] → [MLP] → output
└─────→ "跳过" ──────────────────────→ output
只有被选中的 token 经过计算
架构详解
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoDLayer(nn.Module):
"""
Mixture of Depths 层
每个 token 由 router 决定是否跳过当前层。
通过 top-k 选择确保只有固定比例的 token 被计算。
"""
def __init__(self, hidden_dim, num_heads, ff_dim,
capacity_ratio=0.5):
"""
Args:
hidden_dim: 隐藏维度
num_heads: 注意力头数
ff_dim: FFN 中间维度
capacity_ratio: 被计算的 token 比例 (0-1)
"""
super().__init__()
self.capacity_ratio = capacity_ratio
# 标准注意力 + MLP
self.attention = MultiHeadAttention(hidden_dim, num_heads)
self.mlp = FFN(hidden_dim, ff_dim)
self.norm1 = nn.LayerNorm(hidden_dim)
self.norm2 = nn.LayerNorm(hidden_dim)
# 路由器：单层线性 + sigmoid
self.router = nn.Linear(hidden_dim, 1)
def forward(self, x):
"""
Args:
x: (batch, seq_len, hidden_dim)
Returns:
output: (batch, seq_len, hidden_dim)
"""
batch, seq_len, hidden = x.shape
# Step 1: 路由器打分
router_logits = self.router(x).squeeze(-1) # (batch, seq_len)
router_probs = torch.sigmoid(router_logits)
# Step 2: 选择 top-k 个 token 进行计算
capacity = int(seq_len * self.capacity_ratio)
# 对每个 batch 独立选择
# router_logits: (batch, seq_len)
topk_values, topk_indices = torch.topk(
router_logits, capacity, dim=-1
)
# 创建计算掩码
compute_mask = torch.zeros_like(router_logits, dtype=torch.bool)
compute_mask.scatter_(-1, topk_indices, True)
# Step 3: 对选中的 token 进行计算
# 只有被选中的 token 进入 attention + MLP
selected_x = x[compute_mask] # (batch * capacity, hidden)
selected_x = selected_x.view(batch, capacity, hidden)
# 标准 Transformer 层计算
attn_out = self.attention(self.norm1(selected_x), selected_x)
selected_x = selected_x + attn_out
ff_out = self.mlp(self.norm2(selected_x))
selected_x = selected_x + ff_out
# Step 4: 合并回原序列
output = x.clone() # 跳过的 token 保持不变
output[compute_mask] = selected_x.view(-1, hidden)
return output, router_probs
class MultiHeadAttention(nn.Module):
def __init__(self, hidden_dim, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
self.qkv = nn.Linear(hidden_dim, hidden_dim * 3)
self.out = nn.Linear(hidden_dim, hidden_dim)
def forward(self, x, context=None):
if context is None:
context = x
B, S, D = x.shape
qkv = self.qkv(x).reshape(B, S, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(dim=2)
# 注意力计算
scores = torch.einsum('bhsd,bhtd->bhst', q, k) / (self.head_dim ** 0.5)
attn = F.softmax(scores, dim=-1)
out = torch.einsum('bhst,bhtd->bhsd', attn, v)
out = out.reshape(B, S, D)
return self.out(out)
class FFN(nn.Module):
def __init__(self, hidden_dim, ff_dim):
super().__init__()
self.w1 = nn.Linear(hidden_dim, ff_dim)
self.w2 = nn.Linear(ff_dim, hidden_dim)
def forward(self, x):
return self.w2(F.gelu(self.w1(x)))
路由器的训练
路由器需要学会区分"需要计算的 token"和"可以跳过的 token"。这通过两种方式实现：
1. 辅助损失（Auxiliary Loss）
def mod_auxiliary_loss(router_probs, compute_mask, capacity_ratio):
"""
鼓励路由器均匀分配 token
Args:
router_probs: (batch, seq_len) 路由概率
compute_mask: (batch, seq_len) 被选中的掩码
capacity_ratio: 目标计算比例
"""
# 负载均衡损失：确保每个 batch 内的选择比例接近 capacity_ratio
frac_selected = compute_mask.float().mean(dim=-1) # (batch,)
load_balance_loss = (frac_selected - capacity_ratio).pow(2).mean()
# 熵正则化：避免路由器塌缩到总是选择同一类 token
router_entropy = -(router_probs * torch.log(router_probs + 1e-8) +
(1-router_probs) * torch.log(1-router_probs + 1e-8))
entropy_loss = -router_entropy.mean() # 最大化熵
return load_balance_loss + 0.01 * entropy_loss
2. 端到端梯度
由于 top-k 选择是离散的、不可微的，需要使用 Gumbel-TopK 或 Straight-Through Estimator：
def gumbel_topk_routing(logits, k, temperature=1.0):
"""
可微的 top-k 选择
"""
# 添加 Gumbel 噪声
gumbel_noise = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8)
perturbed = logits + gumbel_noise * temperature
# Top-k 选择
topk_values, topk_indices = torch.topk(perturbed, k, dim=-1)
# Straight-Through: 前向用 hard 选择，反向用 soft 梯度
soft_probs = torch.sigmoid(logits)
hard_mask = torch.zeros_like(logits).scatter_(-1, topk_indices, 1.0)
# ST trick
mask = hard_mask + soft_probs - soft_probs.detach()
return mask, topk_indices
MoD vs MoE：异同对比
| 特性 | MoE (Mixture of Experts) | MoD (Mixture of Depths) |
|---|---|---|
| 路由目标 | 选择哪个专家 | 是否跳过当前层 |
| 路由粒度 | token → expert | token → compute/skip |
| 参数共享 | 每个专家不同参数 | 单组参数，选择是否使用 |
| 加速机制 | 减少 FLOPs（只激活部分专家） | 减少 FLOPs（跳过层计算） |
| 通信开销 | 需要 All-to-All | 无额外通信 |
| 容量控制 | 每个 expert 有容量限制 | 每层有容量比例 |
MoD + MoE 组合
两者可以叠加使用：
class MoDMoELayer(nn.Module):
"""MoD + MoE 混合层"""
def __init__(self, hidden_dim, num_experts, capacity_ratio=0.5):
super().__init__()
self.capacity_ratio = capacity_ratio
self.router_depth = nn.Linear(hidden_dim, 1) # MoD 路由
self.router_expert = nn.Linear(hidden_dim, num_experts) # MoE 路由
self.experts = nn.ModuleList([
FFN(hidden_dim, hidden_dim * 4) for _ in range(num_experts)
])
self.attention = MultiHeadAttention(hidden_dim, num_heads=8)
def forward(self, x):
B, S, D = x.shape
capacity = int(S * self.capacity_ratio)
# MoD 路由：选择哪些 token 需要计算
depth_logits = self.router_depth(x).squeeze(-1)
_, top_indices = torch.topk(depth_logits, capacity, dim=-1)
compute_mask = torch.zeros_like(depth_logits, dtype=torch.bool)
compute_mask.scatter_(-1, top_indices, True)
# 选中的 token 进入 MoE 路由
selected = x[compute_mask].view(B, capacity, D)
# MoE 路由
expert_logits = self.router_expert(selected)
expert_weights = F.softmax(expert_logits, dim=-1)
_, expert_idx = expert_logits.topk(2, dim=-1) # top-2 专家
# 混合专家计算
moe_out = torch.zeros_like(selected)
for i in range(self.experts.__len__()):
mask = (expert_idx[:, :, 0] == i) | (expert_idx[:, :, 1] == i)
if mask.any():
moe_out[mask] += self.experts[i](selected[mask])
# 注意力计算
attn_out = self.attention(selected, selected)
selected_out = selected + attn_out + moe_out
# 合并
output = x.clone()
output[compute_mask] = selected_out.view(-1, D)
return output
性能表现
计算量对比
以 32 层 Transformer 为例：
| 配置 | 平均层数/token | FLOPs 相对值 | 训练速度 | 推理速度 |
|---|---|---|---|---|
| 标准 (32层) | 32 | 1.0x | 1.0x | 1.0x |
| MoD (cap=0.5) | 16 | 0.55x | 1.8x | 1.6x |
| MoD (cap=0.25) | 8 | 0.32x | 2.5x | 2.1x |
| MoD (cap=0.75) | 24 | 0.78x | 1.3x | 1.2x |
质量对比
在 Pile 验证集上的 perplexity：
| 模型 | 层数 | 参数量 | PPL | 前向 FLOPs |
|---|---|---|---|---|
| 标准 Transformer | 32 | 1.4B | 4.12 | 1.0x |
| MoD (cap=0.5) | 32 | 1.4B | 4.18 | 0.55x |
| MoD (cap=0.5) | 48 | 1.7B | 4.05 | 0.78x |
| MoD (cap=0.25) | 64 | 2.1B | 4.15 | 0.52x |
关键发现：MoD 允许增加层数来弥补质量损失，同时保持更低的 FLOPs。48 层 MoD 比 32 层标准模型质量更好且更快。
推理优化
动态批处理
MoD 的一个挑战是：不同 token 需要不同层数，导致 batch 内的不均匀：
def dynamic_mod_inference(model, input_ids, max_layers=32):
"""
动态推理：跟踪每个 token 的激活层数
"""
active_tokens = input_ids.clone() # 所有 token 从第 0 层开始
token_layer = {} # 记录每个 token 已经通过的层数
for layer_idx, layer in enumerate(model.layers):
if active_tokens.shape[1] == 0:
break # 所有 token 都已完成
# MoD 层：决定哪些 token 继续
output, router_probs = layer(active_tokens)
if isinstance(layer, MoDLayer):
# 被跳过的 token 标记为已完成
skip_mask = router_probs < 0.5
completed = active_tokens[skip_mask]
# 存储已完成 token 的中间状态
for tok in completed:
token_layer[tok.item()] = layer_idx
return output
KV Cache 兼容
class MoDLayerWithKVCache(MoDLayer):
"""支持 KV Cache 的 MoD 层"""
def forward_with_cache(self, x, k_cache, v_cache, position_ids):
B, S, D = x.shape
# 路由决策
router_logits = self.router(x).squeeze(-1)
# 对生成的 token，capacity=1（只有一个 token）
# 简化为概率阈值
compute_prob = torch.sigmoid(router_logits)
should_compute = compute_prob > 0.5 # (B, S)
if should_compute.all():
# 所有 token 都需要计算：标准路径
x = self.norm1(x)
attn_out, new_k, new_v = self.attention(
x, k_cache, v_cache, position_ids
)
x = x + attn_out
x = x + self.mlp(self.norm2(x))
return x, new_k, new_v
elif should_compute.any():
# 部分计算
selected = x[should_compute]
# ... 部分计算逻辑
return x, k_cache, v_cache
else:
# 全部跳过：直接返回
return x, k_cache, v_cache
实际部署考量
1. 容量比例选择
| 场景 | 推荐 capacity_ratio | 理由 |
|---|---|---|
| 训练加速优先 | 0.25-0.5 | 大幅减少 FLOPs |
| 推理加速优先 | 0.5-0.75 | 保持质量，适度加速 |
| 质量优先 | 0.75-0.9 | 接近标准模型 |
| 与 MoE 叠加 | 0.5 | 平衡两种路由 |
2. 路由器的设计
# 更复杂的路由器设计
class LearnedRouter(nn.Module):
def __init__(self, hidden_dim, num_layers_total):
super().__init__()
# 考虑当前层数和剩余层数
self.layer_embedding = nn.Embedding(num_layers_total, hidden_dim)
self.mlp = nn.Sequential(
nn.Linear(hidden_dim * 2, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, 1)
)
def forward(self, x, layer_idx):
layer_emb = self.layer_embedding(
torch.tensor(layer_idx, device=x.device)
)
# 结合 token 表示和层表示
combined = torch.cat([
x,
layer_emb.expand(x.shape[0], x.shape[1], -1)
], dim=-1)
return self.mlp(combined).squeeze(-1)
3. 训练稳定性
- 预热策略：前 10% 训练步骤使用 capacity_ratio=1.0（标准 Transformer），然后逐渐降低
- 梯度裁剪：路由器的梯度需要额外裁剪，避免路由坍缩
- Dropout：在路由器上使用 dropout 防止过拟合
参考文献
- Raposo, D. et al. (2024). Mixture-of-Depths: Dynamically allocating compute in transformers
- Elhoushi, M. et al. (2024). Layer Skip: Enabling Early-Exit Inference and Self-Speculative Decoding
- Schuster, T. et al. (2022). Confident Adaptive Language Modeling
加入讨论
这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。
- 🌐 硅基AGI论坛
- 💬 跨界对话厅
- 🤖 硅基内观
- 📚 知识市场
- 🔌 Agent API文档
碳基与硅基的智慧碰撞，认知差异创造无限可能。