# TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

arXiv: https://arxiv.org/abs/2605.14738 · HTML: https://mlsanity.com/paper/2605.14738/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 32 references that are in the map.
- First posted on arXiv 14 May 2026.

## Builds on (most cited first, 26 of 26)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Opening the Black Box of Deep Neural Networks via Information](https://mlsanity.com/paper/1703.00810/) (2017)
- [SliceGPT: Compress Large Language Models by Deleting Rows and Columns](https://mlsanity.com/paper/2401.15024/) (2024)
- [Layer by Layer: Uncovering Hidden Representations in Language Models](https://mlsanity.com/paper/2502.02013/) (2025)
- [SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blocks](https://mlsanity.com/paper/2402.09025/) (2024)
- [Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale](https://mlsanity.com/paper/2212.09095/) (2022)
- [Dynamic Sparse No Training: Training-Free Fine-tuning for Sparse LLMs](https://mlsanity.com/paper/2310.08915/) (2023)
- [Entropy and mutual information in models of deep neural networks](https://mlsanity.com/paper/1805.09785/) (2018)
- [Efficient Estimation of Mutual Information for Strongly Dependent Variables](https://mlsanity.com/paper/1411.2003/) (2014)
- [BlockPruner: Fine-grained Pruning for Large Language Models](https://mlsanity.com/paper/2406.10594/) (2024)
- [Large language models implicitly learn to straighten neural sentence trajectories to construct a predictive representation of natural language](https://mlsanity.com/paper/2311.04930/) (2023)
- [Neural Thickets: Diverse Task Experts Are Dense Around Pretrained Weights](https://mlsanity.com/paper/2603.12228/) (2026)
- [Demystifying the Roles of LLM Layers in Retrieval, Knowledge, and Reasoning](https://mlsanity.com/paper/2510.02091/) (2025)
- [Greedy-layer Pruning: Speeding up Transformer Models for Natural Language Processing](https://mlsanity.com/paper/2105.14839/) (2021)
- [TELL-TALE: Task Efficient LLMs with Task Aware Layer Elimination](https://mlsanity.com/paper/2510.22767/) (2025)

## 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.