TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

• URL: http://arxiv.org/abs/2605.14738v1
• Date: Thu, 14 May 2026 12:01:05 GMT
• Title: TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability
• Abstract summary: Recent work has promoted task-aware layer pruning as a way to improve model performance on particular tasks.<n>We show that such pruning yields no benefit on in-distribution data but consistently improves out-of-distribution accuracy.<n>This leads to a geometric explanation of task-aware pruning.
• Score: 4.56250586729666
• License: http://creativecommons.org/publicdomain/zero/1.0/
• Abstract: 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.

Related papers

• Mitigating the ID-OOD Tradeoff in Open-Set Test-Time Adaptation [52.620992698137236]
  Open-set test-time adaptation (OSTTA) addresses the challenge of adapting models to new environments.<n>In this paper, we first analyze the limitations of entropy in OSTTA and then introduce an angular loss to regulate feature norm magnitudes.<n>Our method achieves strong OOD detection while maintaining high ID classification performance.
  arXiv  Detail & Related papers  (2026-04-02T04:00:57Z)
• Attribution-Guided Model Rectification of Unreliable Neural Network Behaviors [60.06461883533697]
  We leverage rank-one model editing to establish an attribution-guided model rectification framework.<n>We first distinguish our rectification setting from existing model editing, yielding a formulation that corrects unreliable behavior.<n>To target the primary source of misbehavior, we introduce an attribution-guided layer localization method.
  arXiv  Detail & Related papers  (2026-03-08T01:06:07Z)
• Feed Two Birds with One Scone: Exploiting Function-Space Regularization for Both OOD Robustness and ID Fine-Tuning Performance [72.57668440744301]
  We propose a novel regularization that constrains the distance of fine-tuning and pre-trained model in the function space with simulated OOD samples.<n>Our approach could consistently improve both downstream task ID fine-tuning performance and OOD robustness across a variety of CLIP backbones.
  arXiv  Detail & Related papers  (2025-08-31T12:14:34Z)
• Taylor Outlier Exposure [0.7673339435080445]
  Taylor Outlier Exposure (TaylorOE) is an OE-based approach with regularization that allows training on noisy OOD datasets contaminated with ID samples.<n>We demonstrate that the proposed method consistently outperforms conventional methods and analyze our regularization term to show its effectiveness.
  arXiv  Detail & Related papers  (2024-12-10T06:17:07Z)
• Model Reprogramming Outperforms Fine-tuning on Out-of-distribution Data in Text-Image Encoders [56.47577824219207]
  In this paper, we unveil the hidden costs associated with intrusive fine-tuning techniques. We introduce a new model reprogramming approach for fine-tuning, which we name Reprogrammer. Our empirical evidence reveals that Reprogrammer is less intrusive and yields superior downstream models.
  arXiv  Detail & Related papers  (2024-03-16T04:19:48Z)
• Unsupervised Out-of-Distribution Detection by Restoring Lossy Inputs with Variational Autoencoder [3.498694457257263]
  We propose a novel VAE-based score called Error Reduction (ER) for OOD detection. ER is based on a VAE that takes a lossy version of the training set as inputs and the original set as targets.
  arXiv  Detail & Related papers  (2023-09-05T09:42:15Z)
• AUTO: Adaptive Outlier Optimization for Test-Time OOD Detection [79.51071170042972]
  Out-of-distribution (OOD) detection aims to detect test samples that do not fall into any training in-distribution (ID) classes.<n>Data safety and privacy make it infeasible to collect task-specific outliers in advance for different scenarios.<n>We present test-time OOD detection, which allows the deployed model to utilize real OOD data from the unlabeled data stream during testing.
  arXiv  Detail & Related papers  (2023-03-22T02:28:54Z)
• Discover, Explanation, Improvement: An Automatic Slice Detection Framework for Natural Language Processing [72.14557106085284]
  slice detection models (SDM) automatically identify underperforming groups of datapoints. This paper proposes a benchmark named "Discover, Explain, improve (DEIM)" for classification NLP tasks. Our evaluation shows that Edisa can accurately select error-prone datapoints with informative semantic features.
  arXiv  Detail & Related papers  (2022-11-08T19:00:00Z)
• No Shifted Augmentations (NSA): compact distributions for robust self-supervised Anomaly Detection [4.243926243206826]
  Unsupervised Anomaly detection (AD) requires building a notion of normalcy, distinguishing in-distribution (ID) and out-of-distribution (OOD) data. We investigate how the emph geometrical compactness of the ID feature distribution makes isolating and detecting outliers easier. We propose novel architectural modifications to the self-supervised feature learning step, that enable such compact distributions for ID data to be learned.
  arXiv  Detail & Related papers  (2022-03-19T15:55:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.