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

AI summary

Overview

Research area: Large language model efficiency and structured pruning (layer-level model compression), with a secondary connection to model interpretability.

Technical level: Intermediate. The method itself is conceptually simple, but the paper assumes familiarity with transformer residual streams, pruning benchmarks, and evaluation frameworks such as LM-Eval.

Scope: The paper introduces TALE (Task-Aware Layer Elimination), a training-free, greedy, inference-time procedure that deletes transformer layers selected by their effect on task-specific validation accuracy, and evaluates it across 9 tasks and 5 model families.

What This Paper Is About

Large Language Models ship with a fixed architecture, even though evidence suggests that not all layers contribute equally to every downstream task. Existing pruning methods mostly reduce computational cost at the price of downstream accuracy, and they typically rely on task-agnostic signals such as perplexity or representational similarity. The goal of this work is to remove layers that are irrelevant or actively detrimental for a specific task, so that a model becomes both cheaper to run and, in many cases, more accurate, without retraining or changing model weights.

Key Contributions

1. The TALE algorithm: a simple, greedy, iterative layer-elimination procedure that, at each step, evaluates all possible single-layer removals and permanently deletes the layer that maximizes validation accuracy, repeating until performance falls below a user-defined threshold. It is hardware-agnostic, requires no retraining, operates entirely at inference time, and adaptively determines how many layers to remove rather than fixing that number in advance.

2. Broad empirical validation: evaluation across 9 benchmarks and 5 model families (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, showing that TALE matches or surpasses baseline performance while reducing computational cost.

3. Direct comparison to prior training-free pruning: TALE is compared against SLEB, SparseGPT, Wanda, and SliceGPT on LLaMA-2-7B and LLaMA-2-13B, and against task-aware variants of SLEB and BlockPruner that are given the same validation data and metric. The paper also shows that substituting TALE's objective with cosine similarity or perplexity reproduces the degradation seen in those prior methods.

4. Evidence that TALE composes with other adaptation methods: experiments combining TALE with few-shot prompting (on Lucie and LLaMA) and with LoRA fine-tuning, including prune-then-fine-tune and fine-tune-then-prune orderings, plus a Mutual Information analysis of why layer removal can help.

Main Findings

• Accuracy and efficiency can improve at the same time: across all models and benchmarks, TALE produced consistent accuracy improvements after layer removal, showing that selectively removing task-misaligned layers can improve rather than degrade performance. The paper frames the central empirical result as: task-aware layer pruning can simultaneously improve downstream performance and reduce model depth.

• Largest single reported gain: on LLaMA 3.1 8B, GSM8K-HARD improved from 39.0% to 59.0% with only 1 layer dropped.

• Reasoning tasks benefit most: gains ranged from 23% to 51% on MATH500 and GSM8K across models. Knowledge-intensive tasks improved more modestly, though LLaMA showed an 11% gain on BIG-Bench.

• Task-specific magnitudes on ARC-Challenge: improvements were modest for LLaMA (+1.6%) and more pronounced for Qwen 2.5 7B (+6.3%).

• Same behavior under two evaluation protocols: results were consistent under LM-Eval (highest probability among provided options) and the paper's Decoder Eval (accuracy computed from the model's generated answer), indicating improvements are not specific to one evaluation protocol.

• Intermediate layers can outperform the final layer: the authors report that for many tasks, projecting intermediate representations (k < L) through the output projection can yield higher accuracy than the final layer, which they interpret as evidence of representational noise or redundancy in later layers.

• Robustness across seeds: variance across five seeds was consistently low for LLaMA, Qwen, Lucie, and Mistral, and LLaMA 3.1 8B showed the overall lowest variance, indicating TALE does not depend on favorable initialization.

• Validation set requirements are modest: TALE needs 500 to 1500 examples; once the validation set exceeds 500 examples, the set of dropped layers stabilizes across all tasks.

• TALE outperforms prior training-free pruning: on LLaMA-2-7B and LLaMA-2-13B across four zero-shot benchmarks, TALE achieved the highest accuracy while maintaining sparsity equivalent to SLEB, which performed second best. SparseGPT, Wanda, and SliceGPT performed worse.

• Task-agnostic objectives fail: when cosine similarity guided TALE on ARC-Easy, it dropped 2 layers and reduced accuracy from Llama's baseline of 79.5 to 58.5 with a time speedup of 1.32. The paper reports similar drops when perplexity is used as the optimizing objective.

• Comparison against task-aware variants: under Decoder Eval, TALE reached 76.7 on ARC-Easy, 54.3 on ARC-Challenge, and 73.1 on Winogrande, versus SLEB-ta at 61.0, 38.0, and 66.5, and BlockPr-ta at 64.6, 39.6, and 65.59. Under LM-Eval, TALE reached 81, 55, and 78, versus BlockPr-ta at 65, 41, and 66.

• Early-exit comparison (limited): from the limited information available, the paper notes RAEE gives a score of 65% on ARC-Easy, while the Llama baseline is 76% and TALE on ARC-Easy gives 79% using LM-Eval.

• Latency and throughput: measured on one A100-80GB NVIDIA GPU with identical decoding settings, TALE improved first-token latency in 9/9 settings (macro average −14.3%) and throughput in 9/9 settings (macro average +17.9%), using the BEST variant rather than the BSBA variant.

• Cost of running TALE itself: pruning time scales as O(I · L · V · T_layer) for I iterations, L layers, and validation set size V. For LLaMA 3.1 8B (L = 32, V ≈ 500–1500), pruning took approximately 1 to 2 GPU-hours on a single A100.

• Synergy with fine-tuning: on LLaMA 3.1 8B Winogrande, the baseline scored 53.83 with 0 layers deleted, TALE-only scored 56.67 with 4 layers deleted, FT-only scored 85.00, TALE→FT scored 87.06 with 4 layers deleted, FT→TALE scored 86.74 with 7 layers deleted, and (TALE→FT)→TALE scored 87.37 with 8 layers deleted. Comparable staged improvements appear on MMLU, CommonQA, and GSM8K, and on Qwen 0.5B for Winogrande and MMLU.

• Fine-tuning is also cheaper after pruning: pruning LLaMA-3.1 8B before fine-tuning reduced training time by approximately 18.5% (2–2.5 GPU hours on an A100) while improving Winogrande performance by +2.4%. Pruning the fully fine-tuned model yielded a 7-layer reduction while maintaining strong accuracy (86.66%).

• Layer importance is task-specific: removing early layers reduces accuracy to near zero on commonsense reasoning tasks, whereas removing LLaMA's layer 3 improves performance on GSM8K-Hard. Mathematical reasoning tasks benefit from pruning one to three early-to-middle layers (LLaMA layer 3, Mistral layers 6 and 22, Lucie layer 12). Knowledge-intensive tasks (ARC, BoolQ, CommonsenseQA, Winogrande, BIG-Bench) show more modest gains and benefit more from removing later layers.

• Over-pruning is sharply penalized: a specific number of removed layers n may be optimal, while pruning further (n+1) can cause a sharp drop, and this optimal point varies across datasets, which the authors argue makes any fixed pruning budget potentially suboptimal or harmful.

• Pruning dynamics are consistent in shape: the first pruning step often yields a noticeable improvement, followed by smaller gains or mild fluctuations, after which performance typically decreases monotonically. The authors note they did not observe any cases where performance recovered after falling below the baseline.

• Redundancy is not only a multi-task artifact: training a transformer from scratch on a single task (in-context learning of linear functions) still left several layers redundant, and some degraded performance.

• Pretraining-scale hypothesis: Lucie was trained on 3T tokens versus 15T for LLaMA and 13T for Qwen. The authors hypothesize that models trained close to their performance ceiling (large-scale pretraining, instruction tuning, or RLHF) gain less from TALE, while models trained under more limited objectives benefit more. Lucie also tolerated more aggressive pruning than the other models.

• Mutual Information analysis: estimated with MINE, MI profiles showed many layers with pronounced drops. TALE removes some, but not all, of those layers, reducing peaks and valleys in the MI profile. Removing all layers associated with MI decreases leads to very poor performance, suggesting some local MI decreases are necessary for proper functioning.

• Multilingual observation: initial testing on Lucie, which is tuned for French conversational proficiency, using bilingual versions of the same dataset, indicated that optimal pruning is task-specific rather than language-specific.

Methodology in Plain English

TALE works by trial and error on a validation set. Starting from the full pre-trained model, it tries removing each layer one at a time, measures validation accuracy for every resulting candidate, and permanently deletes the single layer whose removal gave the best accuracy. That compressed model then becomes the starting point for the next round, and the process repeats until accuracy drops more than 8% below the original baseline. The 8% slack lets the search pass through slightly worse intermediate models in case a later removal recovers performance. The algorithm then returns the most compressed model still above the threshold.

Removal is implemented with a custom wrapper that keeps the embedding layer, final normalization, and language modeling head, filters out the deleted layers while preserving the order of the rest, and updates the configuration's layer count. The result remains compatible with Hugging Face and LoRA fine-tuning without custom training loops.

Evaluation uses two approaches: the standard LM-Eval framework, which picks the highest-probability option among provided choices, and the paper's Decoder Eval, which extracts the final answer from the model's generated response and compares it to ground truth. The authors argue LM-Eval systematically inflates scores, ignores hallucination, and compresses differences between models.

Fine-tuning experiments used PEFT with LoRA (rank 64, alpha 16, dropout 0.1) over all linear layers, 4-bit NF4 quantization with float16 compute, paged AdamW (32-bit), learning rate 2e-4, weight decay 0.001, cosine schedule with 0.03 warmup, gradient clipping at 0.3, 10 epochs, effective batch size 60 (per-device batch size 2 with 30 gradient accumulation steps), sequences up to 300 tokens, right-side padding with padding token ID 128004, and packing disabled. All experiments ran on NVIDIA A100 GPUs, and all pruning experiments on 1 NVIDIA A100 with 80GB memory. The code is available at https://github.com/omyokun/tale/.

Why This Matters

The paper challenges the

Authors’ abstract

Large Language Models (LLMs) typically come with a fixed architecture, despite growing evidence that not all layers contribute equally to every downstream task. We introduce TALE (Task-Aware Layer Elimination), an inference-time method that improves task performance by selectively removing layers that are irrelevant or detrimental for a given task. TALE optimizes task-specific performance, yielding a task-optimized architecture without retraining. Across 9 tasks and 5 model families, under both zero-shot and few-shot settings, TALE consistently matches or surpasses baseline performance while simultaneously reducing computational costs. TALE also synergizes with fine-tuning, leading to further performance improvements. Computing TALE for a new task requires modest resources, making it a practical and deployable solution for task-specialized LLM inference.