In-Context Learning as a Rehearsal of Recovery for Depth Pruning

Depth pruning removes whole transformer blocks and then repairs the model. Existing criteria choose the blocks by the damage a removal does before repair. But they cannot tell which of two equally damaging removals repair will undo, because that depends on whether the remaining layers can still learn. Repairing and testing every candidate is out of reach (removing 9 of 36 layers alone leaves 94 million sets), so practice falls back on heuristics such as dropping a contiguous block of late layers. We instead find a better set with a few forward passes, drawing on the view that in-context learning is an implicit weight update. We therefore score a removal by how much it hurts the model's in-context learning ability. Over 1.6K prune-and-repair runs on five language models, we find that this probe helps predict post-repair quality better than present-damage scores alone, and more so with demonstrations than without, as our view expects. The same runs show why the choice matters. At a fixed budget, the set of removed layers changes post-repair quality far more than the choice among published criteria, and no single position is safe on every backbone. At 25% removal our method has the best five-model average under both repair objectives, even ahead of a baseline that briefly fine-tunes every candidate, for a few percent of pipeline cost. We also extend the study to vision-language models, where most prior depth pruning cuts only the language tower. We find that how the budget is split between the two towers decides which capabilities survive repair.