ADEPT: Adaptive Dynamic Early-Exit Process for Transformers

TL;DR

ADEPT tackles the high computational cost of Transformer inference and the KV cache bottleneck that limits earlier exit gains. It introduces a token-level adaptive early-exit framework that operates in both prefill and generation by decoupling skipped-layer KV generation through a Hidden State Mapper and a token-level exit policy formulated as a Markov Decision Process with Direct Preference Optimization. The approach achieves significant efficiency while maintaining or improving performance, with reported gains including up to 25% language-generation savings, 4× speedups in classification, and up to 45% performance improvements on challenging downstream tasks, all with minimal or no fine-tuning of base LLMs. By using a lightweight similarity-based exit criterion and decoupled KV generation, ADEPT provides broad applicability to Transformer-based models and scalable inference optimization.

Abstract

The inference of large language models imposes significant computational workloads, often requiring the processing of billions of parameters. Although early-exit strategies have proven effective in reducing computational demands by halting inference earlier, they apply either to only the first token in the generation phase or at the prompt level in the prefill phase. Thus, the Key-Value (KV) cache for skipped layers remains a bottleneck for subsequent token generation, limiting the benefits of early exit. We introduce ADEPT (Adaptive Dynamic Early-exit Process for Transformers), a novel approach designed to overcome this issue and enable dynamic early exit in both the prefill and generation phases. The proposed adaptive token-level early-exit mechanism adjusts computation dynamically based on token complexity, optimizing efficiency without compromising performance. ADEPT further enhances KV generation procedure by decoupling sequential dependencies in skipped layers, making token-level early exit more practical. Experimental results demonstrate that ADEPT improves efficiency by up to 25% in language generation tasks and achieves a 4x speed-up in downstream classification tasks, with up to a 45% improvement in performance.

Figures (11)

• Figure 1: Illustration of early-exit on the token-by-token basis in the prefill. Anticipated last hidden states ($\Tilde{s}_{Nj}^h$) generated by the hidden state mapper (HSM) are propagated to the following layers for the Keys-Values (KVs) generation and the end of the layer chain as the final hidden states.
• Figure 2: reduced computation by avoid query, attention and feed forward block in the skipped layers (index r > $i^*$), where the exit layer index $i^*$
• Figure S1: Architecture of Hidden State Mapper. It takes neighboring hidden states ($h_{l}$ and $h_{l-1}$) and history states ($h_{intertia}$) as inputs and outputs the anticipated final hidden state ($\hat{h}_{l}$).
• Figure S2: Performance of ADEPT with GPT2 Large model and Llama 13B, 33B and 65B on WikiText103.
• Figure S3: Distribution of early-exit per decoder layer of GPT2 X-Large model with ADEPT on WikiText103 benchmark for an input prompt in the prefill phase with the maximum length. Tokens independently early-exit the layer chain depending on the similarities and the threshold.