自监督早期退出机制加速大语言模型推理

背景

大型语言模型（LLM）在实际部署中面临严重的推理延迟和计算成本问题。随着模型参数规模从数亿增长到数千亿，逐 token 自回归生成的计算量线性累积，使得即使单次推理的延迟也可能达到秒级甚至分钟级。传统加速方法包括模型量化、剪枝、知识蒸馏，以及近年来流行的推测解码（speculative decoding）——使用一个轻量级草稿模型快速生成候选 token，再由大模型验证。然而，草稿模型需要额外训练或调参，且与主模型之间存在分布不匹配。本文提出一种模块化、自监督的早期退出（early exit）方法，直接在预训练 LLM 的中间层添加可训练的退出头，让模型自身在合适时机提前终止计算，从而大幅降低推理成本。

核心内容

论文的核心思想是在 LLM 的每个中间 Transformer 层插入一个早期退出头部（early exit head）。每个头部被训练为以自监督方式模仿主模型（完整模型）的最终预测分布——即最小化该层输出的预测 logits 与完整模型最后一层 logits 之间的 KL 散度。训练时，主模型参数冻结，仅更新退出头部的参数。推理时，模型按顺序逐层计算，每经过一个早期退出头，系统计算该头部输出的置信度（confidence）。一旦某个头部输出的置信度超过预设的校准阈值，就立即停止后续层的计算，直接将该头部的预测作为最终输出。

为了确定最可靠的置信度度量，作者比较了多种指标，包括预测概率最大值、概率差值（top-1 与 top-2 之差）、熵（entropy）以及基于温度缩放的校准分数。实验表明，熵在区分正确与错误预测方面表现最稳健：正确预测通常对应较低的熵值（分布集中），而错误预测则对应较高的熵值（分布分散）。通过设定熵阈值，可以在准确率几乎不下降的前提下实现显著的早期退出率。

实验基于 Pythia 模型套件（参数范围为 70M 至 2.8B），在多个文本生成基准（如 LAMBADA、HellaSwag、WinoGrande）上进行评估。结果显示：

• 对于 2.8B 参数的模型，使用早期退出后，平均仅需计算约 30%~50% 的层即可达到与完整模型几乎相同的准确率。
• 推理速度提升约 1.5~2.5 倍（取决于阈值设定和模型大小），同时准确率损失控制在 0.5% 以内。

论文进一步将早期退出思想与推测解码结合，提出了动态自推测解码（Dynamic Self-Speculative Decoding, DSSD）。传统推测解码需要单独训练一个草稿模型（如 LayerSkip 方法使用浅层模型），而 DSSD 直接利用早期退出头部作为草稿模型：在每一步，用早期退出头快速生成多个候选 token，再使用完整模型进行并行验证。与传统方案相比，DSSD 无需额外训练草稿模型，且能自适应地选择退出层作为草稿来源。实验显示，在最小化超参数调优的情况下，DSSD 的 token 接受率比手动调优的 LayerSkip 基线高出 1.66 倍。

关键要点

• 早期退出头部：在每个中间 Transformer 层后添加可训练分类头，以 KL 散度损失自监督学习模仿完整模型的最终预测分布，主模型参数冻结。
• 置信度度量：比较了最大值、差值、熵等指标，熵（entropy）被证实是区分正确与错误预测最可靠的信号——正确预测熵低，错误预测熵高。
• 模块化与零额外延迟：训练完成后，推理时仅需在每个头部计算一次 logits 和熵值，开销极低；提前退出可跳过后续所有层计算。
• 实验结果：在 Pythia 70M~2.8B 模型上，推理速度提升 1.5~2.5 倍，准确率下降 <0.5%；小模型受益更明显，因为退出层比例更高。
• DSSD 方法：将早期退出头部作为推测解码的草稿模型，无需额外训练；自动选择最优退出层进行草稿生成，token 接受率比 LayerSkip 高 1.66 倍。
• 超参数简易：DSSD 仅需设置一个置信度阈值，无需像 LayerSkip 那样手动调整草稿模型层数等超参数。

意义与影响

该工作为 LLM 推理加速提供了一种轻量级、无外部依赖的解决方案。传统早期退出方法多用于分类任务或小型模型，本文首次系统验证了在大型自回归语言模型（70M~2.8B）上自监督早期退出的有效性，并提出了实用化的置信度校准方案（熵阈值）。将早期退出与推测解码无缝结合的 DSSD，消除了训练额外草稿模型的需求，降低了部署复杂度，同时显著提升了 token 接受率。

从工程角度看，该方法可与现成的预训练模型（如 Pythia、Llama 系列）配合使用，仅在推理流水线中插入少量参数（每个退出头约等于一层 MLP 的参数量），对显存和计算资源的增加几乎可忽略。未来可进一步探索动态退出策略（如根据输入难度自适应调整阈值）、多层投票机制，以及扩展到更大规模（10B+）模型的效果。此外，该方法的自监督性质使其适用于任何仅需输出 logits 的 LLM，为加速在线推理服务提供了新的思路。