💬 NLP

Experience-Driven Dynamic Exits for LLMs with Reinforcement Learning

该论文提出了 LEDE，一个利用离线强化学习来动态优化大语言模型退出层和推测长度的框架，实现了相比于标准自回归解码和静态推测基准线的显著推理加速。

原作者： Yanyu Zhu, Hoilam Pao, Niu Hu, Wei Guo, Shaoxiong Zhan, Boyu Lai, Zitai Wang, Yongqin Zeng, Hai-Tao Zheng

发布于 2026-06-03

📖 1 分钟阅读☕ 轻松阅读

原作者： Yanyu Zhu, Hoilam Pao, Niu Hu, Wei Guo, Shaoxiong Zhan, Boyu Lai, Zitai Wang, Yongqin Zeng, Hai-Tao Zheng

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在读一本由超级智能 AI 编写的非常长且复杂的书。每当这个 AI 想要写下一个词时，它都必须进行一场大规模的脑力马拉松，检查其内部“大脑”中的每一个章节，以确保这个词是完美的。这就是当前 AI 模型的工作方式：它为每一个词都要经历整个过程，这使得它们既慢又耗能。

这篇论文介绍了一种名为 LEDE（基于学习的动态退出，Learning-based Dynamic Exit）的新系统，它就像是一个聪明的教练，教导 AI 如何在不损失准确性的情况下寻找捷径。

以下是它的工作原理，使用简单的类比：

问题：“一刀切”的马拉松

目前的 AI 模型使用一种叫做投机采样解码（Speculative Decoding）的方法。你可以把这想象成 AI 有一个“草稿助手”（一个更小、更快的自身版本）来猜测接下来的几个词。然后，主 AI 会检查这些猜测。

然而，目前这种方式是非常僵化的。这就像一位教练在对助手说：“无论句子是什么，永远只猜 4 个词，并且永远在检查完大脑第 5 层后停止。”

问题在于： 有些词很简单（比如“的”或“和”）。AI 不需要通过马拉松来预测它们。而有些词很难（比如复杂的数学或编程）。AI 需要进行深度思考。
结果： 这种僵化的系统在处理简单的词时浪费了能量，而在处理难词时有时又过于匆忙，从而导致错误或错失加速的机会。

解决方案：LEDE 的“聪明教练”

LEDE 改变了游戏规则，它使用强化学习（一种通过试错来学习的 AI 训练方式，就像狗学习耍俏来换取零食一样）。

LEDE 不再使用固定规则，而是训练一个“聪明教练”（一个智能体）来进行实时决策。以下是类比：

状态（检查点）： 当 AI 编写一个词时，它会经过大脑的不同层级。在每一层，聪明教练都会观察 AI 的“信心”。
- 类比： 想象 AI 正在爬一座山。在每一步，教练都会问：“你确定知道前方的路吗？”如果 AI 非常自信，教练会说：“太棒了，就在这里停下！”如果 AI 很困惑，教练会说：“继续向上爬，以获得更好的视野。”
动作（决策）： 教练在每一步有两个选择：
- 退出（Exit）： “停在这里！我们已经有足够的信息来猜测这个词了。”（这节省了时间）。
- 继续（Continue）： “继续深入大脑，以获得更好的猜测。”（这确保了准确性）。
奖励（零食）： 教练通过获得分数来学习。
- 如果它提前停止且猜测正确，它会获得大奖（因为它节省了时间）。
- 如果它提前停止且猜测错误，它会受到惩罚（因为它浪费了时间去修复错误）。
- 如果它在不需要的时候继续前进，它会受到小额惩罚（因为它浪费了能量）。

它是如何学习的（离线训练）

在 AI 真正与人类交流之前，聪明教练会在模拟器中进行练习。它会运行数千个示例，尝试不同的策略。

它尝试提前停止，然后晚一点停止，然后再早一点停止。
它保留了一个“笔记本”（经验回放池/Replay Buffer），记录了哪些做法有效，哪些无效。
随着时间的推移，它学会了一个完美的策略：“对于简单的句子，在第 3 层停止。对于困难的编程任务，要走到第 8 层。”

结果：为 AI 加速

论文在多个 AI 模型（如 Llama-2 和 Llama-3）上测试了 LEDE，发现它比旧的僵化方法快得多。

速度： 它比标准的缓慢方法快了 2.0 到 2.7 倍。
效率： 即使与使用固定规则的其他“聪明”方法相比，LEDE 也快了 17%。
质量： AI 并没有犯更多的错误；它只是学会了何时停止思考并开始写作。

总结

把旧的 AI 想象成一个在回答像“2+2 等于几？”这样简单问题之前，必须读完教科书每一页的学生。
LEDE 则像是一个已经学会识别问题的学生：“噢，这是一个简单的问题，我立刻就知道答案，所以我直接写下来就行了。”但如果问题很难，这个学生知道要读完整个章节。

通过教会 AI 变得灵活——明确知道何时停止以及何时继续——LEDE 让大型语言模型在不需要改变其大脑结构本身的情况下，变得显著更快、更高效。

技术摘要：LEDE —— 基于强化学习的经验驱动型 LLM 动态退出机制

1. 问题陈述

大语言模型（LLMs）在自回归推理过程中面临显著的延迟问题，因为生成的每个 token 都必须遍历整个 Transformer 层堆栈。虽然自投机解码（Self-Speculative Decoding, SSD）已被提出，通过利用模型自身的早期层作为内部“草拟器（drafter）”来并行生成候选 token 以加速这一过程，但目前的 SSD 方法在本质上受限于静态配置。

现有的方法依赖于固定的草拟深度和应用于整个生成过程的固定投机长度。这种僵化性未能考虑到上下文稀疏性（contextual sparsity），即预测 token 的难度会根据局部上下文的变化而变化。因此，静态方法无法在草拟质量（接受率）与投ک式速度之间实现最优平衡，导致推理性能并非最优。

2. 方法论：LEDE 框架

作者提出了 LEDE（基于学习的动态退出，Learning-based Dynamic Exit），该框架将 SSD 的动态控制重新定义为一个 马尔可夫决策过程（MDP）。LEDE 不再使用僵化的预定义规则，而是采用**离线强化学习（Offline RL）**来训练一个策略，该策略能够根据生成序列的局部上下文，动态选择最优的退出层和投机长度。

2.1. 马尔可夫决策过程建模

该框架将单步生成过程中的决策过程建模为一个具有以下组件的 MDP：

状态空间 ( $S$ )： 智能体在每个候选层 $l$ $l$ 观察一个状态向量 $s_l$ $s_{l}$ 。该向量源自模型的内部 token 分布，包括：
- Token 置信度 ( $C_k^l$ )： 前 $k$ 个 token 的平均负对数概率。
- Token 熵 ( $H^l$ )： 下一个 token 预测分布的不确定性。
- Token 置信度平均轨迹 ( $C_{trace}$ ： 在滑动窗口内前一步 top-1 token 的平均负对数概率。
动作空间 ( $A$ )： 在每个候选层，策略（Q 网络）在两个动作之间做出决策：
- 继续 ( $a=0$ )： 进入下一个候选层。
- 退出 ( $a=1$ )： 将当前层指定为草拟阶段的最优退出层 ( $l^*$ )。
奖励函数 ( $R$ )： 旨在优化推理效率。
- 正向奖励： 如果草拟 token 被接受（ $x_t^{l^*} = x_t^L$ ），则给予奖励，并根据节省的深度进行缩放（ $1 + \frac{L-l^*}{L}$ ）。
- 负向奖励： 如果草拟 token 被拒绝，则给予奖励（ $-2 + \frac{L-l^*}{L}$ ）。
- 惩罚： 对“继续”动作应用微小的惩罚（$-0.01$），以鼓励在适当的情况下进行早期退出。

2.2. 通过离线 RL 进行策略优化

LEDE 使用 深度 Q 网络（DQN），通过过去推理经验的经验回放池进行离线训练：

数据收集： 初始探索阶段收集多样化的经验元组 $(s_t, a_t, r_t, s_{t+1})$ ，其中智能体与环境进行交互。
离线训练： 通过最小化 n 步时序差分（TD）误差 来训练 Q 网络。目标值使用稳定的目标网络计算，以平衡立即退出带来的奖励与状态的长期价值。
双重控制机制： 在在线推理期间，训练好的智能体执行两种自适应控制：
- 动态退出层选择： 它确定开始草拟的最优深度 $l^*$ 。
- 自适应草拟长度： 选择 $l^*$ 后，智能体在每个草拟 token 后重新评估状态。一个“退出”动作表示继续草拟（表明高置信度），而一个“继续”动作（表明低置信度）则立即终止草拟以进行验证。

3. 核心贡献

本文概述了三个主要贡献：

MDP 建模： 这是首个将 SSD 中草拟深度和投机长度的动态控制建模为 MDP 问题，并应用强化学习来推导控制策略的工作。
LEDE 框架： 设计并实现了一个利用离线 RL 和经验回放来学习上下文感知策略的框架，克服了静态配置和简单启发式方法的局限性。
实验验证： 通过综合实验证明了其相对于自回归生成的显著加速，以及相对于静态基准线的显著效率提升。

4. 实验结果

该框架在 Llama-2 (7B, 13B) 和 Llama-3 (1B) 模型，以及 CodeLLaMA (7B, 34B) 上进行了评估，涵盖了指令遵循、摘要生成、语言建模和代码生成等多种任务。

加速性能：
- LEDE 相比标准自回归解码实现了平均 2.32× 的加速。
- 峰值加速达到 2.72×（例如在 Llama-2-7B 语言建模任务上）。
- LEDE 比表现最好的静态投机基准（LayerSkip）额外提升了 17% 的速度。
与基准对比：
- LayerSkip (静态)： 性能对固定配置高度敏感，且难以泛化到不同任务。
- 基于规则的方法 (LITE, Draft & Verify)： 受限于预定义的阈值，导致深度与接受率之间的权衡并非最优。
- LEDE： 通过学习细粒度的控制策略，它在保持相当的生成质量（通过 Rouge-L 衡量）的同时，在加速方面始终优于所有基准模型。
消融研究： 移除动态退出层选择或自适应草拟长度中的任何一个，都会显著降低性能（加速降至约 2.0×），证实了这两种机制具有协同作用且至关重要。

5. 重要性与主张

论文声称 LEDE 代表了从僵化的、启发式驱动的推理向自适应的、策略驱动的推理的概念转变。通过学习根据实时上下文共同优化计算深度和投机长度，该框架能够在多种模型和任务中实现更稳健、更高效的解码。

作者将这项工作定位为构建更智能、计算效率更高的 LLM 系统的一步。他们指出，未来的工作将侧重于将这些优势扩展到更大规模的模型（70B 及以上），但目前的贡献确立了使用离线强化学习来动态管理大型语言模型推理流水线的可行性。