A10 — Mixture of Depths、Early Exit 与动态计算

定位：训练 SOTA LLM 时，如何让不同 token 消耗不同的计算量——以更少 FLOPs 达到相同质量，或以相同 FLOPs 达到更高质量。

---

一、动态计算的核心直觉

1.1 观察：token 难度分布极不均匀

在自然语言中，token 的"预测难度"呈重尾分布。对一个 70B 参数、96 层的 Transformer：

• 功能词（"the", "of", "and"）在前 5-10 层就已经被 residual stream 充分编码，后续层对其 logit 贡献接近零

• 常见搭配（"United States", "machine learning"）在约 20-30 层收敛

• 需要长距离推理的 token（数学证明的下一步、代码中的变量类型推断）需要全部 96 层甚至仍然不够

实证证据：Elbayad et al. (2020) 在每一层插入分类器头，发现对超过 50% 的 token，中间层的预测已经与最终层一致。

1.2 类比 MoE：从"宽度路由"到"深度路由"

1.3 理论潜力

设 baseline 模型每个 token 经过 \(L\) 层，每层 FLOPs 为 \(F\)，总 FLOPs = \(LF\)。

若平均只有 \(C\) 比例的 token 需要经过每层：

\(\text{FLOPs}_{\text{MoD}} = C \cdot L \cdot F + (1-C) \cdot L \cdot F_{\text{router}}\)

其中 \(F_{\text{router}} \ll F\)（router 通常是单个线性层），因此：

\(\text{Speedup} \approx \frac{1}{C + (1-C) \cdot \epsilon} \approx \frac{1}{C}\)

当 \(C = 0.5\) 时，理论加速约 \(2\times\)；实际因 router 开销和 load balancing 约为 \(1.5\times\)。

---

二、Mixture of Depths (MoD)

核心论文：Raposo et al., "Mixture-of-Depths: Dynamically allocating compute in language models", 2024. [arxiv:2404.02258]

2.1 机制详解

Router 设计

每个 MoD 层 \(l\) 有一个 router 函数：

\(r_l(x) = w_l^T x\)

其中 \(w_l \in \mathbb{R}^d\) 是可学习参数，\(x \in \mathbb{R}^d\) 是该 token 在第 \(l\) 层的输入 hidden state。

注意：不需要 sigmoid 或 softmax——只需要标量分数用于排序。

Top-k 选择

给定一个 batch 中的 \(S\) 个 token（或一个 sequence 中的 \(S\) 个 token），选择 router score 最高的 \(k = \lfloor C \cdot S \rfloor\) 个 token 进行处理：

\(\mathcal{S}_l = \text{top-}k\big(\{r_l(x_i)\}_{i=1}^{S},\ k\big)\)

• 被选中的 token：经过该层的完整 Transformer block（attention + FFN）

• 未被选中的 token：直接通过 residual connection，\(x_i^{(l+1)} = x_i^{(l)}\)

训练时梯度传播

top-k 操作不可微。解决方案：straight-through estimator (STE)。

前向传播时使用离散的 top-k mask，反向传播时将梯度直通到 router 的连续输出：

\(\frac{\partial \mathcal{L}}{\partial w_l} = \frac{\partial \mathcal{L}}{\partial r_l} \cdot x\)

其中 \(\frac{\partial \mathcal{L}}{\partial r_l}\) 通过 STE 从离散 mask 传回。

2.2 实现

import torch
import torch.nn as nn

class MoDLayer(nn.Module):
    def __init__(self, d_model: int, capacity_ratio: float, transformer_block: nn.Module):
        super().__init__()
        self.router = nn.Linear(d_model, 1, bias=False)
        self.block = transformer_block
        self.capacity_ratio = capacity_ratio

    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None):
        # x: (batch, seq_len, d_model)
        B, S, D = x.shape
        k = int(S * self.capacity_ratio)

        scores = self.router(x).squeeze(-1)  # (B, S)
        topk_indices = scores.topk(k, dim=-1).indices  # (B, k)
        topk_indices_sorted, _ = topk_indices.sort(dim=-1)

        # gather selected tokens
        idx_expanded = topk_indices_sorted.unsqueeze(-1).expand(-1, -1, D)
        x_selected = x.gather(1, idx_expanded)  # (B, k, D)

        # process through transformer block
        x_processed = self.block(x_selected, attention_mask=None)

        # scatter back
        out = x.clone()
        out.scatter_(1, idx_expanded, x_processed)
        return out

注意：上面的实现省略了因果 attention mask 在 subset token 上的正确处理。实际中需要：(1) 维护 token 的原始位置信息用于 RoPE 计算；(2) 被跳过的 token 不参与 attention 计算但仍可被 attend to——这是一个设计选择，Raposo et al. 选择让被跳过 token 完全不参与。

2.3 关键设计决策

2.4 结果分析

Raposo et al. 的核心实验使用 isoFLOP 比较（固定总训练 FLOPs，比较最终 loss）：

• 12.5% capacity, 每隔一层 MoD：与 dense baseline 在 isoFLOP 条件下 loss 相当，但推理时只需 ~50% FLOPs

• 等价理解：把省下的 FLOPs 用来增大模型（更多参数），可以在相同总 FLOPs 下获得更低 loss

• Router 学到的 pattern：功能词和标点几乎在所有层都被跳过；名词和动词在中间层开始被跳过；需要长距离依赖的 token 始终被处理

2.5 为什么 MoD 有效

直觉解释与理论支撑：

1. Residual stream 是主干：Transformer 的 residual connection 意味着跳过一层相当于 identity mapping，信息不会丢失，只是不被进一步精炼

2. 层间冗余：Dalvi et al. (2020) 证明相邻层的表示高度相似（CKA > 0.95），删除中间层对性能影响很小

3. Token 难度的 Zipf 分布：自然语言中简单 token 数量远多于困难 token，所以即使 \(C\) 很小（如 0.125），困难 token 仍被充分处理

---

三、Early Exit

3.1 基本思想

Early Exit 是动态计算的另一个方向：不是在每层选择 哪些 token 被处理（MoD），而是让每个 token 在某一层 提前退出 不再经过后续层。

\(\hat{y}_i = \text{LM\_head}(h_i^{(l^*)})\)

其中 \(l^* \leq L\) 是 token \(i\) 的退出层。

3.2 CALM: Confident Adaptive Language Modeling

论文：Schuster et al., "Confident Adaptive Language Modeling", NeurIPS 2022. [arxiv:2207.07061]

退出判据

在每一层 \(l\)，用一个轻量分类器（或直接复用 LM head）计算该层的输出分布，当满足以下条件之一时退出：

方法 1：Softmax 熵阈值

\(H(p_l) = -\sum_v p_l(v) \log p_l(v) < \tau\)

方法 2：Top-1 概率阈值

\(\max_v p_l(v) > \tau\)

方法 3：相邻层一致性

\(\text{JSD}(p_l \| p_{l-1}) < \tau\)

其中 JSD 是 Jensen-Shannon 散度。

实现的关键细节

class EarlyExitTransformer(nn.Module):
    def __init__(self, layers: nn.ModuleList, lm_head: nn.Linear, threshold: float):
        super().__init__()
        self.layers = layers
        self.lm_head = lm_head
        self.threshold = threshold
        self.exit_classifiers = nn.ModuleList([
            nn.Linear(layers[0].d_model, 1) for _ in layers
        ])

    def forward_inference(self, x: torch.Tensor):
        # only works for batch_size=1 in naive implementation
        for l, layer in enumerate(self.layers):
            x = layer(x)
            confidence = torch.sigmoid(self.exit_classifiers[l](x[:, -1:, :]))
            if confidence.item() > self.threshold:
                return self.lm_head(x)
        return self.lm_head(x)

训练策略

CALM 的训练需要每层都能产出合理的预测。两种方案：

1. 辅助 loss：在每层都加 LM head loss，总 loss 为加权和

\(\mathcal{L} = \sum_{l=1}^{L} \alpha_l \cdot \mathcal{L}_{\text{CE}}^{(l)}\)

1. 一致性蒸馏：用最后一层的输出蒸馏中间层

\(\mathcal{L}_{\text{distill}}^{(l)} = \text{KL}(p_L \| p_l)\)

3.3 Early Exit 的根本困难

3.4 SkipDecode

论文：Del Corro et al., "SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference", 2024. [arxiv:2307.02628]

SkipDecode 试图解决 batched early exit 的问题：为同一 batch 内所有 token 设定统一的退出层（per-position 而非 per-token），使得 batch 内可以高效并行。

关键洞察：token position 越靠后（即越早生成的 token），其退出层可以越浅，因为后续 token 的注意力会"修正"前面 token 的次优表示。

---

四、Attention Residual 与 Layer Skipping

4.1 Attention Residual

核心观察：相邻 Transformer 层的 attention pattern 高度相似。测量 \(\text{cos}(A_l, A_{l-1})\) 通常 > 0.8。

利用这一点，可以复用上一层的 attention weight，减少重新计算的开销：

\(A_l = \alpha A_{l-1} + (1-\alpha)\, \text{softmax}\!\left(\frac{Q_l K_l^T}{\sqrt{d_k}}\right)\)

其中 \(\alpha \in [0, 1]\) 是可学习参数或超参数。

变体：Cross-Layer Attention (CLA)

Brandon et al. (2024) [arxiv:2405.12981] 提出相邻层共享 KV cache：

\(\text{Attn}_l(Q_l, K_{l'}, V_{l'}) \quad \text{where } l' = l - (l \mod s)\)

每 \(s\) 层共享一组 KV，减少 KV cache 内存约 \(s\times\)。

4.2 Stochastic Depth / Layer Dropping

论文：Huang et al., "Deep Networks with Stochastic Depth", ECCV 2016. [arxiv:1603.09382]

训练时以概率 \(p_l\) 跳过第 \(l\) 层：

\(x^{(l+1)} = \begin{cases} x^{(l)} + f_l(x^{(l)}) & \text{w.p. } 1 - p_l \\ x^{(l)} & \text{w.p. } p_l \end{cases}\)

通常 \(p_l\) 随层数线性增长：\(p_l = \frac{l}{L} \cdot p_{\max}\)。

效果：

• 起到正则化作用（类似 Dropout 但在层维度）

• 隐式训练了一个层数可变的模型集合

• 推理时可以直接跳过尾部层，性能退化平滑

4.3 LayerDrop

论文：Fan et al., "Reducing Transformer Depth on Demand with Structured Dropout", ICLR 2020. [arxiv:1909.11556]

LayerDrop = Stochastic Depth 在 Transformer 上的系统化应用：

class LayerDropTransformer(nn.Module):
    def __init__(self, layers: nn.ModuleList, drop_rate: float = 0.2):
        super().__init__()
        self.layers = layers
        self.drop_rate = drop_rate

    def forward(self, x: torch.Tensor):
        for layer in self.layers:
            if self.training and torch.rand(1).item() < self.drop_rate:
                continue
            x = layer(x)
        return x

    def forward_pruned(self, x: torch.Tensor, keep_every_n: int = 2):
        for i, layer in enumerate(self.layers):
            if i % keep_every_n != 0:
                continue
            x = layer(x)
        return x

关键发现：训练时 drop_rate=0.2，推理时可以均匀跳过一半的层（keep_every_n=2），性能仅下降 ~2%。

4.4 与深度扩展的关联

---

五、MoD vs MoE 全面对比

比较维度	MoE	MoD	为什么这个区别重要
路由维度	token → expert（宽度）	token → 是否处理（深度）	正交维度意味着可以独立或联合优化
路由粒度	每个 token 选 \(k\) 个 expert	每层选 \(C\%\) 的 token 进入	MoE 的粒度在 expert 级别（几十个），MoD 在 token 级别（几千个），MoD 路由决策更多
FLOPs 节省来源	每 token 只激活部分参数	部分 token 跳过整层	MoE 节省"宽度"计算，MoD 节省"深度"计算；理论上可以相乘
训练难度	load balancing loss、expert collapse	router 学习什么是"重要"token	MoE 的 failure mode 是 expert 不均匀；MoD 的 failure mode 是 router 学到 trivial 的 pattern（如总是跳过相同位置的 token）
推理加速	需要 EP（expert parallelism）支持	天然减少计算量，无需特殊并行策略	MoE 的推理瓶颈在 all-to-all 通信；MoD 的瓶颈在动态 batch 重组
参数效率	总参数远大于激活参数	总参数 = 激活参数（+ 微小 router）	MoE 用参数换计算；MoD 用智能调度换计算
内存占用	高（所有 expert 需驻留或 offload）	与 dense 模型几乎相同	MoD 的优势：不增加模型内存，只减少计算
成熟度	已被 Mixtral、DeepSeek 等工业化验证	仅有学术验证，尚无大规模部署	MoE 有 Switch Transformer (2022) 以来的工程积累；MoD 在 2024 年才被正式提出

---

六、MoD + MoE 组合：双维度计算分配

6.1 概念框架

最优的计算分配应该同时在两个维度进行：

\(\text{Compute}(x_i, l) = \begin{cases} \text{Expert}_j\ \text{if}\ \text{MoD-router}_l(x_i) > \tau_l\ \text{and}\ j = \text{MoE-router}_l(x_i) \\ \text{Skip (residual)}\ \text{if}\ \text{MoD-router}_l(x_i) \leq \tau_l \end{cases}\)

即：先决定 token 是否需要在这层处理（MoD），如果需要，再决定由哪个 expert 处理（MoE）。

6.2 FLOPs 分析

设 MoE 的激活率为 \(k/N\)（\(k\) 个 active experts / \(N\) 个 total experts），MoD 的 capacity ratio 为 \(C\)：

\(\text{FLOPs}_{\text{MoD+MoE}} = C \cdot \frac{k}{N} \cdot F_{\text{dense}}\)

例如 \(C = 0.5\)，\(k/N = 2/16 = 0.125\)：

\(\text{FLOPs} = 0.5 \times 0.125 \times F_{\text{dense}} = 0.0625 \times F_{\text{dense}}\)

理论上 16x 加速——但实际中 attention 层不受 MoE 影响，且 router 和通信开销不可忽略。更现实的估算（仅 FFN 层使用 MoE，attention + FFN 比例约 1:2）：

\(\text{FLOPs} \approx C \cdot (F_{\text{attn}} + \frac{k}{N} \cdot F_{\text{FFN}}) \approx 0.5 \times (F_{\text{attn}} + 0.125 \times F_{\text{FFN}})\)

6.3 DeepSeek 的启示

DeepSeek-V2/V3 使用了 MoE 但未使用 MoD。然而其设计暗示了组合的可能性：

• MLA (Multi-head Latent Attention) 已经在 attention 层做了压缩（宽度方向）

• DeepSeekMoE 的 fine-grained experts 将路由粒度做到极致（256 个 expert 选 8 个）

• 如果在此基础上加入 MoD，每层只处理 50% 的 token，可以进一步将推理 FLOPs 减半

6.4 实现挑战

---

七、工程成熟度与实践指南

7.1 技术成熟度评估

7.2 对训练基础设施的要求

# MoD-aware data loader: tokens within a batch should have diverse difficulty
# to ensure router learning receives varied signal

def create_mod_aware_batch(dataset, batch_size, seq_len):
    # mix documents of varying complexity
    # avoid batches where all tokens are "easy" or all "hard"
    # this prevents router collapse (always skip / never skip)
    easy = [d for d in dataset if d['perplexity'] < ppl_median]
    hard = [d for d in dataset if d['perplexity'] >= ppl_median]
    batch = []
    for _ in range(batch_size // 2):
        batch.append(random.choice(easy)[:seq_len])
        batch.append(random.choice(hard)[:seq_len])
    return torch.stack(batch)

7.3 推理优化

MoD 的推理优化关键在于避免动态 shape 带来的开销：

---

八、追问延伸

1. MoD 的 router 会退化吗？ 如果 router 学到了基于 token position 而非 token content 的 pattern（如"总是跳过偶数位置"），它就退化为静态 layer skip。需要监控 router decision 的 entropy 和 token-content 相关性。

2. MoD 与 speculative decoding 的关系？ Speculative decoding 用小模型预测、大模型验证；MoD 可以理解为"同一模型的浅层预测、深层验证"——如果浅层已经足够自信，就不需要深层验证。

3. 为什么不直接训练一个更浅的模型？ 因为 不同 token 需要不同的深度。一个 30 层模型对所有 token 都只用 30 层；一个 96 层 MoD 模型对简单 token 用 10 层、对困难 token 用 96 层，平均可能也是 30 层但效果远好于固定 30 层。

4. MoD 训练的 scaling behavior？ 目前缺乏系统性的 MoD scaling law 研究。关键问题：当模型规模增大时，最优 capacity ratio \(C\) 是保持不变还是会变化？直觉上，更大的模型有更多冗余层，\(C\) 应该可以更小。

5. Dynamic compute 在 post-training 阶段的应用？ RL fine-tuning 时，model 应该把更多计算花在 reward 信号强的 token 上。这与 MoD 的直觉一致：RL 阶段可以微调 router 使其关注 reward-relevant token。

---

参考文献

1. Raposo, D., Ritter, S., Richards, B., et al. "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models." arXiv preprint, 2024. [arxiv:2404.02258]

2. Schuster, T., Fisch, A., Gupta, J., et al. "Confident Adaptive Language Modeling." NeurIPS, 2022. [arxiv:2207.07061]

3. Elbayad, M., Gu, J., Grave, E., Auli, M. "Depth-Adaptive Transformer." ICLR, 2020. [arxiv:1910.10073]

4. Fan, A., Grave, E., Joulin, A. "Reducing Transformer Depth on Demand with Structured Dropout." ICLR, 2020. [arxiv:1909.11556]

5. Huang, G., Sun, Y., Liu, Z., Sedra, D., Weinberger, K. "Deep Networks with Stochastic Depth." ECCV, 2016. [arxiv:1603.09382]

6. Del Corro, L., Del Giorno, A., Aber, S., et al. "SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference." arXiv preprint, 2024. [arxiv:2307.02628]

7. Brandon, W., Mishra, M., Nrusimha, A., Panda, R., Ragan-Kelly, J. "Reducing Transformer Key-Value Cache Size with Cross-Layer Attention." arXiv preprint, 2024. [arxiv:2405.12981]

8. Dalvi, F., Durrani, N., Sajjad, H., et al. "Analyzing Redundancy in Pretrained Transformer Models." EMNLP, 2020. [arxiv:2004.04010]

9. Kaplan, J., McCandlish, S., Henighan, T., et al. "Scaling Laws for Neural Language Models." arXiv preprint, 2020. [arxiv:2001.08361]

10. DeepSeek-AI. "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model." arXiv preprint, 2024. [arxiv:2405.04434]

---

↑ 上级 · A. 基础理论