Depth Exploration for LLM Decoding¶

会议: NeurIPS 2026
arXiv: 2606.29223
论文: arXiv PDF
领域: LLM 其他
关键词: LLM 解码加速, 深度探索, 无损推断, 早退机制, 并行解码

一句话总结¶

提出深度探索解码（DEX），摒弃单早退层的猜测与回滚惩罚，在多个候选深度上并行展开候选分支，由最终深度验证并提交真实 token 并塌缩保留有效前缀分支，在保持严格无损自回归等价性的同时显著提升解码吞吐。

研究背景与动机¶

现代自回归（AR）大语言模型在生成每个新 token 时，都必须完整穿透全部网络层才能产出下一个 token，这种严格的串行依赖构成了推断延迟的核心瓶颈。已有研究发现 Transformer 在深度维度存在显著计算冗余，大量常规或易预测的 token 在穿透中间层时其语义和预测分布就已经趋于稳定，无需经过最后几层微调。为利用该冗余，早期的自适应深度方法尝试让模型从中间层直接早退；而近期无损加速方法（如 LayerSkip、AdaDecode、DEL）引入了投机验证范式，即在选定的中间层预测草稿 token，再由完整深度的最终层进行一致性校验，从而确保输出分布与标准 AR 严格一致。

然而，现有无损深度自适应方法普遍受制于「单点深度选择瓶颈（Selection Bottleneck）」。这些方法无论采用固定深度还是启发式分类器，每次都只能赌一个非最终退出深度。一旦选得过深，中间本可早退的浅层冗余计算直接被白白浪费；而一旦选得过浅，草稿 token 极易通不过最终层校验，从而触发最终层全量回滚，并被迫丢弃所有依赖于该错误草稿的后续子代 token。在存在分支级联的情况下，单次早退失误导致的层计算惩罚甚至会显著超过标准单步 AR 解码的原始开销。

本文的切入视角在于：既然单一退出深度的决策在不确定性下难以兼顾激进与保守，何不将「孤注一掷的单点选择」转变为「候选深度集合上的并行探索」？核心 idea：提出深度探索解码（DEX），在预设的多个候选深度上并行展开候选分支，并在每个提交位置以最终深度输出为基准，严格提交最终深度 token、命中可复用分支并塌缩丢弃失效分支，将选层失误的不确定性风险转化为受候选深度粒度严格约束的确定性截断误差。

方法详解¶

整体框架¶

DEX 将具有 $L$ 层的解码器划分为 $K$ 个连续的深度阶段（Stage）$\Pi_k = \text{Layers}(d_{k-1}, d_k]$，其中 $0 = d_{-1} < d_0 < \dots < d_{N-1} = L$，每个阶段由一个硬件计算单元承载并运行独立的深度探索器（Depth Explorer, $E_k$）。DEX 围绕当前待提交位置 $p_c$ 组织探索循环，并在每个周期内执行展开（Expand）、验证与提交（Validate and Commit）以及状态塌缩（Collapse）三步流程。

展开过程中，候选分支在「深度-位置」构成的二维网格（Lattice）上推进：深度方向将隐藏状态跨阶段后传，位置方向则以当前候选 token 为前缀开启下一个 token 的候选探索；由最初前缀持续推至最终深度的分支构成「锚点路径（Anchor Path）」，并在到达最终深度时提供绝对基准输出。提交时，DEX 校验浅层候选是否命中最终深度基准，挑选中标的最浅分支予以保留，同时由有限状态机（FSM）精准裁剪并重置失效分支的 KV Cache 与隐状态。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["已提交上下文序列<br/>Current Committed Prefix"] --> B["阶段并行晶格展开<br/>Parallel Lattice Expansion"]
    B --> C["深度耦合采样 (DECS)<br/>非最终深度去重掩码多样化"]
    C --> D["锚点路径最终深度验证与提交<br/>Final-Depth Reference Commit"]
    D -->|命中浅层候选| E["保留最早一致分支<br/>Preserve Earliest Matching Branch"]
    D -->|未命中浅层| F["回退最终深度分支<br/>Fallback to Final Depth Branch"]
    E --> G["晶格状态塌缩与前缀转移<br/>FSM Lattice Collapse & KV Reset"]
    F --> G
    G --> H["输出严格等价 Token 并推进<br/>Committed Token & Advance"]

关键设计¶

1. 晶格展开与最早可用深度（EAD）有界开销：消除选层失误的回滚级联

传统深度选择方法若在位置 $t$ 选定退出层 $s_t$，其接受概率为 $p_t^{\text{acc}}$；一旦选浅验证失败，不仅要付出完整深度 $L$ 的重算代价，还要报废后续因该投机预测而展开的 $v_t-1$ 个子代候选，总期望开销为： $$W_t^{\text{sel}} = p_t^{\text{acc}} s_t + (1 - p_t^{\text{acc}})\left(L + \sum_{i=t}^{t+v_t-1} s_i\right)$$ 若令稳定最早可用深度 $\text{EAD}^{\text{stab}}(t)$ 为从某层开始至第 $L$ 层预测完全与最终输出一致的最小层号，单点选层相比 Oracle 的开销为 $c_t^{\text{sel}} = W_t^{\text{sel}} - \text{EAD}(t)$，极易产生正向放大。

DEX 引入探索集合 $\mathcal{X} = \{d_0, d_1, \dots, d_{N-1} = L\}$，定义探索分辨率 $\Delta(\mathcal{X}) = \max_{\ell \in [L]} (\lceil \ell \rceil_{\mathcal{X}} - \ell)$。在并行展开模式下，某 token 实际消耗的有效层数等于将其稳定深度向上取整到探索集内的最近深度： $$W_{\mathcal{X}}(t) = \lceil \text{EAD}(t) \rceil_{\mathcal{X}}$$ 由此，DEX 相对理论 EAD Oracle 的开销严格受限于分辨率常数： $$0 \le c_t^{\text{exp}} = \lceil \text{EAD}(t) \rceil_{\mathcal{X}} - \text{EAD}(t) \le \Delta(\mathcal{X})$$ 这意味着只要探索集划分足够精细，DEX 的深度自适应开销就呈现单调递减并逼近 Oracle，彻底杜绝了因单点误判而导致级联报废的极端惩罚。

2. 最终深度基准驱动的提交与晶格塌缩机制：确保数学级无损等价性

在多层分支并行推演时，探索晶格内会同时存在多条潜在的未来生成分支，但真正的自回归解码在逻辑上永远只有单一确定的前缀。DEX 解决该矛盾的核心在于「计算探索解耦，提交绝对锚定」： - 在状态机推进至终态时，锚点路径必定在最高深度 $L$ 产出未被修改的基准输出 $u_{N-1}$； - 无论浅层候选分支 $\{u_0, \dots, u_{N-2}\}$ 是否有命中，系统最终向官方序列追加的 token 恒等于 $y = u_{N-1}$； - 浅层匹配的作用仅在于确定保留状态——若存在 $k < N-1$ 满足 $u_k = u_{N-1}$，则标记 $b_k$ 分支为可用状态并继承其在后续 token 位置预计算的 KV Cache 与隐状态。

随后触发塌缩操作（Collapse）：依据有限状态机规则，仅保留前缀完全匹配官方已提交序列 $Y \| y$ 的节点状态，将其平移更新为新一轮探索周期的根节点；其余不匹配分支对应的探索缓存（Explore Buffer）与孤立 KV 项被彻底清空或覆盖，从而以数学归纳法证明了输出分布和缓存状态与原生串行 AR 完全一致。

3. 深度耦合采样（DECS）：消除同槽重复预测并保留目标采样分布

当同一个采样槽（Sampling Slot，即具有相同前缀假定和目标位置的分支）在多个候选深度分别解码时，若各深度独立贪婪或采样，深层往往会重复浅层已经提议过的 token。这种重复提议毫无收益：若该 token 命中最终基准，浅层分支早已能够承接复用；若未命中，重复提议反而挤占了探索其他潜在正确 token 的宝贵槽位。

DEX 提出了深度耦合采样机制（Algorithm 2）。在同一个槽链（Slot Chain）中，节点 $i$ 在解码非最终候选时，先收集该槽内更浅层节点已提议的集合 $\mathcal{A} = \{\mathcal{B}[a] \mid a \in \text{Prior}(i)\}$，并在当前 logits 上将其屏蔽置为负无穷： $$\ell_v \leftarrow -\infty, \quad \forall v \in \mathcal{A}$$ 但在状态机到达最终步 $q = q_{\max}$ 且针对基准节点 $i_{\text{ref}}$ 时，DECS 强制置空 $\mathcal{A} \leftarrow \emptyset$，恢复完整的未掩码 logits；在随机采样模式下，同槽各节点共享同一份预先生成的 Gumbel 噪声 $\mathcal{G}_{r(i), :}$，根据 Gumbel-Max 技巧，最终深度基准 token 的产生严格遵循目标无偏采样分布，既极大丰富了中间深度的探索广度，又丝毫不损毁最终采样等价性。

损失函数 / 训练策略¶

对于已经经过早退训练的模型（如 LayerSkip 系列），DEX 可直接利用其自然具备的浅层对齐特征，无需额外训练。

对于标准开源 LLM（如 CodeLlama、Llama-2、Qwen3 等），其中间层的隐藏状态未对齐到全局共享的 LM Head，直接输出会导致浅层 EAD 大量偏向末层。DEX 引入参数量极小（仅占主干约 1.15%~2.68%）的残差诱导适配器（Inducing Adapter）。适配器仅挂载在各个非最终候选深度 $\ell \in \mathcal{X} \setminus \{L\}$ 的出口，结构定义为： $$\Delta h_{\ell} = \alpha_{\ell} W_{2, \ell} \sigma\left(W_{1, \ell} \text{RMSNorm}(h_{\ell})\right), \quad \tilde{h}_{\ell} = h_{\ell} + \Delta h_{\ell}$$ 其中隐藏层扩展维度 $m = 2h$，$\alpha_\ell$ 为可学习缩放门控（初始设为 $10^{-3}$）。

在训练阶段，骨干模型完全冻结，使用 ShareGPT（约 70k 样本）通过自蒸馏方式对齐中间表示。损失函数采用基于最终层软标签的掩码全词表 KL 散度： $$\mathcal{L} = \frac{1}{|\mathcal{X}| - 1} \sum_{\ell \in \mathcal{X} \setminus \{L\}} \tau^2 \text{KL}\left(\text{softmax}(z_L / \tau) \;\|\; \text{softmax}(z_\ell / \tau)\right)$$ 其中温度 $\tau = 1$。使用 AdamW 优化器，学习率 $10^{-3}$，训练 2 个 epoch，在 4 块 H100 上训练 70B 模型仅需 12.6 小时。由于最终深度不接入适配器且直接负责判定输出，适配器的任何轻微扰动均不影响输出 token 的严格等价性。

实验关键数据¶

主实验¶

评估在 NVIDIA H100 GPU 上进行，覆盖 LayerSkip 专用早退模型及标准 Instruct 大模型，基准涵盖数学推理（GSM8K）、代码生成（HumanEval）与长文摘要（XSum）。下表展示了 8 块 GPU 环境下 70B 级别大模型的系统级吞吐表现（单位：tokens/s，贪婪解码，最大长度 512）：

数据集 / 模型	指标 (tokens/s)	DEX-1/8 (本文 8-GPU)	Tensor Parallelism (TP-8)	AdaDecode / LSSD (典型基线)	相对 TP-8 提升
LS-Llama-2-70B - GSM8K	解码吞吐	68.4	38.2	44.5	+79.1%
LS-Llama-2-70B - HumanEval	解码吞吐	74.1	39.8	48.2	+86.2%
LS-Llama-2-70B - XSum	解码吞吐	62.7	36.5	42.1	+71.8%
Llama-2-70B-Instruct - GSM8K	解码吞吐	61.8	37.9	36.2 (AdaDecode)	+63.1%
Llama-2-70B-Instruct - HumanEval	解码吞吐	66.5	39.4	38.0 (AdaDecode)	+68.8%
Qwen3-32B (matched resource)	解码吞吐	54.2	34.1	42.6 (EAGLE-3)	+58.9% (vs TP)

注：DEX 在同等硬件卡数资源下（如 3-GPU 对齐 34B、4-GPU 对齐 70B），其端到端吞吐可与 PEARL 和 EAGLE 系列正面匹敌；随着探索分辨率由 DEX-1/3、DEX-1/4 扩展至 DEX-1/8，其吞吐呈现单调攀升并全面超越各类分布式投机解码基线。

消融实验¶

为厘清 DEX 的性能来源，作者在 8 探索器配置下对探索结构、DECS 采样优化以及运行期耗时分布进行了细致剖析：

配置 / 组件剖析	平均耗时 / 阶段 (ms)	耗时占比 (%)	核心消融结论与吞吐表现
全量 DEX-1/8 (展开阶段)	13.830 ms	89.13%	计算主体集中在各阶段前向计算，GPU 计算效率极高
跨 GPU 通信 (NCCL Communication)	1.222 ms	7.88%	仅需阶段间隐状态与 token id 传递，通信开销轻量
晶格塌缩 (Collapse Bookkeeping)	0.309 ms	1.99%	状态裁剪与指针重置极快，不构成运行期系统瓶颈
解码与提交 (Decode/Commit)	0.156 ms	1.01%	最终深度对比开销微乎其微
消融：单退出分支 (Single-Exit Variant)	-	-	禁用中间探索分支退化为双退出后，吞吐下降 28.4%~34.6%
消融：无 DECS (Vanilla 独立采样)	-	-	出现大量同槽重复预测，贪婪与采样下吞吐分别下跌 8.2% 与 11.5%

关键发现¶

分辨率与吞吐单调正相关：实验严格证实了理论推导，随着探索深度划分从 $1/2 \to 1/4 \to 1/8$，模型实际吞吐平滑递增，不断缩小与理论稳定 EAD Oracle 之间的差距，证实粗粒度决策是阻碍早退加速的核心瓶颈。
DECS 在贪婪与采样模式下均稳定增益：DECS 强制深层避开浅层已探索词元，不仅未破坏基准分布，反而显著提高了多深度命中最终词元的概率。
极低的塌缩与同步开销：在全系统运行时剖析中，晶格塌缩和提交判定总耗时仅占约 3%，证明 DEX 的状态机设计极为精炼，绝大部分算力被切实投入到 Transformer 层前向传播中。

亮点与洞察¶

将离散选层的概率风险转化为有界逼近问题：传统早退算法沉迷于训练更高精度的退出分类器，而 DEX 转换了问题结构，通过并行候选覆盖使系统开销天然受限于相邻探索深度的间距常数，从数学上杜绝了回滚惩罚放大。
巧借 carry-lookahead 算子原理重构解码状态机：借鉴超标量体系结构中先行进位加法器的思想，在主干深度未到达之前预先铺设计算晶格，把自回归串行链重构为可部分并行复用的局部前缀流。
解耦探索多样性与目标分布保真度：通过 Gumbel 噪声共享与末层无条件解掩码，既让中间层各展所长覆盖候选，又保证官方提交产物在数学上完全等价于标准自回归解码。

局限与展望¶

计算资源依赖度高：DEX 的实时墙钟加速高度依赖多 GPU 硬件的物理并发。在单卡单卡串行模拟环境下，展开晶格会造成总计算量倍增，无法直接转化为端到端加速。
纯深度探索存在渐进吞吐上限：纯深度维度的理想加速比受限于最浅候选深度，即上限约为 $L / d_0 \approx K$；若要进一步突破性能瓶颈，需打破单一维度限制。
正交维度的融合潜力：未来最值得探索的方向是将深度维度的 DEX 晶格与 token 维度的推测树验证（如 Medusa/EAGLE）进行统一联合编排，构建时空二维协同探索加速框架。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [打破深度早退必做单点选层的惯性思维，首创深度探索与晶格塌缩解码机制]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖专用早退模型与主流通用大模型，多 GPU 真实吞吐对比详实]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严谨自洽，状态机转换逻辑清晰，附录开销剖析透彻]
价值: ⭐⭐⭐⭐⭐ [为多卡推断集群挖掘模型深度冗余提供了极具实用性与理论美感的全新范式]