← 返回 2026-04-21

River-LLM：基于 KV 共享的大语言模型无感早退机制 River-LLM: Large Language Model Seamless Exit Based on KV Share

Yingtao Shen, An Zou 📅 2026-04-20 👍 6 2026-07-13 08:36
Early Exit KV Cache LLM 推理加速 免训练 动态深度
arXiv HuggingFace PDF

通过轻量级 KV 共享退出流，让早退 token 自然补全缺失的 KV 缓存，实现 1.53×–2.16× 免训练加速。

前置知识

KV Cache

自回归推理时，每个 token 的 Key 和 Value 矩阵被缓存下来，避免后续 token 重复计算前面的注意力。这使得解码第 $t$ 步时只需计算一个新 token 的 QKV，而 K/V 直接从历史缓存中读取，是现代 LLM 高效推理的核心机制。

River-LLM 整篇论文都围绕 KV Cache 完整性展开：早退时跳过的层没有产生 KV，导致后续 token 无法在那些层做注意力。读懂本文必须先理解 KV Cache 是怎么被构造和复用的。

Early Exit（早退）

动态推理中的一种范式，允许样本在不同深度提前终止，从而跳过冗余计算。每个 token（或每条序列）在中间层附加一个分类头或退出模块，置信度达标就跳出主网络。在 CNN/BranchyNet 中已经成熟，但 LLM 上的落地受 KV 缓存制约。

本文是 Early Exit 在 decoder-only LLM 上的最新进展，需要理解早退为什么带来加速（跳过层数 × 节省 FLOPs），以及为什么在自回归场景下会引发新的工程问题。

PTQ（训练后量化）

Post-Training Quantization，对预训练模型权重进行低比特（如 4-bit）转换而无需重新训练。本文采用 W4A16 方案：权重 4-bit、激活与 KV 保持 FP16，以减少精度损失并加速推理。

River-LLM 的 Exit River 是骨干层的 4-bit 量化副本，理解 W4A16 才能明白它为什么能在保留 KV 完整性的同时获得 2.4× 的吞吐。

自回归推理与 causal mask

decoder-only 模型按 token 顺序逐位生成，当前 token 只能 attend 到它自己和之前的 token，这一约束通过 causal mask 实现。早退打破了 KV 的时序连续性，等价于把这些位置变成'不存在'，需要相应的 mask 或补全策略。

本文反复出现的 KV Mask、State Propagation、KV Recompute 等方案，本质都是在不同假设下处理因果 mask 与早退的兼容性。

状态转移相似度（state transition similarity）

定义为同一 decoder 块输入隐藏态 $h_t^{(l-1)}$ 与输出隐藏态 $h_t^{(l)}$ 的余弦相似度 $s_t^{(l)} = \frac{h_t^{(l-1)\top} h_t^{(l)}}{\|h_t^{(l-1)}\| \|h_t^{(l)}\|}$，衡量一层对表征的'改变程度'。

这是 River-LLM 退出判定的核心信号：相似度高意味着该层'无事可做'，是潜在的退出候选；本文发现它与累积量化误差呈中等正相关（$r=0.5536$），可作为可靠代理。

研究动机

Token 级别 Early Exit 理论上能让 LLM 推理获得最高 3.3× 的加速（实验在 Llama3.2 1B + GSM8K 上测得），但实际墙钟时间几乎拿不到这种收益。瓶颈来自 decoder-only 架构独有的 KV Cache Absence 问题：当某个 token 在第 $l$ 层早退，被跳过的第 $l+1 \sim L$ 层就不会产出对应的 K/V。后续 token 在自回归时必须在这些层做注意力，但历史 KV 是空的。现有四条补救路径都存在结构性缺陷：Batching Recompute 需要在遇到缺失 KV 时重算历史隐藏态，在长序列下显存访问被打乱、算力节省被吞噬；State Propagation 把当前层的隐藏态当作后续所有被跳过层的输入，直接计算对应 K/V，但近似误差大、准确率下滑明显；KV Mask 直接把这些 token 在注意力中 mask 掉，让 D-LLM 等方法不得不限制最小退出层、严重牺牲退出粒度；Mono-Decreasing Exit 要求整条序列的退出位置单调递减以兼容 KV 时序，虽然规避了缺失问题却把粒度从 token 级退化回 sequence 级。图 3 显示在 GSM8K 上这些方案要么 backbone 延迟更高（Mask ≈ 19ms/token），要么虽然层数更少但延迟几乎不降（Recompute 16.91ms 与 Baseline 15.18ms 接近），验证了'理论加速无法落地'的困境。

本文的目标是本文追求一个称为'无缝退出'（Seamless Exit）的目标，包含两个独立而互补的属性。第一是 Granular Freedom：每个 token 都能独立选择退出层，没有任何单调性或最小退出层的硬约束。第二是 Intrinsic KV Integrity：被跳过层的 KV 缓存必须在退出过程中被自然地生产并写入对应地址，不再需要任何 post-hoc 的恢复、重算或 mask。最终目标是让 Token 级 Early Exit 的理论加速倍数真正反映到 wall-clock 上。

与已有工作不同的是，过去一年的相关工作（SkipDecode、DiffSkip、D-LLM、EE-LLM、Balcony、LayerSkip 等）都聚焦于'怎么决定要不要退'，但都默认接受 KV 缺失作为前提条件，再去修补。River-LLM 反其道而行：从 KV 完整性这一被忽视的视角切入，通过给骨干模型加挂一条与它共享 KV 地址空间的'退出流'（Exit River），让早退动作天然地填充 KV，把'退出'和'恢复'统一成一个操作。这条退出流本身是骨干层经过 4-bit PTQ 后的轻量级副本，吞吐量是骨干块的 2.4×，并且其产出 KV 与骨干原生 KV 的余弦相似度稳定在 0.97 以上。这种'以共享代替恢复'的设计视角，在之前的 Early Exit 文献中尚未被系统提出。

核心方法

River-LLM 的整体思路可以一句话概括：让早退不再是一次'节省层数'的决策，而是一次'换一个更快的子模型继续算'的操作。它给 decoder-only 骨干（L 个块）从起始层 $L_s$ 起逐层挂接一个'退出流'，每个退出层是同构骨干层经过 W4A16 PTQ 后的 4-bit 副本，且 KV 缓存的地址索引与对应的骨干层完全一致。在自回归生成第 $t$ 个 token 时，隐藏态从底层往上跑，每跑完一层就计算该层的 state transition similarity $s_t^{(l)} = \cos(h_t^{(l-1)}, h_t^{(l)})$，若超过阈值 $\tau$ 就立即把后续路径切换到 Exit River 上、由它接力跑到 LM Head。这一跳不仅完成了 token 的剩余推理，还顺带把本该由被跳过骨干层产生的 K/V 写进了对应槽位，后续 token 因此能像正常解码一样读到完整历史。Prefill 阶段（处理 prompt）保持 Sequence-level Exit 以利用并行注意力核；Generation 阶段切到 Token-level Exit 以榨干粒度。配套的 Backbone Offloading 还能把那些几乎用不上的深层骨干块从显存驱逐出去，进一步省 VRAM。

最核心的创新在于把'KV 恢复'这件事从'运行后补救'前置成了'退出时的自然产物'。与已有方法相比：KV Recompute 把缺失当异常、用历史隐藏态重新跑一遍，代价巨大；State Propagation 直接拿当前隐藏态替代后续层输入，近似粗暴；KV Mask 则完全放弃这些层的注意力贡献；Mono-Decreasing 干脆用拓扑约束避开问题。River-LLM 的不同在于它承认 KV 必须真实存在，但不需要由原骨干层来生成——只要有一组权重近似于骨干、且在结构上'占住'对应 KV 槽位就够了。它利用了 decoder 块输入-输出天然的高相似性（图 6b 显示相似度随层数递增），使得 4-bit 量化版本就足以充当'语义补完器'。第二个关键洞察是第一层 state transition similarity 与最终 KV 误差之间存在 $r=0.5536$ 的中等正相关，让退出判定不需要单独训练 head，直接用一个余弦相似度阈值即可，且该阈值能可靠预测累积量化漂移。

方法步骤详情

完整流程可分为四步。第一步是 Exit River 初始化：对每个骨干层 $l \in [L_s, L]$ 复制其 Attention 与 FFN 权重形成对应的 Exit Layer $l$，KV cache 的 addressing scheme（位置索引、head 划分、维度布局）与骨干完全一致；然后对这些副本施加 HQQ 框架下的 4-bit 权重量化、保留 FP16 激活与 FP16 KV 缓存，并通过部分图编译获得 2.4× 吞吐量。整个权重搬运过程无需训练，通常 1 分钟内完成。第二步是 Prefill 阶段处理 prompt：所有 token 采用 Sequence-level Exit，统一在某一固定层切到 Exit River，以最大化并行注意力 kernel 的吞吐。第三步是 Generation 阶段的 Token-level 循环：对当前 token 的隐藏态 $h_t^{(l)}$，在进入下一骨干块前计算 $s_t^{(l)} = \frac{h_t^{(l-1)\top} h_t^{(l)}}{\|h_t^{(l-1)}\| \|h_t^{(l)}\|}$，若 $\min_{b \in B} s_t^{(l)} > \tau$（对当前 batch 内所有正在生成的序列取最小值）则判定可退出，从 Exit River 接力跑到 LM Head，期间由 Exit Layer 写入对应 KV 槽位；否则继续走骨干的下一层。第四步是 Backbone Offloading：统计一段时间内 token 的退出分布，把退出深度之后几乎从未被访问的骨干层从主 VRAM 卸载（论文示例把 14.96 GB 参数压缩到 6.73 GB），仅保留 Exit River 常驻。整条流水线对 batch 内不同 token 独立决策，最大化 Granular Freedom。

技术新颖性

技术新颖性体现在三方面。其一是'共享 KV 寻址'的物理一致性设计——以往方法把 Exit Layer 看作独立的加速器，KV 写入位置任选；River-LLM 把 KV 槽位与骨干层强绑定，使得自回归读写完全对齐，无需任何额外索引或映射层。其二是'PTQ 副本作为语义补完器'的可行性论证——通过图 5 量化展示了 Exit Layer 的 K/V 与骨干原生 K/V 余弦相似度全层稳定在 0.97 以上，证明 4-bit 量化足以承担'KV 占位 + 语义接续'双重任务。其三是退出决策的可解释性——基于 $r=0.5536$ 的相关性证明 state transition similarity 可以作为'累积误差代理'，这是首次把'退出阈值'和'量化漂移'建立定量联系，免去了训练 confidence head 的开销。这三件事组合在一起，使 River-LLM 成为目前唯一在 token 粒度上同时满足 Granular Freedom 和 Intrinsic KV Integrity 的方案。

Seamless exit architecture and inference paradigm: River-LLM. (a) KV-shared exit layer. (b) Inference with seamless exit.
Fig. 4: Seamless exit architecture and inference paradigm: River-LLM. (a) KV-shared exit layer. (b) Inference with seamless exit.
KV Cache similarity between exit layer and backbone decoder of Llama3.2 1B. (a) Key. (b) Value.
Fig. 5: KV Cache similarity between exit layer and backbone decoder of Llama3.2 1B. (a) Key. (b) Value.
(a) Relationship between first layer state transition similarity and last layer backbone-exit value vector similarity (r = 0.5536, p < 0.001). (b) Trend of state transition similarity with the backbone layer index.
Fig. 6: (a) Relationship between first layer state transition similarity and last layer backbone-exit value vector similarity (r = 0.5536, p < 0.001). (b) Trend of state transition similarity with the backbone layer index.

实验结果

在 4 个模型 × 8 个基准的实验矩阵下，River-LLM 几乎在所有任务上都做到了 near-lossless 加速。最具代表性的是 Llama3.2 1B / GSM8K：阈值 $\tau=0.5$ 时平均退出深度仅 3.79 层（总 16 层），准确率 29.3%，相对骨干 33.2% 损失约 4 个点，但拿到 2.16× 加速（84.5→182.9 tokens/s）；把 $\tau$ 调到 0.7，准确率回升到 33.5%（与骨干齐平），加速回落到 ~15 层退出。Llama3.1 8B / GSM8K 同样呈现 $\tau=0.5$ 时 74.4% / 1.78×、$\tau=0.7$ 时 75.6% 的曲线。图 7 给出理论支撑：$\tau=0.5$ 时 GSM8K 上的 next-token 匹配率仍高达 98.03%、KL 散度仅 0.004513。HumanEval 上有一个反直觉的发现——River-LLM 在 Llama3.1 8B 上以 $\tau=0.7$ 完全复现骨干的 57.3%，在 Ministral3 8B 上 $\tau=0.5$ 的 26.2% 也接近骨干的 29.3%，作者将其归因于退出机制在代码任务上'绕过冗余深层的累积噪声与过思考'。Phi4-mini 与 Ministral3 8B 因为架构更新，浅层语义已足够，使用更严的 $\tau=0.8/0.9$：Phi4-mini GSM8K 仍达 81.0%（骨干 82.1%），HumanEval 63.1%（骨干 63.4%）。内存层面（图 9、表 6）River-LLM 16K 上下文仅需 8.73 GB，相比 Balcony 的 19.77 GB、EE-LLM 的 19.55 GB 节省超过一半；64K 上下文也只占 14.73 GB，逼近 Full Quantize 的 12.47 GB。决策逻辑本身只占 8B 模型每 token 推理时间的 0.0688%（约 100 微秒），可以忽略。

Decoder-only transformer early exit techniques.
Table 1: Decoder-only transformer early exit techniques.
Evaluation results of River-LLM across diverse benchmarks.
Table 2: Evaluation results of River-LLM across diverse benchmarks.
Practical generation speedup and accuracy of River-LLM compared with the backbone and full quantization baselines.
Table 3: Practical generation speedup and accuracy of River-LLM compared with the backbone and full quantization baselines.
Different Exit River Quantization on Llama3.1 8B.
Table 4: Different Exit River Quantization on Llama3.1 8B.
Additional evaluation results of River-LLM on Phi4-mini and Ministral3 8B across diverse benchmarks.
Table 5: Additional evaluation results of River-LLM on Phi4-mini and Ministral3 8B across diverse benchmarks.
Peak GPU Memory Usage of Different Methods for Llama3.1 8B, batch_size = 1.
Table 6: Peak GPU Memory Usage of Different Methods for Llama3.1 8B, batch_size = 1.
(a) Average KL Divergence and (b) next token match rate between backbone output and exit prediction, τ = 0.5.
Fig. 7: (a) Average KL Divergence and (b) next token match rate between backbone output and exit prediction, τ = 0.5.
Trade-off between generation throughput and GSM8K accuracy on Llama3.2 1B.
Fig. 8: Trade-off between generation throughput and GSM8K accuracy on Llama3.2 1B.
Peak GPU memory usage of Llama3.1 8B with different methods, batch_size = 1.
Fig. 9: Peak GPU memory usage of Llama3.1 8B with different methods, batch_size = 1.
查看结构化数据
任务指标本文基线提升
GSM8K (Llama3.2 1B) Accuracy / Speedup 29.3% (τ=0.5) / 2.16×；33.5% (τ=0.7) 33.2% / 1.00× (84.5 tok/s) 吞吐量从 84.5 提升到 182.9 tok/s，准确率最多牺牲 3.9 个百分点
GSM8K (Llama3.1 8B) Accuracy / Speedup 74.4% (τ=0.5) / 1.78×；75.6% (τ=0.7) 78.2% / 1.00× (25.3 tok/s) 准确率最多下降 3.8 点但拿到 1.78× 加速，τ=0.7 时几乎无损
GSM8K (Phi4-mini) Accuracy / Speedup 81.0% (τ=0.9) / 1.61× 82.1% / 1.00× (71.4 tok/s) 仅下降 1.1 个百分点，吞吐量 71.4→115.0 tok/s
HumanEval (Llama3.1 8B) Accuracy / Speedup 55.5% (τ=0.5) / 1.77×；57.3% (τ=0.7) 57.3% / 1.00× (25.2 tok/s) τ=0.7 时与骨干完全持平、加速 1.77×，体现'绕过冗余'的反直觉收益
MATH (Llama3.2 1B) Accuracy / Speedup 14.6% (τ=0.5) / 1.88× 17.8% / 1.00× (100.5 tok/s) 数学题准确率略降 3.2 点，吞吐量翻近 2×
HellaSwag (Llama3.2 1B) Accuracy 44.3% (τ=0.5) 45.1% 常识推理几乎无损，平均退出深度仅 3.71 层
Peak GPU Memory (16K, Llama3.1 8B) Total Memory (GB) 8.73 GB（含 6.73 GB 参数 + 2 GB KV） Backbone 16.96 / Balcony 19.77 / EE-LLM 19.55 / Full Quant 6.47 比 Backbone 省 48%、比 Balcony/EE-LLM 省 56%，逼近 Full Quant 但保留关键 token 的 FP16 精度
Exit Decision Overhead % of per-token latency ≈0.0688%（~100 μs on 8B） 理想上限 100% 开销几乎为零，时间复杂度仅 $O(d)$

局限与改进

作者明确承认两点。一是规模验证不足：所有实验在 ≤8B 模型上完成，更大的 24B / 70B 模型是否仍能保持 KV 余弦相似度 >0.97、Exit River 量化副本是否能跟得上更大隐藏维度，均未经验证。二是阶段不均衡：River-LLM 主要为 Generation 阶段优化，对 Prefill 占主导的任务（如 MMLU 这类判别式 benchmark）只能切换到 Sequence-level Exit，加速效果有限（MMLU 在 Llama3.1 8B 上 τ=0.5 仅 66.1%、退出深度 3.04，依然得到加速但天花板不高）。此外从附录 A.4 的对比看，River-LLM 在'极致位宽压缩'上不如混合精度方案——比如 DP-LLM 在 3.25–4.75 bit 区间能做到极致吞吐，但代价是 GSM8K 在 Llama3 8B 上准确率掉到 36.7–46.9%，而 River-LLM 维持在 74.4–78.2%，二者处于不同的 Pareto 区域。还有一个隐含限制：River-LLM 的 Exit River 是 4-bit 量化副本，对量化本身敏感的任务（如安全/对齐敏感场景）可能放大 drift，需要根据表 4 的对比选择 AWQ 等更稳健的量化后端。

独立分析的弱点

独立分析发现三处可改进之处。其一是退出判定只用了第一层 state transition similarity 的最小值（$\min_{b \in B}$），在 batch 内不同样本难度差异大时会偏向保守——一个高难度样本会拖住整批，导致其他简单样本也无法提前退出。可以借鉴 AdaInfer 的 SVM 思路或 DEER 的 confidence 概念，引入 per-sample 独立的退出判定、或者基于 batch 内中位数/分位数的鲁棒聚合，以在长尾 batch 中拿到更多加速。其二是 Exit River 量化方案固定为 W4A16，对不同层的敏感度未做差异化处理——表 4 显示换成 AWQ 后 GSM8K 立刻从 74.4 升到 77.3，证明量化后端选择本身就是显著的精度杠杆。可以考虑对深层 Exit River（更容易累积误差）使用更高 bit-width，做'深度感知的混合精度退出流'。其三是 Backbone Offloading 依赖运行时统计，对响应长度方差大的 workload（如对话流）可能频繁触发 swap，反而拖慢首次响应；可改为基于 prompt 分类的预判式卸载，或者干脆让 Exit River 覆盖所有层、完全去掉骨干以最大化内存节省。

未来方向

作者提出的方向包括向 24B / 70B 更大规模验证、本文框架与混合精度量化或 speculative decoding 的协同（附录 A.4 已给出与 DP-LLM 的对比，River+AWQ 比静态 AWQ 更准且吞吐相当）。基于本文成果可延伸的方向有：第一，把 KV 共享思想泛化到 MoE 架构，让早退的 token 能落到对应专家的 KV 槽位，避免 D-LLM 在 MoE 上遇到的路由一致性难题；第二，把 Exit River 与 Speculative Decoding 结合，让 Exit River 同时充当 draft model、复用 KV 共享带来的缓存亲和性，进一步压缩验证阶段的开销；第三，研究退出判定的可学习增强——虽然本文强调免训练，但 state transition similarity 可作为 fine-tuning 时的辅助 loss，引导骨干层在浅层就形成稳定的 hidden state；第四，把 Backbone Offloading 与 PagedAttention / KV 卸载结合，做一个面向 64K+ 长上下文的完整推理引擎，在吞吐、显存、延迟之间给出可配置的多目标调度。

复现评估

复现评估整体偏友好。算法层面是完全免训练的——权重搬运 + PTQ + 阈值即可落地，作者声称整组权重搬运在 1 分钟内完成；退出判定只需一次 $O(d)$ 的余弦相似度，不依赖任何外部分类器或校准数据。模型与 benchmark 全部公开：Llama3.2 1B、Llama3.1 8B、Phi4-mini、Ministral3 8B 均为开源权重；GSM8K、MATH、HumanEval、BoolQ、HellaSwag、ARC-c/e、MMLU 是标准公开数据集。论文明确报告了硬件（NVIDIA A40，48GB）、batch_size=1、阈值 $\tau$、shot 数（GSM8K 5-shot、MATH 4-shot、其余 0-shot）、评估协议（10 次平均）等关键超参。代码与量化框架依赖 HQQ，文中列出版本与配置。论文未开源代码仓库，但附录给出了足够详细的实现描述（KV addressing scheme、Exit River 串行拓扑、Prefill vs Generation 切换策略、Backbone Offloading 触发条件）。中等难度团队可在 1–2 周内复现核心结果；主要门槛在于 Exit River 的 partial graph compilation（需要熟悉 vLLM/TensorRT-LLM 的算子融合）以及对不同模型重新校准 $\tau$。

论文图表

KV Cache Absence problem for Early Exit in decoder-only LLM.
Fig. 1: KV Cache Absence problem for Early Exit in decoder-only LLM.

用三栏对比展示 KV 缺失问题的根因：(a) 简化示例说明前序 token 在早退后、当前 token 在被跳过层找不到 K/V；(b) 已有 Past Solution 把退出视为一次 Latency–Overhead–Recovery 的权衡；(c) River-LLM 通过 KV 共享让退出路径自然产出 KV。

这是全文问题的可视化锚点，让读者一眼看到'早退→KV 缺失→后续 token 无法 attention'的因果链；动机章节必须引用此图。

(a) Distribution of optimal Token-level Exit position for Llama3.2 1B on GSM8K. Score = 0.26. (b) Token-level Exit significantly outperforms Sequence-level Exit on GSM8K.
Fig. 2: (a) Distribution of optimal Token-level Exit position for Llama3.2 1B on GSM8K. Score = 0.26. (b) Token-level Exit significantly outperforms Sequence-level Exit on GSM8K.

(a) 柱状图显示 Llama3.2 1B 在 GSM8K 上最优退出层的分布，平均层数 4.81，约 33% 的 token 可在第 1 层退出；(b) 对比 Sequence-level Exit 平均退出深度 ~15 层、Token-level Exit 仅 ~5 层，仍保持相近准确率。

用数据量化'Token-level 退出潜力 3.3×'，是讨论粒度重要性的关键证据。

Average ms/token of Token-level Exit using different KV Cache Strategy on GSM8K. (a) Relaxed threshold, Score ≈ 0.15. (b) Strict threshold, Score ≈ 0.25.
Fig. 3: Average ms/token of Token-level Exit using different KV Cache Strategy on GSM8K. (a) Relaxed threshold, Score ≈ 0.15. (b) Strict threshold, Score ≈ 0.25.

柱状图对比 Baseline、Mask、Mono-Decreasing、State Propagation、Recompute 五种策略在 Llama3.2 1B 上的平均每 token 延迟与退出层数。两种阈值下 Recompute、State Propagation 都无法显著低于 Baseline，Mask 甚至更慢。

用延迟数据证明既有 KV 修复策略都不能真正提速，论证'无缝退出'的必要性。