Adaptive Layer-skipping in Pre-trained LLMs

TL;DR

FlexiDepth proposes adaptive per-token depth in pre-trained LLMs by inserting a router and adapter at each decoder layer, enabling dynamic layer skipping without modifying original parameters. The gating mechanism computes $G = \sigma(\text{Router}(\text{Norm}(X)))$ with a threshold $\tau$ to route tokens through full the path or an efficient adaptor, while maximizing efficiency via the layer-skipping loss $\mathcal{L}_{\text{skip}}$ and the joint loss $\mathcal{L} = \alpha \mathcal{L}_{\text{skip}} + \mathcal{L}_{\text{lm}}$; a KV-Cache-preserving attention skipping strategy maintains context for autoregressive decoding. Empirically, FlexiDepth on Llama-3-8B-Instruct can skip 8 of 32 layers yet retain 100.7% average benchmark performance and outperforms baselines on multi-token tasks, while large models show greater inherent redundancy enabling more skipping. However, no wall-clock speedup is observed yet due to execution overhead and irregular skipping, highlighting a gap between theoretical FLOPs reductions and practical throughput, and motivating future hardware-aware optimizations and broader MoE integrations. The authors provide open-source code and a DepthMap-style dataset documenting per-token layer usage to foster further study of adaptive compute in LLMs.

Abstract

Various layer-skipping methods have been proposed to accelerate token generation in large language models (LLMs). However, limited attention has been paid to a fundamental question: How do computational demands vary across the generation of different tokens? In this work, we introduce FlexiDepth, a method that dynamically adjusts the number of Transformer layers used in text generation. By incorporating a plug-in router and adapter, FlexiDepth enables adaptive computation in LLMs without modifying their original parameters. Applied to Llama-3-8B, it skips 8 out of 32 layers while maintaining full benchmark performance. Our experiments reveal that computational demands in LLMs significantly vary based on token type. Specifically, generating repetitive tokens or fixed phrases requires fewer layers, whereas producing tokens involving computation or high uncertainty requires more layers. Despite the computational savings, FlexiDepth does not yet achieve wall-clock speedup due to varied skipping patterns and I/O overhead. To inspire future work and advance research on practical speedup, we open-sourced FlexiDepth and a dataset documenting its layer allocation patterns.

Figures (4)

• Figure 1: DepthMap illustrating layer-skipping patterns when applying FlexiDepth to Llama-3-8B-Instruct. The light-to-dark blue gradient indicates layer usage ranging from 16 layers to 32 layers.
• Figure 2: The FlexiDepth layer. Left: Full-processing path where hidden states undergo the pre-trained attention and FFN modules. Right: Skipping path where hidden states bypass the attention module and processed by a lightweight adapter. The router and adaptor (in red) are the only trainable components.
• Figure 3: Comparison of attention masks. Left: the vanilla attention mask. Middle: the attention mask without KV Cache, where hidden states $x_1$ and $x_3$ skip this layer. Right: the attention mask with KV Cache, where despite $x_1$ and $x_3$ skipping this layer, their corresponding keys and values are still computed.
• Figure 4: Percentage of tokens processed by transformer layers 17 to 32. The x-axis represents the layer index, and the y-axis represents the percentage of tokens processed by the layer.