Experience-Driven Dynamic Exits for LLMs with Reinforcement Learning

背景

大型语言模型（LLMs）在推理阶段面临的一个主要瓶颈是自回归（autoregressive）生成的速度缓慢。为了加速这一过程，研究人员提出了多种优化策略，其中**自推测解码（Self-Speculative Decoding）**是一种备受关注的技术。

自推测解码的核心思想是利用一个较小的“草稿模型”并行生成多个 token，然后由主模型进行验证。如果验证通过，这些 token 可以被一次性接受，从而显著减少主模型的迭代次数。然而，现有的自推测解码方法通常依赖于静态配置，例如固定的退出层（exit layers）和固定的推测长度（speculation lengths）。这种“一刀切”的配置无法适应不同输入序列的复杂性差异，导致在简单任务上浪费算力，或在复杂任务上推测失败率高，从而限制了整体效率的提升。

核心内容

为了解决上述静态配置的局限性，本文提出了一种名为 LEDE（Experience-Driven Dynamic Exits）的新框架。该框架将推理加速问题重新定义为马尔可夫决策过程（Markov Decision Process, MDP），并采用**离线强化学习（Offline Reinforcement Learning）**来训练一个智能策略。

1. 问题重构：从静态到动态

LEDE 不再使用固定的参数，而是让模型根据当前生成序列的局部上下文（local context），在每一步动态选择两个关键参数：

• 退出层（Exit Layer）：决定在 Transformer 网络的哪一层输出结果。较浅的层计算成本低但语义理解可能不足，较深的层则相反。
• 推测长度（Speculation Length）：决定草稿模型一次性生成多少个 token。

2. 核心机制：离线强化学习策略

LEDE 通过离线强化学习训练一个策略网络（Policy Network）。该策略网络的目标是平衡计算成本与草稿质量：

• 状态（State）：当前生成的序列局部上下文。
• 动作（Action）：选择退出层和推测长度。
• 奖励（Reward）：基于推理速度提升和生成质量保持（如困惑度或准确率）的综合指标。

通过这种方式，LEDE 能够根据输入内容的难易程度和当前生成状态，自适应地调整推理路径。对于简单或已确定的上下文，模型可以选择浅层退出和较短的推测以节省时间；对于复杂或模糊的上下文，则选择深层退出和更长的推测以保证准确性。

3. 实验评估

研究者在 Llama-2 和 Llama-3 模型上进行了全面的评估。实验结果表明：

• 与传统的自回归解码相比，LEDE 实现了 $2.0\times \sim 2.7\times$ 的加速。
• 与基于静态配置的推测解码基线相比，LEDE 提供了额外的 17% 的速度提升。

这些结果证明了动态调整退出层和推测长度在平衡效率与质量方面的有效性。

关键要点

• 动态优化而非静态配置：LEDE 摒弃了固定的退出层和推测长度，转而根据生成序列的局部上下文动态选择最优参数。
• 强化学习驱动：将推理加速建模为马尔可夫决策过程，利用离线强化学习训练策略网络，以平衡计算成本与草稿质量。
• 双重动态控制：策略同时控制“何时退出”（退出层）和“推测多少”（推测长度），实现了更细粒度的资源分配。
• 显著的性能提升：在 Llama-2 和 Llama-3 模型上，LEDE 相比自回归解码加速 2.0 至 2.7 倍，相比静态推测基线额外加速 17%。
• 解决自推测解码的痛点：有效克服了传统自推测解码因静态配置导致的效率瓶颈，提升了其在不同复杂度任务上的适应性。

意义与影响

LEDE 框架为大型语言模型的推理加速提供了一种新的范式。它表明，通过引入强化学习来实现推理过程的动态化，可以显著突破传统静态优化方法的性能上限。

1. 提升部署效率：对于需要低延迟响应的实际应用（如实时对话、代码生成等），LEDE 能够在不牺牲生成质量的前提下大幅降低推理延迟，提高单位算力的吞吐量。
2. 资源自适应：动态退出机制使得模型能够根据任务难度灵活分配计算资源，避免在简单任务上过度计算，或在复杂任务上计算不足，从而优化整体资源利用率。
3. 推动动态推理研究：本文将强化学习应用于动态退出和推测长度的联合优化，为后续研究提供了可借鉴的 MDP 建模思路，可能启发更多基于动态控制的 LLM 推理优化方法。

总之，LEDE 展示了通过数据驱动的策略学习来实现高效、自适应的 LLM 推理的潜力，是迈向更智能、更高效的大模型部署的重要一步。