Source-linked AI summary
Demystifying When Pruning Works via Representation Hierarchies
Shwai He, Guoheng Sun, Haichao Zhang, Yun Fu, Ang Li
TL;DR
Pruning is attractive for reducing the costs of large language models, but its effectiveness differs across task types. The paper analyzes embedding, logit, and probability spaces and finds robust earlier representations but amplified probability perturbations that degrade generation while leaving non-generative tasks more effective.
Problem
Pruned models often perform well on non-generative tasks but frequently fail during generation, creating a need to understand when pruning succeeds across language tasks.
Method
The paper combines empirical and theoretical representation-level analyses of pruning perturbations across embedding, logit, and probability spaces.
Results
Embedding and logit spaces remain relatively robust, while nonlinear probability-space effects and temporal accumulation degrade generation; non-generative tasks remain comparatively effective.
Takeaways & Limitations
Pruning should be evaluated and applied according to target task type, with non-generative settings more suitable than generative ones under the analyzed conditions.
Takeaways & Limitations
The study focuses on training-free pruning; post-training or pruning-time fine-tuning is left for future work as a possible mitigation.
Abstract
from arXiv · showhide
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
1. Introduction
Network pruning improves efficiency but affects language tasks unevenly: non-generative performance often remains strong, whereas generation frequently degrades. A representation-hierarchy analysis attributes this discrepancy to perturbation behavior across embedding, logit, and probability spaces.
- Network pruning removes less important parameters or architectures to improve computational efficiency as language models scale.
- Pruned models often retain strong non-generative performance but frequently fail on generative tasks that sample from predicted probability distributions.
- The analysis decomposes language-model computation into embedding, logit, and probability spaces to trace pruning-induced perturbations.
- Embedding representations remain largely robust after substantial parameter removal, while the linear mapping to logits preserves comparable representational similarity.
- The nonlinear logit-to-probability projection amplifies perturbations, producing large output-distribution deviations that accumulate across generation steps and degrade generation quality.
- Stable categorical-token probability subspaces, together with embedding robustness, support pruning effectiveness for retrieval and multiple-choice classification.
2. Related Works
Large language models deliver broad capabilities but impose substantial efficiency costs. Network pruning addresses these costs by removing redundant or less important components through unstructured or structured compression.
- Scaling large language models has produced strong capabilities across tasks but creates significant memory and computational overhead.
- Network pruning reduces memory footprint and inference cost by removing less important model components.
3. Background on Language Modeling
Language models transform discrete tokens through continuous representations and latent spaces before producing vocabulary-wide token probabilities. Non-generative tasks use embeddings or a small candidate-token subset, whereas generation repeatedly feeds sampled outputs back into the model.
- Language models tokenize text into discrete indices and map those tokens to continuous embedding vectors.
- Deep neural layers transform embeddings into hidden representations, which are projected through the language-model head into logits.
- Softmax converts logits into a probability distribution over the vocabulary, from which the next token is sampled.
- Autoregressive generation feeds each generated token back with prior context, creating a feedback loop across decoding steps.
- Non-generative tasks produce embeddings or probabilities for a small candidate-token set without iterative decoding.
- Table 1 compares Mistral models after dropping eight attention or MLP layers across embedding, multiple-choice, and generative benchmarks.
4. Inconsistent Effects of Pruning
Pruning can preserve single-step decision quality while destabilizing multi-step generation. The discrepancy is associated with output dimensionality, nonlinear probability projection, and autoregressive error accumulation.
- Pruning operates through fine-grained intra-layer sparsification or coarse-grained inter-layer removal of transformer blocks.
- The same pruned architecture can show markedly different behavior on multiple-choice and generative tasks after layers are removed.
- Generative tasks face a substantially higher-dimensional vocabulary output space than non-generative tasks using embeddings or a small label set.
- Nonlinear projection to token probabilities can amplify pruning perturbations, while autoregressive decoding propagates early errors over time.
- Figure 3 evaluates Wanda pruning on HellaSwag and GSM8K under unstructured 50%, 4:8, and 2:4 sparsity patterns.
5. Hierarchical Effects of Pruning
The analysis tracks pruning-induced deviations through embedding, logit, and probability spaces using controlled layer replacement. Embedding and logit representations remain comparatively stable, whereas probability representations fluctuate substantially.
- Method: The study replaces one layer at a time with its pruned counterpart under a shared dense-model context to isolate representation shifts.Deviations are measured across decoding steps using angular deviation.
- Method: Pruning-induced deviations are evaluated by deriving logits and probabilities from embedding representations across the inference pipeline.This comparison characterizes how one perturbation evolves across representation spaces.
- Representation-space behavior: Embedding representations remain largely stable under layer dropping, except at the first and last layers.The first and last layers exhibit substantially larger transformations.
- Representation-space behavior: Probability representations fluctuate substantially despite comparable embeddings, while logit similarity remains comparable to embedding similarity.Thus, dimensionality increase from embeddings to logits alone does not explain the task-performance gap.
6. Representation-level Analysis
The representation-level analysis models pruning perturbations locally across embedding, logit, and probability spaces. It finds that the LM head preserves similarity, while softmax sensitivity can amplify relatively small logit deviations into large distributional shifts.
- Analysis framework: A Taylor-based local analysis studies how pruning-induced perturbations propagate and amplify across representation spaces.The analysis leverages localized layer-wise deviations.
- Embedding and logit spaces: The embedding-space approximation assumes the orthogonal perturbation component is sufficiently small and local, except at the first and last layers.This assumption supports the local cosine-similarity characterization.
- Embedding and logit spaces: The relative orthogonal magnitude is significantly reduced after the LM head, indicating that pruning perturbations remain limited in logit space.This is consistent with comparable logit similarity before and after pruning.
- Probability space: Softmax converts logits into probability distributions and amplifies differences even when the underlying logits remain relatively similar.Probability-space deviation is governed by temperature and the weighted variance of logit perturbations, with pronounced effects especially in later layers.
- Probability space: Theoretical estimates closely match ground-truth angular deviation and KL-divergence trends in the vocabulary space.The KL divergence highlights large distributional discrepancies between original and pruned models.
7. Multi-Scale Effects of Pruning
Across generation steps, pruning preserves early-step similarity but deviations persist and grow during autoregressive decoding. Non-generative tasks remain more robust because they use single-step decisions or stable categorical-token subspaces.
- Multi-scale task effects: With eight attention layers removed, the pruned model performs comparably on non-generative tasks but fails on generative tasks.The same setting exposes the contrasting effects of pruning across task types.
- Persistent divergence in generation: Cosine similarity is significantly higher at the first generation step than at later steps across all feature spaces.This supports pruning for tasks relying on embedding or logits at the first decoding step.
- Persistent divergence in generation: Differences in sampled tokens cause the baseline and pruned models to condition on diverging histories, sharply increasing later-step deviation.The first step remains low-deviation because both models receive the same prompt tokens.
- Persistent divergence in generation: Pruning produces persistently high divergence across decoding steps and substantially greater degradation in generative than non-generative tasks.The reported pattern links iterative decoding with progressively worsening output differences.
- Robustness of probability subspaces: Multiple-choice tasks depend on a small categorical-token subset, whose log-likelihood often preserves the same argmax despite shifts in top-token probabilities.Candidate tokens commonly lie in the distribution tail, where probability shifts are milder.
8. Discussion of Effective Pruning
Pruning effectiveness depends on representation space, task-relevant subspaces, and temporal dependence. The study’s conclusions apply to training-free pruning, while post-pruning training is left for future work.
- Representation Space: Embedding and logit spaces are relatively robust, making tasks that operate directly on them more amenable to pruning.The conclusion identifies representation space as a determinant of post-pruning performance.
- Task-Relevant Subspace: Task-specific probability subspaces can remain stable and preserve predictions even when the global probability distribution shifts.Many tasks use only low-dimensional subsets of the full vocabulary.
- Temporal Dependence: Temporal dependence compounds pruning errors during autoregressive generation, whereas single-step tasks avoid this amplification.This distinction helps explain greater robustness outside generation.
- Beyond Training-Free Pruning: The study focuses on training-free pruning; post-training or fine-tuning after pruning is proposed as a complementary mitigation approach for pruning-induced collapse.That approach is left for future work.
9. Conclusion
The paper shows that pruning robustness is task-dependent: large language models often remain effective on non-generative tasks but frequently fail in generative settings. A representation-hierarchy analysis explains how robustness varies across embedding, logit, and probability spaces.
- Pruning preserves performance more reliably on non-generative tasks than on generative tasks.
Impact Statement
Pruning-induced perturbations remain comparatively manageable in embeddings and logits but become more consequential in probability space, helping explain divergent effects across language tasks. The analysis further identifies error paths through context-dependent operators that can accumulate across decoding steps.
- Task impact: Pruning mainly affects probability space, explaining why generation degrades while non-generative performance remains stable.The paper frames this discrepancy through embeddings, logits, and probabilities, and evaluates generative, multiple-choice, and retrieval settings.
- Probability-space analysis: The probability shift is modeled by comparing p = softmax(z/T) with q = softmax((z + ∆z)/T).Here z is the original logit vector, ∆z is the compression-induced perturbation, and T is the temperature.
- Probability-space analysis: Softmax compresses the probability-space nonlinearity into a single log–sum–exp expectation term.The resulting closed form is used to characterize how compression shifts the vocabulary probability distribution.
- Logit-space analysis: Without softmax, cosine deviation is governed by the perturbation magnitude in the subspace orthogonal to the logits under uniform weighting.The corresponding logit-space analysis contrasts with the probability-dependent reweighting in the softmax case.
- Error propagation: Self-attention creates value and weight error paths, whereas operators without historical dependency involve only parameter perturbations and the current input.For non-history-dependent operators, pruning-induced deviations do not include accumulated representation errors from previous steps.
E.3. Error Sources in Autoregressive Decoding
During autoregressive decoding, pruning errors include both current-step parameter effects and deviations accumulated from historical representations. Because generated tokens alter later attention contexts, these errors can propagate and compound across decoding steps.
- The decoding deviation depends on both effective parameter perturbations and accumulated historical representation perturbations.
- Self-attention converts perturbations in past activations into additional contributors to the current output.
- Unlike nonhistorical operators, self-attention introduces an error source driven by historical representations.
- Prompt-token perturbations are fixed after prefill, whereas generated-token perturbations depend on prior decoding deviations.
- Generated-context deviations cause later self-attention to use different historical contexts, allowing pruning errors to propagate and compound.
- Compressed models can suffer severe degradation in generative settings even on relatively simple prompts, indicating failures are not merely due to task difficulty.
H. Ablation Study on Temperature Factors
The ablations validate the analysis across temperature settings, compression methods, layers, and representation spaces. Theoretical estimates track measured deviations, while probability-space deviations are larger than embedding- and logit-space deviations.
- Estimated cosine-similarity and KL-divergence trends consistently align with ground-truth measurements across temperature settings.
- Quantization produces higher representation similarity and lower deviations than pruning because it approximates parameters rather than removing them entirely.
- Wanda pruning shows a layer-wise representation-similarity trend similar to layer dropping, although magnitudes differ across pruning strategies.
- In probability space, the estimator closely tracks ground-truth KL-divergence and angular-deviation trends, with deeper layers showing larger deviations.
- Embedding and logit spaces exhibit substantially smaller angular deviations, with theoretical curves closely aligned to ground truth; the first and last layers are exceptions.
- The LM head projection substantially reduces relative orthogonal energy, consistent with limited pruning-induced perturbations in logit space.
- The visualizations support progressive amplification of pruning-induced perturbations across representation spaces and distinguish linear from nonlinear transformation effects.