基于强化学习的经验驱动式LLM动态退出策略

# 基于经验驱动的强化学习大语言模型动态退出
来源: https://arxiv.org/html/2606.03113
###### 摘要

大语言模型存在自回归推理缓慢的问题。虽然自推测解码能加速这一过程，但其效率受限于固定配置（如固定的退出层和推测长度）。我们将此优化问题重新表述为马尔可夫决策过程，并提出 LEDE，一个基于离线强化学习的框架。LEDE 学习一个策略，该策略根据生成序列每一步的局部上下文动态选择最优退出层和推测长度，从而平衡计算开销与草稿质量。在 Llama-2 和 Llama-3 模型上的全面评估表明，LEDE 相比自回归解码可实现高达 2.0× 到约 2.7× 的加速，并且比静态推测基线额外提供 17% 的加速。

索引词—LLM 推理加速, 推测解码, 动态提前退出, 强化学习

参见说明

图 1：LEDE 框架架构。该框架包含两个阶段：(a) 离线训练阶段，其中 Q 网络从经验回放缓冲区学习退出策略；以及在线推理阶段 (b-c)。在推理过程中，训练好的智能体首先选择最优退出层 \(l^*\) (b)，然后自适应地确定该层的草稿长度 (c)。

## 1 引言

大语言模型 (LLMs) [1 (https://arxiv.org/html/2606.03113#bib.bib1),2 (https://arxiv.org/html/2606.03113#bib.bib2),3 (https://arxiv.org/html/2606.03113#bib.bib3)] 已变得越来越深，以处理包括问答、摘要、编码和数学推理在内的各种任务。虽然这种深度增强了模型能力，但它也增加了自回归解码过程中的显著延迟，因为每个词元都必须经过完整的 Transformer 层堆栈。为缓解这一低效问题，自推测解码 (SSD) 被提出，它利用模型的浅层作为内部“起草者”，以并行方式生成候选词元序列 [4 (https://arxiv.org/html/2606.03113#bib.bib4)]。

然而，当前 SSD 方法的效率从根本上受到其静态配置的限制——在整个生成过程中统一应用固定的草稿深度和推测长度 [5 (https://arxiv.org/html/2606.03113#bib.bib5)]。这种僵化方法与 LLM 推理中上下文稀疏性的存在 [6 (https://arxiv.org/html/2606.03113#bib.bib6),7 (https://arxiv.org/html/2606.03113#bib.bib7),8 (https://arxiv.org/html/2606.03113#bib.bib8),9 (https://arxiv.org/html/2606.03113#bib.bib9)] 从根本上不匹配——并非所有词元都同样难以预测。因此，现有方法（包括使用简单启发式的方法 [10 (https://arxiv.org/html/2606.03113#bib.bib10),11 (https://arxiv.org/html/2606.03113#bib.bib11),12 (https://arxiv.org/html/2606.03113#bib.bib12)]）无法在草稿质量和推测速度之间取得最优平衡，导致次优性能。

为解决这一局限性，我们引入了 LEDE（基于学习的动态退出），这是一个将这种动态控制挑战表述为马尔可夫决策过程 (MDP) [13 (https://arxiv.org/html/2606.03113#bib.bib13)] 的框架。具体来说，LEDE 应用离线强化学习 (RL) [14 (https://arxiv.org/html/2606.03113#bib.bib14)] 来训练一个策略，该策略动态控制草稿深度和推测长度。在每一步，智能体观察模型的内部状态以评估生成难度，并根据其学到的策略执行动作。这种基于策略的方法用复杂、自适应的策略替代了僵化的预定义规则，该策略能从过去的推理经验中泛化，从而实现更稳健、更高效的解码。

我们的贡献主要有三点：(1) 我们首次将 SSD 中草稿深度和推测长度的动态控制问题表述为 MDP，并应用 RL 推导出控制策略。(2) 我们设计并实现了 LEDE，一个利用离线 RL 和经验回放来学习其上下文感知策略的框架，克服了静态和基于简单启发式配置的局限性。(3) 我们的实验表明，相比自回归生成可获得高达 \(\bm{2.7\times}\) 的显著加速，并且相比静态基线可获得可观的效率提升（最高 17%）。

## 2 相关工作

### 2.1 动态计算方法

为降低推理的高昂成本，动态计算方法自适应地调整分配给每个词元的计算量。这通常通过提前退出（从中间层预测词元）或完全跳过层来实现。开创性工作如 CALM 和 Mixture-of-Depths 根据词元难度动态改变模型深度 [7 (https://arxiv.org/html/2606.03113#bib.bib7),15 (https://arxiv.org/html/2606.03113#bib.bib15)]。其他方法探索统一的层跳过策略 [16 (https://arxiv.org/html/2606.03113#bib.bib16)]、仅激活必要子网络的上下文稀疏性 [9 (https://arxiv.org/html/2606.03113#bib.bib9),17 (https://arxiv.org/html/2606.03113#bib.bib17),18 (https://arxiv.org/html/2606.03113#bib.bib18)]，或使用 k-NN 搜索优化层选择 [19 (https://arxiv.org/html/2606.03113#bib.bib19)]。这些方法共享一个共同目标：避免对“简单”词元进行冗余计算。

### 2.2 自推测方法

推测解码 (SD) 通过使用更快、更小的草稿模型生成词元，再由更大的目标模型并行验证，从而加速推理 [20 (https://arxiv.org/html/2606.03113#bib.bib20),21 (https://arxiv.org/html/2606.03113#bib.bib21),22 (https://arxiv.org/html/2606.03113#bib.bib22),23 (https://arxiv.org/html/2606.03113#bib.bib23)]。一种非常有效的变体——SSD，使用目标模型自身的浅层作为起草者，从而消除了对外部模型的需求 [4 (https://arxiv.org/html/2606.03113#bib.bib4)]。为了更好地对齐中间层与最后一层的预测以提高接受率，一些工作 [12 (https://arxiv.org/html/2606.03113#bib.bib12),7 (https://arxiv.org/html/2606.03113#bib.bib7)] 训练了语言头适配器，以便从中间深度进行词元预测。LayerSkip [5 (https://arxiv.org/html/2606.03113#bib.bib5)] 采用了一种训练方案，结合了层丢弃和加权提前退出损失，以获得稳健的提前退出能力。我们的工作建立在这一基础之上，但引入了基于学习的策略来动态控制起草过程。

## 3 LEDE 框架

我们引入了基于学习的动态退出，这是一个用学到的自适应策略替代静态启发式来控制 SSD 过程的框架。如图 1 (https://arxiv.org/html/2606.03113#S0.F1) 所示，我们的方法包括两个阶段：一个离线学习过程，其中提前退出智能体从过去推理经验库中训练（图 1 (https://arxiv.org/html/2606.03113#S0.F1)a）；以及一个在线推理过程（图 1 (https://arxiv.org/html/2606.03113#S0.F1)b-c），其中训练好的智能体被部署用于做出上下文感知的退出决策。以下小节详细描述该框架的每个组件。

### 3.1 马尔可夫决策过程

我们将退出层 \(l^*\) 的动态选择问题表述为 MDP 框架。与选择固定退出层的先前工作不同，我们的智能体在单个生成步骤的前向传播过程中做出一系列决策，以确定起草的最优深度。

#### 3.1.1 状态空间

智能体观察到的状态是动态的，并依赖于计算深度。在单个生成步骤的前向传播过程中，当到达候选退出层 \(l\) 时，我们提取一个状态向量，记为 \(s_l\)。该向量由从模型在第 \(l\) 层的内部词元分布中导出的几个特征组成，作为捕获模型在该特定深度下的置信度和不确定性的紧凑内在信号。这些特征源自下一刻词元概率分布 \(P_t^{(l)}\)，该分布是从第 \(l\) 层的隐藏状态计算得到的。

词元置信度 \(C_k^l\)。定义为来自 \(P_t^{(l)}\) 的 top-k 词元的平均负对数概率：
\[
C_k^l = -\frac{1}{k}\sum_{j=1}^k \log P_t^{(l)}(j)
\]
其中 \(P_t^{(l)}(j)\) 是 top-k 预测的概率，在我们的实验中 \(k=5\)。

词元熵 \(H^l\)。衡量第 \(l\) 层下一刻词元预测分布的不确定性，计算如下：
\[
H^l = -\sum_{j \in V} P_t^{(l)}(j) \log P_t^{(l)}(j)
\]
其中 \(V\) 表示词表集。

词元置信度平均迹 \(C_{trace}\)。定义为窗口内 top-1 预测词元的平均负对数概率：
\[
C_{trace} = -\frac{1}{T}\sum_{k=1}^T \log P_{t-k}
\]
其中 \(T\) 表示窗口大小，\(P_{t-k}\) 是前一步 \(t-k\) 输出分布中的最大概率，在我们的实验中 \(T=20\)。

#### 3.1.2 动作空间

在每个候选层 \(l\)，策略 Q 网络将状态 \(s_t\) 作为输入，并决定动作 \(a_t \in \{\text{continue}, \text{exit}\}\)。“exit”动作 (\(a_t=1\)) 将当前层 \(l\) 指定为起草阶段选定的退出层 \(l^*\)。“continue”动作 (\(a_t=0\)) 则前进到集合中的下一个候选层。如果智能体到达最后一个候选层，则强制退出。

#### 3.1.3 奖励函数

奖励函数 \(R\) 旨在直接优化推理效率。对于在 \(l^*\) 处每次退出，如果草稿词元被接受，我们给予正奖励以鼓励正确性和效率。相反，如果草稿词元被拒绝，我们给予负奖励。此外，我们还给予一个小惩罚以激励提前退出。奖励函数定义如下：
\[
R = \begin{cases}
1 + \frac{(L - l^*)}{L}, & \text{if } a_t = 1 \cap x_t^{l^*} = x_t^L, \\
-2 + \frac{(L - l^*)}{L}, & \text{if } a_t = 1 \cap x_t^{l^*} \neq x_t^L, \\
-0.01, & \text{if } a_t = 0.
\end{cases}
\]
其中 \(L\) 表示模型的总层数，\(x_t^l\) 表示在第 \(l\) 层预测的词元。

### 3.2 基于离线 RL 的策略优化

我们的目标是训练一个深度 Q 网络 (DQN) [24 (https://arxiv.org/html/2606.03113#bib.bib24)] 来优化退出决策。该框架的关键组件包括：

目标 Q 网络。我们使用一个小型 MLP 作为 Q 网络。它以状态向量 \(s_t\) 作为输入，并输出对应“continue”和“exit”动作的 Q 值 \(Q(s_t, a_t)\)。然后基于最大 Q 值贪心地选择当前层的动作。

经验收集与回放。如图 1 (https://arxiv.org/html/2606.03113#S0.F1)a 所示，在初始探索阶段，策略智能体不断与环境交互，并收集多样化的经验。每次层级别前向传播中的决策都作为经验元组 \((s_t, a_t, r_t, s_{t+1})\) 存储在回放缓冲区中。这里，\(s_t\) 是候选层 \(l\) 处的上下文状态，\(a_t\) 是执行的动作（“continue”或“exit”），\(s_{t+1}\) 是下一个候选层的状态。奖励 \(r_t\) 依据 3.1.3 (https://arxiv.org/html/2606.03113#S3.SS1.SSS3) 中定义的奖励函数分配。这个经验回放缓冲区为离线策略优化提供了全面的数据集。

离线策略更新。Q 网络通过从回放缓冲区中采样小批量经验进行离线训练。策略优化通过最小化 n 步时序差分 (TD) 误差 [25 (https://arxiv.org/html/2606.03113#bib.bib25)] 来实现。这通过最小化预测 Q 值与目标值 \(y_t\) 之间的均方误差来完成。损失函数定义如下：
\[
L(\theta) = \mathbb{E}\left[ (y_t^{(n)} - Q_\theta(s_t, a_t))^2 \right]
\]
其中 n 步目标回报 \(y_t^{(n)}\) 是 n 步的折扣奖励之和，加上 n 步后到达状态的价值（由稳定的目标网络 \(Q_{\bar{\theta}}\) 估计）：
\[
y_t^{(n)} = \sum_{k=0}^{n-1} \gamma^k r_{t+k+1} + \gamma^n \max_{a'} Q_{\bar{\theta}}(s_{t+n}, a')
\]
其中 \(\gamma\) 是折扣因子，\(a'\) 是来自目标网络 \(Q_{\bar{\theta}}\) 的贪心动作。在我们的实验中 \(\gamma=0.99\)，\(n=3\)。

### 3.3 动态退出层与自适应起草

我们训练好的策略智能体同时动态控制退出层和草稿长度，如我们框架的在线推理阶段所示。首先，智能体通过在初始前向传播中进行层级别评估，做出“exit”决策来选择最优退出层 \(l^*\)（图 1 (https://arxiv.org/html/2606.03113#S0.F1)b）。随后，它从 \(l^*\) 开始生成草稿词元。每次起草后，智能体重新评估状态：“exit”动作表示高置信度并继续起草过程，而“continue”动作表示低置信度并立即终止起草以进行验证（图 1 (https://arxiv.org/html/2606.03113#S0.F1)c）。这种双控制机制协同优化计算深度和推测长度，以最大化推理速度。

表 1：LEDE 与基线方法在 LLaMA2 和 LLaMA3 模型上的多项任务中的比较。加速比相对于自回归解码。最佳结果以**粗体**表示，次佳结果以下划线表示。

| 模型 | 方法 | CNN/DM（摘要） | CNN/DM（语言） | TOPv2（指令） | Alpaca（指令） |
|------|------|----------------|----------------|---------------|----------------|
|      |      | 接受率 | 加速比 | R-L | 接受率 | 加速比 | R-L | 接受率 | 加速比 | R-L | 接受率 | 加速比 | R-L |
| LLaMA-3.2-1B | AR | - | - | 1.00×\(1.00\times\) | 0.117 | - | - | 1.00×\(1.00\times\) | 0.108 | - | - | 1.00×\(1.00\times\) | 0.068 | - | - | 1.00×\(1.00\times\) | 0.153 |
| | LS | 4.00 | 4.00 | 0.632 | 1.60×\(1.60\times\) | 0.113 | 4.00 | 4.00 | 0.740 | 1.81×\(1.81\times\) | 0.108 | 4.00 | 4.00 | 0.771 | 1.86×\(1.86\times\) | 0.068 | 4.00 | 4.00 | 0.751 | 1.73×\(1.73\times\) | 0.153 |
| | LITE | 10.36 | 11.83 | 0.922 | 1.07×\(1.07\times\) | 0.117 | 9.77 | 11.63 | 0.844 | 1.30×\(1.30\times\) | 0.108 | 9.27 | 11.76 | 0.863 | 1.45×\(1.45\times\) | 0.680 | 10.07 | 11.59 | 0.849 | 1.09×\(1.09\times\) | 0.155 |
| | DV | 6.00 | 5.61 | 0.882 | 1.44×\(1.44\times\) | 0.116 | 6.00 | 7.26 | 0.868 | 1.65×\(1.65\times\) | 0.108 | 3.00 | 6.94 | 0.905 | 1.78×\(1.78\times\) | 0.068 | 6.00 | 6.05 | 0.867 | 1.23×\(1.23\times\) | 0.156 |
| | **LEDE** | 5.70 | 6.84 | 0.924 | **2.04×\(2.04\times\)** | 0.116 | 4.62 | 6.40 | 0.881 | **2.28×\(2.28\times\)** | 0.107 | 6.88 | 6.40 | 0.911 | **1.98×\(1.98\times\)** | 0.068 | 3.96 | 4.70 | 0.867 | **2.04×\(2.04\times\)** | 0.156 |
| LLaMA-2-7B | AR | - | - | 1.00×\(1.00\times\) | 0.194 | - | - | 1.00×\(1.00\times\) | 0.225 | - | - | 1.00×\(1.00\times\) | 0.095 | - | - | 1.00×\(1.00\times\) | 0.207 |
| | LS | 4.00 | 4.00 | 0.254 | 1.22×\(1.22\times\) | 0.193 | 4.00 | 4.00 | 0.553 | 1.95×\(1.95\times\) | 0.225 | 4.00 | 4.00 | 0.732 | 2.47×\(2.47\times\) | 0.094 | 4.00 | 4.00 | 0.724 | 2.37×\(2.37\times\) | 0.206 |
| | LITE | 15.39 | 11.82 | 0.955 | 1.54×\(1.54\times\) | 0.193 | 19.17 | 11.88 | 0.806 | 1.32×\(1.32\times\) | 0.225 | 18.00 | 0.881 | 0.417 | 1.52×\(1.52\times\) | 0.094 | 18.49 | 11.83 | 0.881 | 1.37×\(1.37\times\) | 0.204 |
| | DV | 8.00 | 7.57 | 0.858 | 2.09×\(2.09\times\) | 0.194 | 8.00 | 5.60 | 0.865 | 2.13×\(2.13\times\) | 0.225 | 8.00 | 6.29 | 0.779 | 1.94×\(1.94\times\) | 0.096 | 7.00 | 7.37 | 0.854 | 2.14×\(2.14\times\) | 0.207 |
| | **LEDE** | 6.71 | 7.27 | 0.918 | **2.70×\(2.70\times\)** | 0.193 | 6.10 | 7.64 | 0.878 | **2.27×\(2.27\times\)** | 0.225 | 5.78 | 6.81 | 0.856 | **2.66×\(2.66\times\)** | 0.095 | 6.90 | 7.38 | 0.887 | **2.48×\(2.48\times\)** | 0.207 |