TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

TL;DR

This paper investigates task-aware layer pruning, showing it does not improve in-distribution performance but enhances out-of-distribution accuracy by realigning model geometry with task-adapted representations.

Contribution

It provides a geometric explanation for task-aware pruning's effectiveness on OOD data and demonstrates consistent improvements across models and controlled experiments.

Findings

Pruning improves OOD accuracy but not ID performance.

OOD inputs distort representation geometry compared to ID inputs.

Removing layers that cause distortion realigns OOD data with task geometry.

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.