Technical Summary: Depth Exploration Decoding (DEX)

1. Problem Statement

Autoregressive (AR) decoding in Large Language Models (LLMs) is inherently expensive because every generated token must pass through the entire layer stack before the next token can be produced. Recent studies have identified significant redundancy along the depth dimension: many tokens become predictable at intermediate layers, meaning they do not strictly require the full model depth to be accurate.

Existing lossless depth-adaptive methods attempt to exploit this by employing a depth-selection strategy. These methods draft a token from a single, pre-selected intermediate depth and verify it against the final-depth model. However, the authors identify a critical selection bottleneck:

• Late Exit: If the chosen exit depth is too deep, computation is wasted on layers that were already redundant.
• Early Exit: If the chosen exit is too shallow and the prediction mismatches the final-depth output, the system triggers a fallback. This not only discards the current token but also invalidates any dependent "descendant" drafts, leading to significant efficiency losses.

The paper posits that the gap between current selection-based methods and an ideal "oracle" (which knows the exact earliest depth a token becomes stable) is substantial. The core question is whether this redundancy can be exploited more effectively by reducing the risk of a wrong exit-depth decision.

2. Methodology: Depth Exploration Decoding (DEX)

The authors propose Depth Exploration Decoding (DEX), a lossless algorithm that replaces single-depth selection with parallel exploration over multiple candidate depths.

Core Concepts

• Earliest Available Depth (EAD): The authors define EAD as the shallowest layer where a token's prediction matches the final-depth output and remains stable through subsequent layers.
• Exploration Set ($X$): Instead of choosing one exit, DEX defines a set of candidate depths $\{d_0, d_1, \dots, d_{N-1}\}$, where $d_{N-1}$ is the final depth.
• Resolution: The overhead of DEX is bounded by the "resolution" of this set (the maximum gap between any layer and the next explored depth). As the set becomes finer, the overhead approaches zero.

The Algorithm: Expand–Commit–Collapse

DEX operates on a depth–position lattice managed by a Finite State Machine (FSM). The process consists of three phases:

1. Lattice Expansion (Parallel Exploration):

   • The model is partitioned into $N$ depth stages, each handled by a separate "depth explorer" (potentially on different hardware units).
   • At each commit position, the system expands candidate branches in parallel. An exploration token can undergo depth expansion (passing its hidden state to the next layer explorer) or position expansion (using the current candidate as a prefix to generate the next token position).
   • This creates a tree of potential futures, where different branches represent predictions made at different depths.

2. Commit (Validation):

   • Once the "anchor path" (the branch traversing all layers) reaches the final depth, it produces the final-depth reference token.
   • DEX validates all non-final candidate tokens from the anchor path against this reference.
   • It selects the shallowest branch whose token matches the final-depth reference. If no intermediate match is found, it falls back to the final-depth branch.
   • Crucially, the token appended to the output is always the one produced by the final-depth reference, ensuring lossless equivalence to standard AR decoding.

3. Collapse (State Pruning):

   • After committing a token, the system collapses the lattice. Branches that are inconsistent with the new committed prefix are discarded.
   • Only the states (KV caches, hidden states) from the selected branch are preserved for the next step.
   • This ensures that while the system explores multiple depths, it maintains the exact state required for the next AR step.

Optimizations

• Depth-Coupled Sampling (DECS): To prevent different depths in the same slot chain from proposing the same token (which wastes exploration), DEX masks tokens already proposed by shallower depths when decoding deeper candidates. This encourages diversity in intermediate proposals without altering the final sampling distribution.
• Inducing Adapter: For standard LLMs where intermediate layers are not naturally aligned with the output head, the authors propose attaching small residual adapters to intermediate depths. These are trained via self-distillation to align intermediate hidden states with the final LM head, improving the "readiness" of tokens at shallower depths.

3. Key Contributions

1. Formalization of Token Readiness: The paper defines Earliest Available Depth (EAD) to quantify when a token becomes predictable. It uses this to mathematically demonstrate the "selection bottleneck" in existing methods, showing that their overhead is driven by the risk of wrong selection rather than just the cost of computation.
2. The DEX Algorithm: A lossless decoding algorithm that instantiates depth exploration via an expand–commit–collapse procedure. It preserves exact equivalence to standard AR decoding while reducing the cost of committing tokens.
3. Empirical Validation: The authors demonstrate that DEX outperforms representative depth-selection baselines (LayerSkip, AdaDecode, DEL) and achieves competitive end-to-end throughput against speculative and distributed decoding methods (Lookahead, PEARL, EAGLE).

4. Experimental Results

The authors evaluated DEX on both early-exit-trained models (LayerSkip-trained CodeLlama-34B and Llama-2-70B) and standard LLMs (CodeLlama-34B, Llama-2-70B, Qwen3-32B) across tasks including math reasoning (GSM8K), code generation (HumanEval), and summarization (XSum).

• Throughput Gains: DEX consistently achieved higher tokens-per-second (TPS) than depth-selection baselines. On LayerSkip models, DEX realized a larger fraction of the theoretical "oracle" headroom than selection methods.
• Scaling with Resolution: Throughput improved as the exploration resolution became finer (i.e., more explorers were added). For example, on Llama-2-70B, increasing explorers from 3 to 8 significantly narrowed the gap to the oracle speedup.
• Competitiveness: In end-to-end comparisons, DEX (with 3–4 explorers) was competitive with multi-GPU tensor parallelism (TP) and distributed speculative decoding (PEARL, EAGLE) on matched hardware resources.
• Ablation: Removing intermediate branches (single-exit variant) showed that the performance gain comes specifically from exploring multiple depths, not just parallelizing the layers.

5. Significance and Claims

The paper claims that parallel depth exploration provides a scalable way to exploit the underused depth axis of LLM decoding.

• Paradigm Shift: Unlike previous methods that gamble on a single exit point, DEX treats depth adaptation as a resolution-controlled problem. By exploring a set of depths, it converts the uncertainty of "when to exit" into a bounded overhead determined by the spacing of the exploration set.
• Lossless Efficiency: DEX achieves these gains without compromising the correctness of the output (lossless) and without requiring changes to the model architecture (except for optional adapters in standard models).
• Complementarity: The authors note that DEX operates on the "depth axis" (how early a token is determined), which is complementary to "token-axis" acceleration methods (how many future tokens are predicted). They suggest that while combining them is possible, it requires complex coordination.

In conclusion, the paper establishes that by replacing single-depth selection with parallel exploration, it is possible to significantly reduce decoding latency while maintaining the exact behavior of standard autoregressive models.