TELL-TALE: Task Efficient LLMs with Task Aware Layer Elimination

TL;DR

Tale introduces Task-Aware Layer Elimination, an inference-time greedy pruning algorithm that removes entire transformer layers to optimize task-specific validation performance without retraining. By exhaustively evaluating single-layer deletions at each step, Tale often achieves higher accuracy with moderate speedups across diverse tasks and models, and its benefits persist when combined with finetuning. An information-theoretic lens shows that pruning can remove bottleneck layers, increasing the flow of task-relevant information through the remaining network. The approach reveals task- and model-specific layer importance patterns and provides a versatile framework for efficient deployment and interpretability of transformer-based LLMs.

Abstract

In this paper we introduce Tale, Task-Aware Layer Elimination, an inference-time algorithm that prunes entire transformer layers in an LLM by directly optimizing task-specific validation performance. We evaluate TALE on 9 tasks and 5 models, including LLaMA 3.1 8B, Qwen 2.5 7B, Qwen 2.5 0.5B, Mistral 7B, and Lucie 7B, under both zero-shot and few-shot settings. Unlike prior approaches, TALE requires no retraining and consistently improves accuracy while reducing computational cost across all benchmarks. Furthermore, applying TALE during finetuning leads to additional performance gains. Finally, TALE provides flexible user control over trade-offs between accuracy and efficiency. Mutual information analysis shows that certain layers act as bottlenecks, degrading task-relevant representations. Tale's selective layer removal remedies this problem, producing smaller, faster, and more accurate models that are also faster to fine-tune while offering new insights into transformer interpretability.

Figures (6)

• Figure 1: Illustration of Tale layer elimination. Candidate layers (yellow) are tested for removal, and the best-performing ones above the threshold are permanently dropped (red) until no further improvement is possible.
• Figure 2: Accuracy progression of Tale across 9 benchmark datasets for LLaMA 3.1 8B. Each curve represents the accuracy at successive iterations. The $\star$ denotes the best-performing layer drop configuration, while the $\square$ highlights the Best Speed up with at least Baseline Accuracy (BSBA) configuration.
• Figure 3: Evolution of mutual information (MI) across transformer layers for different benchmark datasets and different models. Each subplot shows how information is processed and transformed as it flows through the network layers, demonstrating distinct patterns of information propagation for (a) ARC-Easy on Qwen 0.5B, (b) BoolQ on Lucie 7B, and (c) BigBench on LLama 8B.
• Figure 4: Nine benchmark tasks indicating performance after one layer is dropped from different positions in Llama3-8B.
• Figure 5: Layer-wise output performance for LLaMA models: results when generating predictions from intermediate layers 1 through 32 on three different datasets.