TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability
Not on the map. Why?
Abstract (arXiv)
Recent work has promoted task-aware layer pruning as a way to improve model performance on particular tasks, as shown by TALE. In this paper, we investigate when such improvements occur and why. We show first that, across controlled polynomial regression tasks and large language models, such pruning yields no benefit on in-distribution (ID) data but consistently improves out-of-distribution (OOD) accuracy. We further show empirically that OOD inputs induce layerwise norm and pairwise-distance profiles that deviate from the corresponding ID profiles. This leads to a geometric explanation of task-aware pruning: each task induces a task-adapted geometry, characterized empirically by the representation profiles observed on ID inputs. OOD inputs can introduce a distorted version of the task-adapted geometry. Task-aware pruning identifies layers that create or amplify this distortion; by removing them, it shifts OOD representational norms and pairwise distances toward those observed on the adapted distribution. This realigns OOD inputs with the model's task-adapted geometry and improves performance. We provide causal evidence through controlled distribution shifts and residual-scaling interventions, and demonstrate consistent behavior across model scales.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 26 of its 32 references that are in the map. First posted on arXiv 14 May 2026.
Builds on
26 of 32 references.
References · 26
- 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,084
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
- 2019BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questionscited by 2,140
- 2016Understanding intermediate layers using linear classifier probescited by 1,279
- 2023SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shotcited by 1,090
- 2019BERT Rediscovers the Classical NLP Pipelinecited by 1,046
- 2016Deep Variational Information Bottleneckcited by 915
- 2021Probing Classifiers: Promises, Shortcomings, and Advancescited by 831
- 2015Deep Learning and the Information Bottleneck Principlecited by 793
- 2023A Simple and Effective Pruning Approach for Large Language Modelscited by 776
- 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
- 2018MINE: Mutual Information Neural Estimationcited by 680
14 more
- 2017Opening the Black Box of Deep Neural Networks via Informationcited by 494
- 2024SliceGPT: Compress Large Language Models by Deleting Rows and Columnscited by 378
- 2025Layer by Layer: Uncovering Hidden Representations in Language Modelscited by 293
- 2024SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blockscited by 96
- 2022Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scalecited by 64
- 2023Dynamic Sparse No Training: Training-Free Fine-tuning for Sparse LLMscited by 61
- 2018Entropy and mutual information in models of deep neural networkscited by 50
- 2014Efficient Estimation of Mutual Information for Strongly Dependent Variablescited by 46
- 2024BlockPruner: Fine-grained Pruning for Large Language Modelscited by 35
- 2023Large language models implicitly learn to straighten neural sentence trajectories to construct a predictive representation of natural languagecited by 32
- 2026Neural Thickets: Diverse Task Experts Are Dense Around Pretrained Weightscited by 25
- 2025Demystifying the Roles of LLM Layers in Retrieval, Knowledge, and Reasoningcited by 15
- 2021Greedy-layer Pruning: Speeding up Transformer Models for Natural Language Processingcited by 11
- 2025TELL-TALE: Task Efficient LLMs with Task Aware Layer Eliminationcited by 2
Built on by
Not cited yet.