IG-Pruning: Input-Guided Block Pruning for Large Language Models

AI summary

Overview

• Research area: Efficient inference for large language models (LLMs), specifically depth/block pruning.
• Technical level: Intermediate — assumes familiarity with transformer architecture (attention, FFN, KV cache) and standard pruning terminology, but the core ideas are explained clearly enough for a motivated reader.
• Scope: A single conference-style paper proposing an input-adaptive block pruning method that dynamically selects which transformer blocks to execute at inference time, validated on Llama-3 and Qwen-3 model families across six zero-shot benchmarks.

What This Paper Is About

Most depth pruning methods for LLMs compute one fixed set of layers to remove based on a global importance score, then apply that same pruned structure to every input. The authors show this static approach is flawed: different masks with identical perplexity can perform very differently on downstream tasks, meaning a single mask cannot serve all inputs well. IG-Pruning instead trains several specialized block masks and picks the best one per input at inference time based on semantic similarity.

Key Contributions

1. Input-aware dynamic block pruning. A framework that selects a layer mask per input rather than committing to one fixed pruned architecture, addressing the task- and input-dependence of static mask quality.
2. A two-stage pipeline combining clustering with learnable masks. Semantic clustering of calibration inputs produces cluster-specific masks trained via hard concrete distributions and L0 optimization, balancing candidate diversity (across clusters) with quality (within each cluster).
3. A training-efficient, router-free inference mechanism. Only mask parameters are optimized while base model weights stay frozen, and routing at inference reduces to a nearest-neighbor lookup in embedding space — no per-layer router networks or additional base-model training.
4. Fine-grained block-level granularity with attention-aware skipping. Attention and FFN blocks are masked independently, and skipping an attention block still computes K and V projections to preserve the KV cache while bypassing the expensive scaled dot-product operation.

Main Findings

• Consistent gains over static pruning baselines. IG-Pruning outperforms SLEB, ShortenedLlama, and EvoPress at every tested sparsity level on both Llama-3-8B and Qwen-3-8B.
• Large margins at moderate sparsity. On Llama-3-8B at 25% sparsity, IG-Pruning retains 87.18% of dense performance versus 76.32% for the best baseline — a 10.86 percentage point gap. At 12.5% sparsity it retains 98.29%, beating EvoPress by 6.36 points. On Qwen-3-8B at 13.9% sparsity it retains 96.01% versus 90.37% for the best baseline.
• Block-level pruning beats layer-level pruning. Independently masking attention and FFN components helps most around 20% sparsity; the advantage narrows at very high sparsity (above 40%), and on Winogrande specifically, layer-level pruning occasionally wins.
• Attention blocks are pruned more aggressively than FFN blocks. Visualizations show Llama's middle-layer attention blocks are frequently masked out, consistent with prior findings about representational similarity in those layers. Qwen-3 shows a more balanced attention/FFN distribution with fewer consecutive masked blocks.
• More clusters yield better performance. Going from 4 to 16 clusters improves average zero-shot scores, and the gap widens with sparsity — at 37.5% sparsity, 16 clusters beat 4 clusters by 10.64 points.
• Calibration data quality matters, and the gain is not a data artifact. Fineweb-edu (59.40 average) outperforms instruction data (58.20) and wikitext2 (55.85). Running SLEB on fineweb-edu did not improve its score, confirming the advantage comes from the method rather than privileged data.
• Meaningful FLOP reductions. At 25% sparsity, computation drops to 89.8% of the dense model; at 37.5% sparsity, to 75.8%.
• Cheap to train. Mask training takes roughly 15 minutes across 16 clusters at sparsity below 25%, and about an hour at 37% sparsity. Because base weights are frozen, an 8B model can be trained on a single RTX 3090.
• Generalizes across model sizes. Appendix results on Llama-3.2-3B, Llama-3.2-1B, and Qwen-3-4B show the same pattern: sizable gains at moderate sparsity, convergence with baselines at extreme sparsity.

Methodology in Plain English

The approach runs in two stages.

Stage 1 — building a menu of masks. The authors take a calibration corpus and encode each sentence into an embedding vector. K-means groups these embeddings into N semantic clusters (typically 16). For each cluster, they train a separate mask over all transformer blocks — one mask entry per attention block and per FFN block. These masks are parameterized with a hard concrete distribution, which lets the mask values stay continuous (and differentiable) during training but snap to 0 or 1 at inference. A Lagrangian penalty pushes each mask toward a target sparsity. Only the mask parameters are learned; the model's actual weights never change. Because each mask is trained only on its own cluster's data, it specializes in that cluster's input patterns.

Stage 2 — routing at inference. When a new input arrives, the same encoder embeds it, and the system finds the nearest cluster center by Euclidean distance. The mask belonging to that cluster is applied. Skipped FFN blocks are bypassed entirely. Skipped attention blocks still compute key and value projections so the KV cache stays intact for autoregressive generation, but the expensive query-key attention computation is skipped. The routing step is a single distance comparison — no learned router network, no extra training.

The design deliberately avoids the token-level routing used by prior dynamic methods, which the authors argue lacks sentence-level semantic coherence and incurs constant per-token overhead. Routing once per input at the cluster level is both cheaper and more globally consistent.

Why This Matters

Impact on research. The paper reframes depth pruning as a routing problem rather than a one-shot importance-ranking problem, and provides evidence that perplexity — the field's default saliency metric — can be misleading because different masks with similar PPL diverge on downstream tasks. It also shows that mask-only training with frozen weights is sufficient to obtain good pruned models, which lowers the compute barrier for pruning research.

Real-world applications:

• Serving LLMs under strict latency budgets where a fixed pruned model would be too lossy for some inputs but unnecessarily expensive for others.
• On-device or edge deployment, where a single 8B model can be mask-trained on consumer hardware and run with reduced FLOPs.
• Batch or mixed-workload inference, where semantically different request types get different computational paths automatically.
• Cost-controlled API deployments that need to trade accuracy for throughput dynamically without maintaining multiple separate model checkpoints.

Industry relevance. The value proposition is practical: near-dense accuracy at moderate sparsity, minuscule training cost relative to full retraining, no router network to maintain, and a routing decision cheap enough to run per request. For teams already doing post-training compression, this is an incremental layer on top of existing pipelines rather than a replacement.

Future Directions

• Reducing dependence on clustering and calibration data. Performance degrades when semantic clusters overlap or when calibration data lacks diversity (wikitext2 and instruction data both underperformed fineweb-edu), so more robust clustering or adaptive cluster counts are open problems.
• Testing factuality and hallucination effects. The authors explicitly flag that removing computational blocks may strip components needed for factual recall, and suggest pairing dynamic pruning with truthfulness methods such as TruthX or Truth-Aware Context Selection.
• Broadening evaluation beyond zero-shot benchmarks. All headline results are on six zero-shot tasks; generalization to reasoning, long-context, multimodal, or domain-specific settings remains unvalidated.
• Tuning sensitivity to hyperparameters. L0 regularization strength, Lagrangian multipliers, and cluster count all affect outcomes, and the paper does not offer a principled procedure for setting them across new models.

Target Audience

Researchers and engineers working on LLM inference efficiency, model compression, or adaptive computation. It is most useful to practitioners who already understand transformer internals and pruning basics, and who are evaluating whether dynamic routing is worth the added complexity over a simple static pruning baseline. Readers new to the area can follow the high-level argument but will need background on the hard concrete distribution and L0 optimization to engage with the method section in detail.

Authors’ abstract

With the growing computational demands of large language models (LLMs), efficient inference has become increasingly critical for practical deployment. Depth pruning has emerged as a promising approach for reducing the computational costs of large language models by removing transformer layers. However, existing methods typically rely on fixed block masks, which can lead to suboptimal performance across different tasks and inputs. In this paper, we propose IG-Pruning, a novel input-aware block-wise pruning method that dynamically selects layer masks at inference time. Our approach consists of two stages: (1) Discovering diverse mask candidates through semantic clustering and L0 optimization, and (2) Implementing efficient dynamic pruning without the need for extensive training. Experimental results demonstrate that our method consistently outperforms state-of-the-art static depth pruning methods, making it particularly suitable for resource-constrained deployment scenarios.