PaperlayerView exact revision on arXiv

Source-linked AI summary

Repair Changes Which Layers to Remove: Exact Optima for Depth Pruning of Language Models

Jianchang Su, Dawei Xiang, Yifan Zhang, Wei Zhang

arXiv:2610.30552v1cs.DC

TL;DR

Depth pruning removes transformer layers, but the distance between heuristic layer choices and the best subset is unknown, and the best choice may depend on repair. The paper evaluates subsets exactly and develops selection methods that account for the repair. Without repair, removals are distributed across runs; linear repair makes the optimum contiguous, and unrepaired choices miss it by up to 189%.

  • Problem

    The distance from existing heuristic layer choices to the best subset is unknown, and proxies score layers in the unrepaired model even though repair changes their interactions.

    Source: §1 · pp. 1–2

  • Method

    The paper enumerates layer-group subsets with and without ReplaceMe repair, then selects on repaired candidates or scores repaired blocks using a second-order expansion of the loss.

    Source: §6.2 · p. 6

  • Results

    Published methods choosing layers without repair miss the repaired optimum by up to 189%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Selection should evaluate candidates under the repair that will be applied; a second-order score can select a repaired block from intact-model gradients.

    Source: §8 · p. 8

    Exact optima are available only at two-layer granularity up to six removed groups and single-layer granularity up to three; the study tests one repair and English text and multiple-choice tasks.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Depth pruning removes entire transformer layers to make large language models faster, and recent methods recover quality by replacing the removed layers with a training-free linear map. Which layers to remove, however, is chosen with heuristic scores whose distance from the best choice is unknown. We evaluate every subset of layer groups exactly (about 53,000 pruned models across four families, with and without repair) and find that the best layers to remove depend on the repair: linear repair absorbs the interaction among nearby layers and turns the optimum from a split layout into a contiguous block, so methods that choose layers without the repair miss the repaired optimum by up to 189%. We therefore select layers under the repair that will be applied. A greedy search at the cost of SLEB lowers perplexity by 20% to 66% relative to ReplaceMe on five models up to 27B parameters, and with a task objective it is up to 25 points more accurate given the same data. A second-order expansion of the loss yields a score that selects the repaired block from one forward and one backward pass of the intact model.

1 Introduction

Depth-pruning methods must choose which layers to remove, but existing proxies score layers in the unrepaired model and their distance from the optimum is unknown. Exact evaluation shows that the optimum changes with linear repair, motivating selection on the repaired model that will be deployed.

1 Introduction

  • Linear repair makes the optimum a block, or a block and one short run, while the unrepaired optimum spreads removals across separate runs.A least-squares map absorbs interactions within a removed run.Source: §1 · p. 2
  • About 53,000 subsets across four model families were evaluated exactly, with and without ReplaceMe repair.The study uses layer groups of two for the main enumeration.Source: §1 · p. 2
  • Published methods that choose layers without repair miss the repaired optimum by up to 189%.The gap between a proxy and the optimum depends on the repair applied to the removed layers.Source: §1 · pp. 1–2
  • A greedy search that evaluates repaired candidates costs as much as SLEB and reaches the exact optimum or comes within 0.2% at every enumerated budget on three of four models.The paper also reports 20% to 66% lower perplexity than ReplaceMe on four models and 54% to 56% lower on a 27B hybrid-attention model.Source: §1 · p. 2
  • At 4 to 6 removed groups, the best contiguous block without repair is 18% to 588% worse than the optimum.A few critical layers have joint loss increases below the sum of their individual increases, while local super-additivity among other layers decays over 3 to 6 layers.Source: §1 · p. 2
  • With a task objective, the search is more accurate at every budget, by up to 25 points on the 27B model.The evaluation also includes tasks unseen by the objective, C4, LoRA recovery, and a 27B model not used for design choices.Source: §1 · pp. 2–3

2 Related work

Prior work selects layers using proxies, search, or contribution estimates, while repair methods generally act after selection or do not score the repaired model. This paper computes exact optima, analyzes interactions, and develops a score for repaired blocks.

2 Related work

  • Layer-selection work spans similarity scores, greedy or evolutionary search, and game-theoretic contribution estimates.Kim et al. (2026) call the subset problem NP-hard and report little gain from search over one-shot scores without repair.Source: §2 · p. 3
  • Repair methods include least-squares linear maps, interface corrections, low-rank or shifted maps, and submodule selection by replacement residual.These methods repair after layer choice or use a criterion that does not evaluate the repaired model; ReplaceMe’s cosine criterion is 11% to 189% worse than the repaired optimum across block lengths and four families.Source: §2 · p. 3
  • The repair Fisher score applies a second-order loss expansion to the residual remaining after repair, weighting it by downstream Fisher information.The paper relates this weighting to balanced truncation, where the loss gradient plays the role of observability; unweighted residual and first-order rankings are inconsistent across models.Source: §2 · p. 3

3 Problem setup

The setup partitions a decoder-only model into contiguous groups and represents removed runs with fitted linear maps. It defines the loss and evaluation data, then enumerates subsets to compare search methods with exact optima.

3 Problem setup

  • Layers are partitioned into contiguous groups of one or two, and a run is a maximal set of adjacent removed groups.Removing a set skips its layers; a run is replaced by a linear map on the MLP output of the last kept layer before it.Source: §3 · p. 3
  • Each removed run gets its own map, folded into the down-projection of the preceding kept layer; a run starting at the first layer cannot be repaired.The map is fitted using intact-model activations at calibration tokens.Source: §3 · p. 3
  • The selection loss is mean token loss on 16 WikiText-2 calibration sequences, while perplexity is reported on 16 other test sequences.Maps are fitted on 128 WikiText-2 training sequences; accuracy averages five multiple-choice tasks, with BoolQ reported separately.Source: §3 · pp. 3–4
  • Every subset of up to six removed groups is evaluated for four models, with and without repair; Qwen3-14B is evaluated up to four groups.The study also enumerates Qwen3-1.7B at single-layer granularity up to three removed layers and stores subset losses to evaluate search against the exact optimum.Source: §3 · p. 4

4 The structure of the exact optimum

Removal losses are non-additive: nearby non-critical groups interact, while critical groups have strongly negative pairwise interactions. Without repair, exact optima split removals across runs, and contiguous blocks can be much worse.

Interaction decomposition

  • Interaction shares rise from 21% to 50% on Qwen3-1.7B, 15% to 31% on Gemma, 7% to 25% on Llama, and 11% to 22% on Qwen3-14B.These ranges correspond to 2 to 6 removed groups, except Qwen3-14B, where they cover 2 to 4.Source: §4 · p. 4

Critical layers

  • A group is critical if removing it alone more than doubles perplexity; every exact optimum keeps these groups.There are 4 of 14 critical groups in Qwen3-1.7B, 5 of 17 in Gemma, 3 of 16 in Llama, and 3 of 20 in Qwen3-14B; they are usually among the first layers.Source: §4 · p. 4

Local super-additivity

  • Among non-critical groups, pairwise interactions are positive and decay with depth distance, with fitted decay lengths of 2.9 layers on Gemma to 6.4 on Llama.Pairwise terms explain 52% to 76% of non-additivity at 3 removed groups; higher-order terms among nearby groups grow at heavier budgets.Source: §4 · p. 4, Figure 3 · p. 4

A contiguous block is close to the worst choice

  • At 4 or more removed groups, every exact optimum splits into 2 or 3 runs, while the best contiguous block is 18% to 588% worse in perplexity.A block contains the most pairs within the interaction range, so local super-additivity raises its interaction cost.Source: §4 · p. 4, Table 1 · p. 4

Implications for existing methods

  • One-shot scoring is 33% above the optimum on Qwen3-1.7B and 581% on Gemma at 4 removed groups; SLEB finds the optimum in 13 of 20 budgets.A pairwise model can be 189 to 226,000 times worse, while the interaction remains when removed layers are replaced by their mean update.Source: §4 · pp. 4–5

5 Linear repair makes the optimum contiguous

Linear repair absorbs interactions among layers in a removed run, changing the exact optimum from split runs to a contiguous block or a block plus a short run. Layer selection must therefore evaluate the repaired model that will be deployed.

Repair absorbs the interaction inside a run

  • At the heaviest budget, repair lowers interaction share from 62% to 36% on Qwen3-1.7B, 66% to 52% on Gemma, 57% to 28% on Llama, and 25% to 15% on Qwen3-14B.Separately repaired runs still interact: their summed repaired loss explains subset loss with R2 of 0.8 to 0.9 at 2 removed groups, but only −0.09 to 0.44 at 6.Source: §5 · p. 5

The repaired optimum is contiguous

  • Without repair, every optimum at 4 or more removed groups splits into 2 or 3 runs; under repair, the optimum is one block or one block plus another run.A contiguous block shifts from close to the worst choice without repair to close to the best choice with repair.Source: §5 · p. 5

Selection must account for the repair

  • Repairing the best plain-removal subset is 5% to 33% worse than the repaired optimum, while ReplaceMe’s cosine-distance block is 11% to 189% worse.On Llama at 8 of 32 layers, the repair-selected subset is 79% lower in perplexity and 4.0 points more accurate with a task objective than ReplaceMe’s block.Source: §5 · p. 5

6 Selecting layers under repair

The paper selects layers by evaluating candidates under repair, then gives a second-order score for ranking repaired blocks from the intact model. The greedy method approaches the exact repaired optimum, while the score weights repair residuals by downstream loss sensitivity.

6.1 Greedy selection under repair

  • RGS finds the repaired optimum at every budget on Qwen3-1.7B and Qwen3-14B, and is within 0.2% on Llama and 3.4% on Gemma.It evaluates each candidate after fitting and applying repair maps, at the cost of SLEB; its objective can be any loss on the repaired model.Source: §6.1 · p. 5
  • With a task objective using crossentropy of the correct choice among the choices on 200 training items per task, RGS-T selects for accuracy.Source: §6.1 · p. 5

6.2 A second-order score from the intact model

  • The second-order score ranks repaired blocks using the repair residual and intact-model gradients, weighting residual directions by downstream Fisher information.The empirical Fisher uses observed tokens, with each position playing the role of a sample.Source: §6.2 · pp. 5–6
  • The least-squares map minimizes the Euclidean norm of the repair residual, whereas the Fisher-weighted score measures that residual according to downstream loss response.Unweighted residual and first-order rankings are inconsistent across models.Source: §6.2 · p. 6

7 Experiments

Across five models, repair-aware selection improves perplexity and task accuracy over existing criteria, while the repair Fisher score offers a cheaper way to select a block. The experiments also test calibration changes, recovery fine-tuning, ablations, and deployment cost.

7.1 Experimental setup

  • The experiments use four enumerated model families and Qwen3.8-27B, pruning 25% or 37.5% of layers and comparing methods with the same least-squares repair.Metrics are WikiText-2 perplexity, C4 perplexity, and average accuracy on five zero-shot tasks.Source: §7.1 · p. 6

7.2 Main results

  • 20% to 66% lower perplexity than ReplaceMe on four models, and 54% to 56% lower on Qwen3.8-27B; RGS-T beats same-data ReplaceMe by 0.9 to 14.1 accuracy points.RGS-T’s accuracy gain transfers to Qwen3.8-27B and tasks outside its objective; task-objective selection also beats ReplaceMe on ARC-Challenge and WinoGrande at both budgets on three of four models.Source: §7.2 · pp. 6–7, Table 3 · p. 7

7.3 The repair Fisher score

  • The repair Fisher score’s lowest-scoring block is within 5% of the best at 19 of 22 block lengths, while ReplaceMe’s cosine distance reaches that range at 2.Without evaluating a pruned model, its block is within 3% of RGS at 4 of 8 settings and beats ReplaceMe at 7 of 8; greedy search helps at heavier budgets.Source: §7.3 · p. 7

7.4 ReplaceMe’s default calibration

  • With maps fitted on SlimOrca conversations without a ridge, RGS has 16% to 73% lower perplexity; its perplexity is lower at 7 of 8 settings and accuracy at 6 of 8.Figure 5 compares block-ranking scores with exact repaired loss increases and reports the share of block lengths where the lowest-scoring block is within 5% of the best block.Source: §7.4 · pp. 7–8, Figure 5 · p. 8

7.5 Recovery fine-tuning and a second corpus

  • After 200 LoRA steps on C4, RGS retains lower mean perplexity at 7 of 8 settings, with gains of 16% to 30% on Llama, 2.5% to 19% on Gemma, and 1% to 2% on Qwen3-1.7B.RGS also has lower C4 perplexity than ReplaceMe at both budgets on all four models despite selecting on WikiText.Source: §7.5 · p. 8, Figure 6 · p. 8

7.6 Ablations

  • Removing the repair from the search raises perplexity by 19% on average for SLEB, while one-shot scoring under repair is 21% worse and ReplaceMe’s criterion is 73% worse.The ablation reports the largest SLEB penalty as 41% and a fourfold increase when RGS’s selected layers are removed without repair.Source: §7.6 · p. 8, Table 4 · p. 8

7.7 Cost

  • Pruning 25% or 37.5% of layers gives 1.3× or 1.55× speedup; removing 24 of 64 layers from Qwen3.8-27B reduces memory from 53.8 to 35.6 GB.RGS evaluates as many pruned models as SLEB, while RFS needs one forward and one backward pass.Source: §7.7 · p. 8

8 Conclusion

The conclusion links the exact optimum to the repair and recommends evaluating candidate subsets under the repair intended for deployment. It frames testing other repair types and jointly optimizing selection and repair as future work.

8 Conclusion

  • The exact pruning optimum depends on the repair: local super-additivity makes a contiguous block nearly worst without repair and nearly best with it.Evaluating candidates on the repaired model lowers perplexity 20% to 66% versus ReplaceMe at SLEB’s cost; a second-order expansion yields a score from the intact model.Source: §8 · p. 8
  • Evaluating each candidate subset on the repaired model intended for deployment is the paper’s stated selection principle.The conclusion also presents the second-order score as a way to assess repaired loss from the intact model.Source: §8 · p. 8
  • The core finding applies when a repair alters the cost landscape of removed layers; whether it holds for low-rank or shifted repairs remains for future work.The repair Fisher score measures repair quality from the intact model alone, pointing toward joint optimization of layer selection and repair objective.Source: §8 · p. 8

Limitations

These appendices bound the study’s scope and report proof, algorithm, mechanism, and search analyses. The results include limits on exact enumeration and partial evidence for one explanation of local super-additivity.

Limitations

  • Exact optima are available only through six removed two-layer groups or three single layers; heavier budgets compare methods rather than exact optima.The structural results concern perplexity, and only ReplaceMe’s least-squares repair is tested.Source: Limitations · p. 9

A Proof of Proposition 1

  • The second-order expansion has an O(∥δ∥3) remainder when the downstream network is three times differentiable in the residual stream.The stated condition holds for transformer layers with smooth activations and normalisation away from zero norm; the empirical-Fisher and block-diagonal approximations are also described as standard pruning approximations.Source: A Proof of Proposition 1 · p. 11

B Algorithms

  • RGS adds one layer at a time by evaluating each candidate subset on its repaired model, then returns the selected layers and folded maps.The first layer is always kept because a run starting there cannot be repaired; the first group is critical in every enumerated setting.Source: B Algorithms · pp. 11–12

C Mechanism of local super-additivity

  • Functional similarity has partial support as an explanation: update CKA predicts pairwise terms on Qwen3-1.7B and Llama, weakly on Qwen3-14B, and not on Gemma.The tests are listed in Table 5.Source: C Mechanism of local super-additivity · p. 12

D Search without repair

  • SLEB finds the optimum in 13 of 20 enumerated budgets; a refinement designed on Qwen3-1.7B finds it in 16 budgets in every seed.In three SLEB misses, one shift of a removed group would reach the optimum; the refinement improves on SLEB on Gemma.Source: D Search without repair · p. 12

E Additional results

Additional results examine MMLU selection, selection time, calibration-set size, and repair fitting. The task objective leads the MMLU results at most settings, while the reported repair-fitting changes are small relative to differences between selection methods.

MMLU

  • Each objective favours the metric it optimises: MMLU subsets have 1.7 to 4.9 times the WikiText perplexity of WikiText subsets.Source: E Additional results · p. 12
  • On Qwen3.8-27B, the MMLU objective reaches 57.6% at 16 layers against 54.8% for SLEB and 24.1% for ReplaceMe, and 31.2% at 24 layers against 30.2% for SLEB.Across settings, RGS with the MMLU objective is most accurate at 7 of 8; the exception is Qwen3-1.7B at 7 layers, where ShortGPT reaches 36.1% and RGS 31.3%.Source: E Additional results · p. 12

Selection time

  • On Llama-3.1-8B pruning 12 of 32 layers, selection takes 384 seconds for RGS, 343 for SLEB with repair, and 8 for ReplaceMe on one RTX A6000.Each repaired method also needs 29 seconds to collect map activations.Source: E Additional results · p. 12

Calibration size

  • With 16 calibration sequences, the calibration-set optimum is also the test-set optimum in 13 of 17 enumerated settings; agreement reaches 100% at 64 sequences on Qwen3-1.7B and Gemma-3-4B.Resampling from a pool of 64 indicates the disagreement is sampling noise.Source: E Additional results · p. 12

Repair fitting

  • Refitting later maps on the already repaired network did not improve on ReplaceMe’s map; the reported differences were small compared with selection-method differences.Fitting toward intact-network activations made multi-run subsets 1% to 16% worse on Qwen3-1.7B and Llama, while fitting toward removed-run outputs was about 5% better on Qwen3-1.7B and within 1.5% either way on Llama.Source: E Additional results · p. 12

F Per-task results

Per-task results show that the WikiText objective and task accuracy can favour different layer subsets. The task objective recovers some task-specific accuracy losses, while the reported outcomes vary across tasks and models.

Per-task results

  • Figure 7 compares each search method’s calibration perplexity with the exact optimum without repair at every enumerated budget.Randomised methods are averaged over 100 seeds.Source: Figure 7 · p. 13
  • At 25% removed on Llama, RGS-T reaches 67.6 on WinoGrande, versus 54.4 for RGS with the WikiText objective.The deep blocks selected by ShortGPT and ReplaceMe reach 71.1 and 69.7, respectively; the task objective has the best ARC-Easy accuracy on all four families at both budgets.Source: F Per-task results · pp. 12–13
  • RGS with the task objective has the best ARC-Challenge accuracy at 6 of 8 settings; Table 5 tests explanations for local super-additivity.The per-task table reports WikiText-2 perplexity and zero-shot accuracy across ARC-Easy, ARC-Challenge, HellaSwag, PIQA, and WinoGrande.Source: F Per-task results · p. 13, Table 5 · p. 13

G Implementation details

Implementation details describe precision, model-specific repair behavior, and the scope of supplementary tables and figures. The captions also report structural and perplexity comparisons for repaired selection methods.

Implementation details

  • Models run in bfloat16; repair maps are solved in float32 with a ridge, or in float64 with a pseudo-inverse without a ridge because Gemma’s SlimOrca MLP outputs are rank-deficient.For Gemma 3, the map acts on post-feedforward-normalisation output and cannot be folded into a weight, so the repaired model has one extra d × d matrix.Source: G Implementation details · p. 13
  • The supplementary SlimOrca tables describe ReplaceMe-calibration results, task-objective accuracy comparisons, and accuracy when the task objective uses three tasks.Table 7 uses SlimOrca selection and fit text; Table 8 compares against ReplaceMe’s criterion on the same task-training texts, and Table 9 reports ARC-Challenge and WinoGrande.Source: Table 7 · p. 13, Table 8 · p. 14, Table 9 · p. 14
  • The study reports fixing wrong Gemma attention types for removed or moved layers and a rank-assumption fault in unregularised map solving; all reported numbers use runs after the fixes.The data behind the tables and figures is exported to supplementary CSV files.Source: G Implementation details · p. 13
  • Figure 8 shows repair changing the heaviest-budget optimum from separate runs toward one block or a block and short run, at 17% to 67% lower perplexity.Light cells mark critical groups; grey denotes no repair and red denotes repair.Source: Figure 8 · p. 14
  • Figure 9 reports RGS perplexity 20% to 66% lower than ReplaceMe on four enumerated families and 54% to 56% lower on Qwen3.8-27B, with every method repaired by the same maps.The figure gives perplexity relative to ReplaceMe on a log scale.Source: Figure 9 · p. 14
  • Table 10 reports mean results over three C4 seeds for ReplaceMe and RGS, with seed spread of 0.01 to 0.3 in C4 perplexity, 0.06 to 0.66 in WikiText perplexity, and at most 0.5 accuracy points.Table 11 describes inference costs and speedups for folded RGS models; Gemma-3 is omitted because its maps cannot be folded.Source: Table 10 · p. 15, Table 11 · p. 15
  • Table 12 covers 5-shot MMLU accuracy for subsets in Table 3 and RGS with an MMLU objective; Table 13 gives per-task results behind Table 3.Table 14 specifies exact enumeration of subsets with and without repair, totaling about 53,000 subsets.Source: Table 12 · p. 15, Table 13 · p. 16, Table 14 · p. 17
Loading 2610.30552v1…