TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Technical Summary: TAPIOCA – Why Task-Aware Pruning Improves OOD Model Capability

Problem Statement
Recent work has demonstrated that removing layers from Large Language Models (LLMs) at inference time—specifically through task-aware pruning methods like TALE (Task-Aware Layer Elimination)—can improve performance on specific tasks without retraining. However, the mechanisms driving these gains remain unclear. A central question is whether these improvements stem from general compression benefits, regularization effects, or specific interactions between the model's internal representations and the input distribution. Furthermore, it is unknown whether these gains are universal or specific to certain types of data shifts.

Methodology
The authors investigate this phenomenon using a combination of controlled synthetic experiments and large-scale evaluations on fine-tuned LLMs.

1. Controlled Regression Setting: The study begins with a 12-layer transformer trained from scratch on in-context linear regression. The training distribution is explicitly defined (coefficients sampled from $U(-1, 1)$). The authors evaluate the model on both in-distribution (ID) data (same distribution) and out-of-distribution (OOD) data (coefficients sampled from wider intervals $U(-\sigma, \sigma)$ where $\sigma > 1$). They apply TALE, a greedy pruning algorithm that removes layers to maximize validation accuracy on a specific distribution.
2. Large-Scale LLM Evaluation: The findings are extended to fine-tuned LLMs (Llama 3.1 8B and GPT-OSS 120B) specialized in mathematical reasoning and code generation. Here, the "adapted" distribution is the fine-tuning task (e.g., MATH500), and OOD tasks include general benchmarks like MMLU and BoolQ.
3. Geometric and Causal Analysis: To explain why pruning works, the authors analyze the geometry of hidden states. They measure layer-wise representational norms and pairwise token distances. They also employ local linear surrogates to approximate layer transformations and conduct causal interventions, such as inserting inverse surrogate maps and rescaling residual updates, to isolate the effect of geometric distortion.

Key Findings

• OOD-Specific Gains: The primary empirical finding is that task-aware pruning yields performance gains only under distribution shift. When evaluated on the distribution to which the model is adapted (ID), TALE either fails to improve performance or actively degrades it by removing layers necessary for optimal precision. Conversely, on OOD data, TALE consistently identifies and removes layers that improve accuracy.
• Representation Geometry Mismatch: The authors demonstrate that OOD inputs induce a distorted representation geometry compared to ID inputs. Specifically, OOD inputs cause layer-wise hidden-state norms and pairwise distances to deviate systematically from the profiles observed on ID data. In the regression setting, OOD inputs cause distances to inflate across layers; in LLMs, OOD trajectories diverge from ID trajectories in middle-to-late layers.
• Amplification Mechanism: The study identifies that certain layers act as "distribution-dependent amplifiers." On ID inputs, these layers perform useful refinements. However, on OOD inputs, the same layers exhibit high-gain, low-rank amplification, pushing representations further away from the task-adapted geometry. TALE improves performance by removing these specific layers, thereby preventing the amplification of geometric distortion.
• Alignment, Not Just Contraction: Pruning does not simply suppress activation norms. In experiments where the direction of distribution shift is reversed (training on wide distributions, testing on narrow ones), pruning expands representation distances to align with the ID profile. The core mechanism is the realignment of OOD representations toward the geometry learned during adaptation, regardless of whether this requires contraction or expansion.
• Causal Evidence: Causal interventions confirm that geometric distortion is the driver of performance loss. Inserting an inverse surrogate map after a pruned layer (to undo the layer's geometric effect) restores OOD performance nearly to the level of the pruned model. Similarly, rescaling the residual update of a selected layer continuously improves accuracy as the layer's contribution is attenuated, ruling out topological artifacts or simple redundancy as the sole cause.

Significance and Claims
The paper claims to provide a novel geometric explanation for the success of inference-time pruning. It argues that layer importance is not fixed but is conditional on both the task and the input distribution. The authors posit that pretrained models encode multiple task-specialized functions, each associated with a local representation geometry. ID inputs fit this geometry, while OOD inputs violate it.

The significance of this work lies in:

1. Clarifying the Limits of Pruning: It establishes that task-aware pruning is not a universal performance booster but a corrective mechanism for distribution shift.
2. Geometric Interpretation: It shifts the understanding of layer redundancy from static parameter counts to dynamic, distribution-dependent geometric distortions.
3. Mechanistic Insight: It provides causal evidence that performance gains arise from correcting the trajectory of internal representations to match the model's adapted geometry, rather than from increasing model capacity or reducing noise.

The authors conclude that pruning improves OOD performance by attenuating layers that amplify the mismatch between the input distribution and the model's learned task geometry, effectively realigning the model's internal state with the conditions under which it was optimized.