FlexEE：用于卸载感知LLM推理的自推测与KV兼容早期退出框架

# 自我推测与KV兼容的早期退出机制用于卸载感知的LLM推理
来源：https://arxiv.org/html/2609.17008
李子伟 单位：中国科学技术大学，合肥，中国  
康毅 单位：中国科学技术大学，合肥，中国  
单位：合肥综合国家科学中心人工智能研究所，合肥，中国  
邮箱：\{xieqihu, heinzt\}@mail.ustc.edu.cn，[email protected]

###### 摘要
大型语言模型（LLM）的推理常受计算与内存双重限制，尤其在基于卸载的部署中，模型权重需在自回归解码时跨内存层次传输。在此场景下，减少执行层数既能降低每token延迟，又可避免昂贵的权重搬运。基于此观察，我们提出FlexEE——一种面向资源受限及卸载环境的LLM推理早期退出框架。FlexEE通过层级退出监督实现可靠的中间层预测、基于Top-K本地词汇表的自我推测解码以实现低成本退出决策，以及动态隐状态管理以实现KV缓存正确且内存感知的执行，使LLM解码的早期退出成为实用技术。在生成与下游任务中，FlexEE以最小精度损失实现高效早期退出，在Llama2-7B与Llama3-8B上分别实现0%/50%权重卸载下最高1.27×/3.16×和1.25×/2.83×的端到端加速。

22脚注：通讯作者。## 1引言

基于Transformer的大型语言模型（LLM）在机器人（Zitkovich等人, 2023；Li等人, 2025；Black等人, 2024）、图像生成（Chen等人, 2024；Xiao等人, 2025）和代码生成（Ghaemi等人, 2024；Luo等人, 2025；Du等人, 2024）等多领域取得显著成功。其进步主要得益于经验缩放定律（Kaplan等人, 2020；Bian等人, 2025），该定律表明性能随模型规模、数据量和计算量增长而提升，推动现代LLM达到前所未有的参数规模。

然而，这种增长使模型在资源受限环境中部署面临挑战，因其参数量常超出片上内存容量。常见解决方案是将部分层参数卸载至片外内存，并在推理时按需加载（Rajbhandari等人, 2020；Rasley等人, 2020；Sheng等人, 2023；Kwon等人, 2023）。这使推理瓶颈从计算转向片外内存访问，加速性能关键取决于减少参数加载的规模与频率。

先前降低LLM推理硬件开销的研究涵盖快速解码与内核级优化（Dao等人, 2022；Shah等人, 2024；Dai等人, 2025）、量化（Dettmers等人, 2023；Frantar等人, 2022；Dettmers等人, 2022）、剪枝（Ma等人, 2023；Sun等人, 2023；Xia等人, 2023；Kwon等人, 2022；Frantar与Alistarh, 2023）以及KV缓存或上下文选择方法（Zhang等人, 2023；Li等人, 2024；Hao等人, 2025）。尽管这些技术减少了计算或内存使用，但许多依赖专用内核、硬件假设或独立缓存优化，可能限制其在异构边缘设备上的直接应用。相比之下，早期退出提供了一种互补且灵活的范式：通过在中间层动态终止计算，减少每个输入的执行层数。在基于卸载的推理中，这同时减少了从片外内存加载的层数，直接缓解数据移动瓶颈。

图1：（a）LLM早期退出的三大挑战；（b）CPU-GPU卸载下的早期退出：左图展示Llama2-13B的逐层计算与权重加载延迟；右图展示FlexEE在0%与50%层卸载下相对无早期退出基线的端到端加速比。
尽管早期退出理论上可减少不必要的计算，但它在算法与系统层面引入了现有方法未充分解决的挑战，如图1所示。

*在算法层面*，大多数早期退出方法（Fan等人, 2024；Xu等人, 2025）依赖将中间隐状态投影回完整词汇空间以评估置信度或预测稳定性。然而，LLM中的中间表示通常无法可靠解码，导致直接应用于现成模型时产生严重质量下降。此外，早期退出预测需通过大型语言建模头（LM头）将中间状态投影至高维词汇空间（如Llama2-7B中的32K词元），计算开销高且搜索困难，常抵消早期退出的潜在加速效果（Fan等人, 2024）。

*在系统层面*，早期退出会破坏高效自回归解码所必需的键值（KV）缓存构建。在中间层退出时，后续层无法生成未来词元所需的键值对。此前解决方案（Schuster等人, 2022；Chen等人, 2023）通过将隐状态复制至后续层以完成缓存更新，但该策略仍需加载所有剩余层的权重。在基于卸载的推理中，此行为抵消了早期退出的优势，因片外权重加载仍主导延迟。

图2：FlexEE框架。针对卸载层采用两阶段设计实现早期退出预测：先通过自我推测层缩减候选空间，再由早期退出层在缩减空间内操作，该设计基于三项统计洞察。
为应对上述挑战，我们提出FlexEE——一种基于系统-算法联合协同设计的卸载感知早期退出框架，如图2所示。FlexEE通过将退出行为与实际内存及执行特性对齐，实现权重卸载下的高效早期退出。我们的核心思想与贡献如下：

首先，我们系统分析LLM中间层，发现其语义解码能力远弱于最终层，限制了现有早期退出方法。基于此观察，我们设计了无需训练、基于自我推测的卸载层早期退出预测器。如图2c所示，其采用两阶段设计：自我推测层提取紧凑候选词汇表，早期退出层在此缩减空间内进行预测，在保持精度的同时降低预测开销。

其次，我们通过动态隐状态管理解决自回归早期退出中的KV缓存冲突，该方法保留并复用关键隐状态，无需强制复制或重构所有剩余层的KV缓存，即可维持跨词元语义连续性，使FlexEE更适合基于卸载的推理。

最后，我们在多项生成与下游任务中评估FlexEE。FlexEE在保持与密集解码几乎相当性能的同时实现有效早期退出。在Llama2-7B上，FlexEE在无卸载与50%权重卸载下分别实现最高1.27×和3.16×的端到端加速；在Llama3-8B上，相同设置下分别实现最高1.25×和2.83×的加速。

## 2相关工作

早期退出通过在达到足够置信度时自适应终止中间层计算以加速神经网络推理。早期方法如BranchyNet（Teerapittayanon等人, 2016）使用辅助分类器进行图像分类，而基于Transformer的方法如DeeBERT（Xin等人, 2020）和FastBERT（Liu等人, 2020）将该思想扩展至BERT类编码器，用于分类与理解任务。这些方法通常依赖额外预测头、联合训练以及基于熵或置信度的退出标准。近期研究探索了LLM与生成推理的早期退出：LayerSkip（Elhoushi等人, 2024）与EE-LLM（Chen等人, 2023）通过层级监督或专门训练提升中间层可退出性；AdaInfer（Fan等人, 2024）使用推理时统计实现无训练退出决策；SpecEE（Xu等人, 2025）、Kangaroo（Liu等人, 2024）和Mixture-of-Depths（Raposo等人, 2024）在推理时减少有效深度或搜索空间。然而，现有方法仍存在高昂的嵌入解码开销，且未充分考虑KV缓存构建与权重卸载等系统约束。我们的方法通过轻量级、系统感知的早期退出设计解决了这些局限。

## 3预备知识与分析

大多数基于Transformer的LLM采用仅解码器架构（Vaswani等人, 2017；Brown等人, 2020；Touvron等人, 2023a），包含分词器、嵌入层、解码器层堆栈及最终语言建模头。每个解码器层包含自注意力模块与前馈网络。自回归推理时，词元逐层处理，而KV缓存存储先前词元的键值对以避免冗余计算（Pope等人, 2023）。

由于延迟与计算量均与解码器层数大致呈线性关系，跳过不必要层可直接提升推理效率。这推动了LLM的早期退出研究，即通过LM头解码中间隐状态以决定是否提前终止计算。然而，标准LLM仅训练从最终层解码，导致中间层与输出分布对齐较弱，语义解码能力不足。

图3：（a）平均熵与（b）困惑度（PPL）反映Llama2-7B与Llama3-8B在WikiText数据集上中间层经层级早期退出监督前后的变化。
图3使用WikiText数据集（Merity等人, 2016）的熵与困惑度评估Llama2-7B和Llama3-8B中间层的语义可解码性。熵衡量词汇表上的预测不确定性，PPL反映层级语言建模质量。监督前，浅层与中间层呈现高熵且PPL持续较差，证实其隐状态提供不可靠的词元预测。此局限解释了为何直接应用于现成LLM时常导致生成质量下降。

为实现早期退出推理效率与模型性能的有利权衡，增强中间层的语义表达力与可解码性至关重要。一种有效方法是层级早期退出监督，即在训练时将共享LM头连接至多个关键Transformer层以预训练模型（Elbayad等人, 2019；Elhoushi等人, 2024）。该方法对不同深度的隐状态应用统一语言建模目标，将中间表示与最终层解码能力对齐，且不增加额外推理开销。

形式上，令 $x_{l}$ 表示第 $l$ 层Transformer产生的隐表示， $g(\cdot)$ 为共享LM头。训练步骤 $t$ 的目标为：

$\mathcal{J}(X,Y,t)=\sum_{l=0}^{L-1}\tilde{e}(t,l)\;\mathcal{J}_{\mathrm{CE}}\big(g(x_{l+1}),Y\big)$, (1)

其中 $\mathcal{J}_{\mathrm{CE}}$ 为交叉熵损失， $\tilde{e}(t,l)$ 为归一化层级权重。该调度逐步将监督从深层转向浅层，避免对浅层表示过度正则化，同时保持最终层性能。

如图3所示，层级监督显著降低了Llama2与Llama3中间层的熵与PPL，同时最终层性能几乎不变。这些结果表明最终层语义能力可有效迁移至浅层，为可靠的早期退出推理提供坚实基础。

图4：（a）-（d）由公式(2)推导的不同特征阈值及对应词元覆盖率；（e）最终层Top-1词元出现在各层Top-N词元中的比例；（f）各层最终层Top-1概率与Top-N词元内局部Top-1概率的平均差值。
## 4轻量级设计

基于上述分析与监督策略，我们设计了一种受三项统计观察指导的轻量级早期退出预测器。

### 4.1退出特征选择

随着LLM表示随深度增加解码性增强，有效早期退出需选择跨层且跨模型可靠稳定的特征。我们考虑四种常见候选：Top Prob（最大softmax概率）