Demystifying When Pruning Works via Representation Hierarchies
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 demystify how such discrepancies arise under pruning, we analyze network pruning from a representation-hierarchy perspective, decomposing the internal computation of language models into three sequential spaces: \textit{embedding} (hidden representations), \textit{logit} (pre-softmax outputs), and \textit{probability} (post-softmax distributions). While representations in the embedding and logit spaces are largely robust to pruning-induced perturbations, the subsequent nonlinear transformation from logits to the probability space amplifies such deviations, whose persistence across time steps leads to substantial degradation during generation. By 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 classification. Our representation-level analysis disentangles the effects of pruning across tasks and offers practical guidance for applying pruning effectively. The code is available in the \href{https://github.com/CASE-Lab-UMD/Pruning-on-Representations}{project repository}.
Lay Summary
Pruning removes parts of a large language model to make it cheaper and faster to run. However, pruning does not affect all tasks in the same way: pruned models may remain reliable on tasks with fixed answer choices, such as multiple-choice questions, while becoming much worse at generating text word by word. This paper studies why pruning has these different effects. We examine what happens inside a language model after pruning, following the computation from internal representations to the scores assigned to possible next words, and finally to the probabilities used for generation. We find that the earlier internal representations often remain stable, but small changes can be amplified when scores are turned into next-word probabilities. During generation, these changes can accumulate over many steps and reduce output quality. These findings help explain when pruning is likely to be useful and when compressed models need more careful evaluation before deployment.