Demystifying When Pruning Works via Representation Hierarchies
| Authors | Shwai He et al. |
| Year | 2026 |
| HF Upvotes | 19 |
| arXiv | 2603.24652 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Network pruning, which removes less important parameters or architectures, is often expected to improve efficiency while preserving performance. However, this expectation does not consistently hold across language tasks: pruned models can perform well on non-generative tasks but frequently fail in generative settings. To understand this discrepancy, we analyze network pruning from a representation-hierarchy perspective, decomposing the internal computation of language models into three sequential spaces: embedding (hidden representations), logit (pre-softmax outputs), and probability (post-softmax distributions). We find that representations in the embedding and logit spaces are largely robust to pruning-induced perturbations. However, the nonlinear transformation from logits to probabilities amplifies these deviations, which accumulate across time steps and lead to substantial degradation during generation. In contrast, the stability of the categorical-token probability subspace, together with the robustness of the embedding space, supports the effectiveness of pruning for non-generative tasks such as retrieval and multiple-choice selection. Our analysis disentangles the effects of pruning across tasks and provides practical guidance for its application. Code is available at https://github.com/CASE-Lab-UMD/Pruning-on-Representations
Engineering Breakdown
Plain English
This paper investigates why network pruning—removing unimportant parameters to reduce model size—works well for non-generative tasks like classification but fails for generative tasks like language generation. The authors decompose language model computation into three sequential representation spaces: embeddings, logits (pre-softmax), and probabilities (post-softmax). They discover that embeddings and logits are robust to pruning perturbations, but the nonlinear softmax transformation from logits to probabilities is highly sensitive to pruning-induced noise, causing performance collapse in generative settings. This representation-hierarchy perspective provides a mechanistic explanation for when and why pruning succeeds or fails across different task types.
Core Technical Contribution
The core novelty is a principled framework for analyzing pruning effects through representation hierarchy: decomposing model computation into three distinct spaces (embedding → logit → probability) and measuring sensitivity at each stage. Rather than treating pruning as a black-box optimization problem, the authors identify that the softmax bottleneck—the nonlinear transformation from logits to probability distributions—is the critical failure point in generative models. This is the first work to pinpoint that robustness to pruning is not uniform across a model's computational pipeline; information bottlenecks at probability space cause cascading failures in token sampling and autoregressive decoding. The framework enables predictive diagnosis: measuring representation perturbations in early spaces can predict whether pruning will succeed or fail before deployment.
How It Works
The method analyzes how pruning-induced perturbations propagate through three sequential transformation stages: (1) Input tokens → embedding space via token embeddings and transformer layers, capturing dense semantic information; (2) Embedding space → logit space via output projection heads, converting representations into class-specific scores; (3) Logit space → probability space via softmax, normalizing scores into probability distributions for token selection. For each stage, the authors measure how pruning—removing neurons, attention heads, or weights—perturbs intermediate representations using metrics like KL-divergence and cosine distance. They apply this analysis to both discriminative tasks (classification) and generative tasks (language generation), measuring how small perturbations in logit space get amplified by the softmax nonlinearity into large probability distribution shifts. The key empirical finding is that embedding and logit spaces show bounded perturbations even under aggressive pruning (e.g., 50% weight removal), but probability distributions shift dramatically, causing incorrect token sampling in generation.
Production Impact
For engineers deploying pruned language models in production, this work provides a diagnostic tool: before pruning a model for a specific task, measure representation perturbations in the logit space to predict success probability. For generative applications (chatbots, translation, summarization), this suggests that naive magnitude pruning is unsafe and alternative strategies are needed—either pruning only non-critical layers, using pruning-aware fine-tuning on generative objectives, or hybrid approaches that protect probability-space bottlenecks. For non-generative tasks (sentiment analysis, classification), aggressive pruning remains safe and delivers real efficiency gains (memory reduction, inference latency). The implication for model serving: generative models may need 20-30% higher parameter counts than classification models post-pruning to maintain quality, changing cost-quality trade-offs. Teams should incorporate representation-space diagnostics into their pruning pipelines as a gating check before production deployment.
Limitations and When Not to Use This
The paper's analysis is restricted to transformer-based language models; applicability to CNNs, RNNs, or multimodal architectures remains unexplored. The framework identifies softmax as a critical bottleneck but does not propose concrete mitigation strategies beyond avoiding pruning—engineers still need separate techniques (careful layer selection, magnitude-aware thresholds, training-based pruning) to successfully compress generative models. The work assumes pruning happens post-hoc after full training; structured pruning during training or pruning guided by task-specific fine-tuning may show different representation dynamics and are not analyzed. Additionally, the paper does not quantify the practical overhead of measuring representation perturbations at inference or provide hardware-specific cost estimates, limiting adoption guidance for resource-constrained deployment scenarios.
Research Context
This paper advances the pruning literature by moving beyond task-agnostic pruning heuristics (magnitude, gradient-based) toward mechanistic understanding of when pruning fails. It builds on recent work analyzing internal representations in transformers (e.g., representation probing, information bottleneck theory) and applies these lenses to the pruning problem. The work relates to broader efficiency research (quantization, knowledge distillation) by showing that not all compression techniques are equally applicable across task types—a crucial insight for practitioners. It opens research directions on pruning-aware training, probability-space regularization, and task-specific compression strategies that account for representation robustness.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::