Tuning Language Models by Mixture-of-Depths Ensemble

技术摘要：通过混合深度集成（Mixture-of-Depths Ensemble, MoDE）进行语言模型微调

问题陈述
传统的基于 Transformer 的大语言模型（LLMs）在微调过程中，完全依赖于最终层的表示来进行预测和损失计算。虽然诸如“逻辑透镜”（logit lens）之类的可解释性工具已经表明，后期层的隐藏状态已经包含了大量成型的、与任务相关的预测，但这些中间表示通常被视为隐性的副产品，而非显性的训练信号。因此，后期层的潜在预测能力未能得到充分利用，且标准的微调仅优化了最后一层，导致中间表示只能间接地得到优化。此外，现有的利用中间层的方法（例如用于提高效率的早期退出或表示转向）通常并未将层特定的逻辑值（logits）视为最终预测的直接、可训练组件。

方法论：混合深度集成（MoDE）
作者提出了混合深度集成（Mixture-of-Depths Ensemble, MoDE），该框架将 LLM 最后 $k$ 层的输出视为单个模型内的一组“专家”的集成。与将不同 Token 路由到不同模型的混合专家模型（MoE）不同，MoDE 将最终预测路由到同一模型内不同深度的逻辑值加权组合。

该框架由三个核心部分组成：

1. 后期层逻辑值提取（Late-Layer Logit Extraction）： 对于一个具有 $n$ 层的模型，选择最后 $k$ 层（$K = \{L_{n-k+1}, \dots, L_n\}$）。将词表头 $\phi(\cdot)$ 应用于这 $k$ 层中的每一层的隐藏状态，以生成逻辑值 $\ell_i(x)$。
2. 学习型路由网络（Learned Routing Network）： 一个轻量级的路由器 $G(x)$（实现为一个单层线性层并接 softmax），根据最后 $k$ 层之前的隐藏状态 $x = H_{n-k}$ 计算路由权重。最终的逻辑值通过以下方式计算：
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   路由器可以被设计为稀疏的（Top-K），以实现推理过程中的早期退出。
3. 通过归一化与蒸馏进行适配（Adaptation via Normalization and Distillation）： 直接将预训练的头部应用于中间层往往会导致较高的初始损失。为了在不增加沉重参数（不同于每层仿射矩阵）的情况下解决这一问题，作者引入了：
   • 逐层归一化（Per-layer Normalization, $N_k$）： 添加可训练的归一化模块（匹配预训练架构，例如 LayerNorm），以使表示适配到头部输入空间。
   • 蒸馏损失（Distillation Loss, $L_{distill}$）： 最后一层充当“教师”。应用 KL 散度损失，将较早的后期层的输出分布向最后一层的分布拉近，从而稳定训练并使集成与任务目标保持一致。

总损失为 $L_{MoDE} = L_{task} + \lambda L_{distill}$。MoDE 被设计为一个插件机制，可以叠加在任何现有的参数高效微调（PEFT）方法（如 LoRA、DoRA）之上，或者独立使用。

核心贡献

• 可训练的后期层信号： 论文证明了通过蒸馏进行单独优化后的后期层逻辑值，其收敛损失可以与最终层相当，并能解决互补的一系列问题。
• MoDE 框架： 一种简单、低开销的微调机制，通过学习到的路由器对后期层逻辑值进行集成。它增加的训练参数极少（约为模型大小的 +0.04%），同时能够以相当的性能替代显著更大的可训练模块（如全层 LoRA）。
• 效率与灵活性： 该框架支持稀疏路由（Top-K），从而实现早期退出和高达 1.6$\times$ 的生成加速，且精度损失极小。它能与现有的 PEFT 方法结合使用，提升其性能。

实验结果
实验在 LLaMA2-7B 和 LLaMA3-8B 上针对算术推理、常识推理和指令遵循任务进行了测试。

• 算术推理： MoDE 持续提升了基础方法（Lo-RA, DoRA）的性能。轻量级变体（将 LoRA 应用于除最后 $k$ 层以外的所有层，并结合 MoDE）在使用的参数显著减少的情况下，实现了与全层 LoRA 仅 0.3–0.5 点差距的性能。
• 常识推理： MoDE 在基于权重的算法（LoRA, DoRA）之上提供了增益。虽然基于表示的方法（ReFT）本身表现强劲，但 MoDE 在与基于权重的方法结合时提供了可靠的改进。
• 指令遵循： MoDE 在 MT-Bench 上带来了微小但一致的提升，尽管在指令遵循任务上的改进不如推理任务那么显著，这表明指令遵循高度依赖于后期层的格式化处理，而这种集成仅部分捕捉到了这一点。
• 消融研究与分析：
  • 路由器自然地将权重集中在最后的一两个层上，但蒸馏和归一化组件使得较早的后期层也能有效地做出贡献。
  • 集成规模 $k \approx 3\text{--}4$ 是最优的；较大的 $k$ 有助于简单数据集，但会损害复杂的推理任务。
  • 稀疏 Top-K 路由在保持大部分精度的同时，实现了显著的加速。

意义与主张
作者将 MoDE 视为一种机制研究，证明了后期层逻辑值可以通过简单的适配（归一化和蒸馏）变得直接可用。论文声称：

1. 后期层表示拥有显著的预测潜力，这种潜力可以通过简单的适配（归一化和蒸馏）来释放。
2. 这些中间信号可以替代显著更大的可训练模块（如全层 LoRA），并提供相当的性能，从而成为一种参数高效的替代方案。
3. MoDE 提供了一个互补的方向，将中间层表示转化为可操作的训练信号，这与推理时的探测或基于权重的编辑方法截然不同。

论文对其范围保持了审慎的态度，指出了局限性，例如依赖于固定的 $k$（而非动态选择）、在指令遵循任务上的增益较小，以及评估仅限于解码器（decoder-only）架构的 LLM。未来的工作建议包括动态 $k$ 选择、更广泛的模型族以及更大的模型。