PaperlayerView exact revision on arXiv

Source-linked AI summary

XMerge: Cross-Axis Selection and Reconstructive Layer Merging for LLM Depth Compression

Jundong Hu, Shekar Ramachandran

arXiv:2609.02083v1cs.LGcs.CL

TL;DR

Existing depth compression can damage quality unpredictably even while preserving a standard serving architecture. XMerge combines cross-axis selection with local boundary reconstruction, and under aggressive removal it generally outperforms evaluated baselines while retaining a standard shallower transformer. Its main practical limitation is construction cost and limited validation scope.

  • Problem

    Depth compression preserves standard serving interfaces but can cause substantial, model-dependent quality degradation, with prior methods differing in both layer selection and layer absorption.

  • Method

    XMerge selects blocks quiet on relative-magnitude and angular axes, then re-fits an adjacent existing standard block to match the original two-block output without labels, end-to-end fine-tuning, or added inference parameters.

  • Results

    At k=4, XMerge ranks first on six of seven backbones on CORE and separately on six of seven on MMLU, while avoiding the extreme perplexity blow-ups seen in several competing operators.

  • Takeaways & Limitations

    XMerge’s advantage grows with more aggressive removal, and local reconstruction supplies most of the gain while cross-axis fusion helps when selection axes disagree.

  • Takeaways & Limitations

    The study uses one compression seed, covers only dense decoder-only Llama/Qwen models, and does not evaluate hallucination or refusal and safety behavior.

Abstract

from arXiv · show

Removing complete transformer layers preserves a standard serving architecture, but existing depth-compression methods can lose substantial quality, and the loss varies unpredictably across models. We introduce XMerge, a post-training method with two components. Cross-axis selection identifies a block with low relative-magnitude and angular hidden-state change, and local boundary reconstruction re-fits the adjacent surviving block to match the original two-block output. XMerge uses no task labels or end-to-end fine-tuning, and it introduces neither architectural changes nor additional inference-time parameters. Across seven Llama and Qwen backbones (0.5B-8B), five published baselines, and three layer-reduction levels, its advantage over baselines is largest at the most aggressive removal: at k=4 it ranks first on six of seven backbones on CORE (a 22-task aggregate) and, separately, on six of seven on MMLU (five of seven on both at once), while avoiding the large perplexity increases of several competing operators. In a task-level bootstrap, the 95% confidence intervals for the three largest CORE margins exclude zero; the remaining margins are consistent with ties. Across the 14 (model, regime) cells it is also the only evaluated operator that never collapses, ranking top-2 in both zero-shot and in-context regimes; on a first calibration probe (one backbone) it is the best-calibrated operator. Ablations show that local reconstruction provides most of the gain, while cross-axis fusion helps when the two selection axes disagree. The additional construction cost is recovered through per-token decode savings after roughly tens of thousands of requests.

1 Introduction

XMerge addresses depth-compression quality loss by separating block selection from adjacent-block merging. It combines cross-axis selection with local reconstruction while preserving a standard shallower transformer interface.

  • Depth compression reduces latency and KV-cache traffic while preserving hidden size, attention, vocabulary, and the standard serving interface.
  • Depth-compression methods consist of selecting a block to remove and choosing an operator that absorbs it into a neighboring block.
  • XMerge selects blocks quiet on both relative-magnitude and angular axes, then re-fits an adjacent surviving block to reproduce the pair’s original boundary states.
  • The method re-fits an existing standard block, producing an ordinary L−k-layer transformer with no new modules or added inference-time parameters.

2 Related Work

XMerge is positioned as serving-preserving post-training reconstruction within depth compression. It differs from prior operators in how the removed block is absorbed and in whether reconstruction uses an existing standard block or a new module.

  • Prior depth-compression work includes activation- or importance-based block removal and analytic neighboring-parameter collapse, while magnitude-based selection has precedent.
  • XMerge uses layer-local post-training calibration against the dense model’s own activations rather than instantaneous pruning or end-to-end recovery.
  • XMerge uniquely re-fits a full existing nonlinear block to the adjacent pair’s output while keeping the served model a module-free L−k transformer.
  • Table 1 compares operator design properties rather than accuracy, including absorption strategy, gradient fitting, and whether new inference modules are introduced.

3 Method: XMerge

XMerge decomposes depth compression into cross-axis block selection and reconstructive merging. It selects a quiet transition using relative-magnitude and angular signals, then locally fits one surviving block to the original adjacent-pair mapping.

  • Pipeline: The selector marks a block for removal, while the operator replaces its adjacent pair with one standard block, yielding L−k blocks after k merges.
  • Cross-axis selection: Low relative-magnitude score indicates small residual displacement, whereas high score indicates a strong activation-magnitude change.
  • Cross-axis selection: Low Block-Influence score indicates little directional change, whereas high score indicates substantial representation rotation.
  • Cross-axis selection: Cross-axis fusion combines relative-magnitude and Block-Influence signals as a safeguard when the two axes disagree.
  • Cross-axis selection: The selector standardizes the axes and uses their maximum, so a block is selected only when neither magnitude nor angular change is large.
  • Reconstructive merging: Reconstruction caches each pair’s input and original output, then distills the two-block mapping into one existing neighboring standard block.
  • Reconstructive merging: Adam runs for 300 steps at learning rate 10−5 on unlabeled WikiText-2 sequences, without downstream supervision, end-to-end optimization, or new inference parameters.
  • Multiple merges: For multiple merges, selection order comes from the original model while reconstruction targets are refreshed after each merge.

4 Experimental Setup

The evaluation spans seven 0.5B–8B backbones, three absolute layer-removal levels, five published baselines, and CORE, MMLU, and WikiText-2 perplexity metrics under a fixed protocol.

  • Models and compression levels: The benchmark covers Llama-3, Qwen2.5, and Qwen3 backbones ranging from 0.5B to 8B parameters.Depth reductions are evaluated at k=1, 2, and 4 removed layers.
  • Baselines and protocol: The comparison includes ShortGPT, LaCo, MKA, SWM, and CoMe under matched L−k architectures.Setting A evaluates the constructed checkpoints without follow-on recovery, fixing the deployed architecture, parameter count, and latency.
  • Metrics: CORE is the primary 22-task centered aggregate, supplemented by zero-shot MMLU and WikiText-2 test perplexity.CORE is used because it is more discriminative in this study and less directly tied to the reconstruction corpus than perplexity.
  • Reproducibility: All 126 backbone cells are complete under a single compression seed, fixed reconstruction hyperparameters, and float16 evaluation.The protocol comprises 7 backbones × 6 methods × 3 compression levels.

5 Main Results

XMERGE’s advantage grows with aggressive depth removal: at k=4 it leads on six of seven backbones for both CORE and MMLU, while avoiding the severe perplexity failures of several baselines. Bootstrap results support the largest margins but treat narrower differences as ties, and matched-depth comparisons give operators identical inference costs.

  • Main benchmark results: At k=4, XMERGE is strongest on six of seven backbones for CORE and separately for MMLU, winning both metrics on five of seven.The CORE win count rises from four, to five, to six of seven as k increases from 1 to 2 to 4.
  • Perplexity: Only XMERGE and LaCo avoid the extreme k=4 perplexity blow-ups of the other four methods, with XMERGE lower than LaCo on six of seven backbones.MKA illustrates the accuracy–perplexity trade-off with competitive CORE on Qwen3-1.7B but perplexity 440.
  • Uncertainty: The three largest k=4 CORE margins have 95% confidence intervals excluding zero: +.070, +.033, and +.048.The three narrow wins and the single loss have intervals spanning zero and should be read as ties; the six-of-seven sign test gives p=0.06.
  • Retention trend: At k=4, XMERGE retains 69% of dense CORE, versus 60% for the best competing merge baseline and 54% for ShortGPT.At k=1, XMERGE and ShortGPT each retain approximately 94%, while the best competing merge baseline retains approximately 88%.
  • Quality–latency trade-off: At matched depth, measured batch-1 decode speedup averages 1.04/1.09/1.17× for k=1/2/4, so quality differences are obtained at identical inference cost.The quality–latency Pareto frontier is therefore determined by quality retained per removed layer; XMERGE is best at the k=4 endpoint on six of seven backbones.

6 Component Analysis

The ablations separate reconstruction from selection and show that reconstruction supplies most of XMERGE’s gain, while cross-axis fusion improves robustness when selection axes disagree. Across zero-shot and in-context regimes, XMERGE avoids collapses and achieves the strongest mean scores.

  • 6.1 Reconstruction Provides Most of the Gain: At k=4, reconstruction raises CORE from .344 to .406 on Llama-3.2-3B and from .136 to .184 on Qwen3-0.6B under fixed RM selection.In both backbones, selected dropping falls below published baselines while reconstruction exceeds them.
  • 6.2 Cross-Axis Fusion Reduces Disagreement Risk: BI-only selection collapses Qwen3-1.7B with PPL 5.1×10^4, while plain dropping collapses Qwen3-0.6B with PPL 226.These failures motivate combining relative-magnitude and angular selection signals.
  • 6.2 Cross-Axis Fusion Reduces Disagreement Risk: On Qwen3-0.6B, fusion trades a small CORE decrease for better perplexity and MMLU, improving worst-case behavior rather than every metric.The authors therefore treat reconstruction as the primary operator contribution and fusion as a robustness mechanism.
  • 6.3 Zero-Shot vs. In-Context Regimes: Across 14 model-regime cells, XMERGE is the only evaluated operator with zero collapses and ranks top-2 in every zero-shot and in-context cell.A collapse is defined as centered CORE below 0.10.
  • 6.3 Zero-Shot vs. In-Context Regimes: XMERGE has the highest mean score in both regimes: .291 versus .236 for ShortGPT in zero-shot and .300 versus .242 in ICL.The Qwen3-1.7B and Qwen3-0.6B exceptions occur only in the ICL slice.

7 Recoverability Under a Shared Post-Compression Budget

Under a shared post-compression recovery budget, XMERGE provides the strongest or tied-strongest starting point and retains competitive recovered performance. This advantage is observed after applying the same follow-on WikiText-2 distillation procedure across methods.

  • Shared Recovery Budget: Under the shared follow-on budget, XMERGE starts from the strongest unhealed model at both evaluated scales.The recovery stage uses identical data, tokens, LoRA rank, learning-rate policy, steps, and teacher for each method.
  • Recovered Performance: XMERGE achieves the best recovered CORE on Llama-3.2-1B, with a +0.007 margin over the next method.This result is reported under the common post-compression recovery setting.
  • Recovered Performance: On Qwen3-8B, XMERGE ties SWM at CORE .358 versus .358, with an unrounded gap of approximately 0.0002.The reported gap is within single-seed noise, while XMERGE retains the best perplexity among CORE-competitive methods.

8 Discussion and Limitations

XMERGE’s advantage is strongest under aggressive compression, while its reliability extends across task regimes and calibration in the reported tests. The study also limits claims through single-seed evaluation, restricted model coverage, construction-cost accounting, and incomplete safety evaluation.

  • Discussion: XMERGE exceeds the strongest baseline in 15/21 cells using only label-free reconstruction and is best-or-tied under the shared recovery budget.The strongest single baseline is most often ShortGPT, while XMERGE’s advantage becomes more pronounced at aggressive removal.
  • Reliability: At k=4, XMERGE has the lowest ECE degradation on Llama-3-8B: ∆ECE=+0.010 over the dense model’s 0.125.All methods remain somewhat overconfident, so the result indicates preserved rather than fully restored calibration.
  • Limitations: The calibration result is limited to one backbone and one metric and does not establish a general safety guarantee.Hallucination and refusal or safety behavior are not evaluated for the base models studied.
  • Limitations: Construction costs range from minutes to 4.4 hours and are recovered after approximately 1.9k–24k requests, but total construction-plus-recovery compute is not claimed equivalent to training-free baselines.The study also uses a single compression seed, leaving reconstruction-seed variance unmeasured.
  • Limitations: The evaluation covers dense decoder-only Llama and Qwen models from 0.5B to 8B, not MoE, encoder–decoder, multimodal, or state-space models.Reconstruction uses WikiText-2, making WikiText-2 test perplexity partially in-domain.

C Task-Bootstrap Confidence Intervals on the CORE Margins

The task-bootstrap evaluates whether k=4 CORE margins over the strongest per-model baseline are separable from task-sampling noise. Three largest margins exclude zero, while narrower wins and the single loss remain consistent with ties.

  • Bootstrap Procedure: 20,000 bootstrap resamples of the 22 CORE tasks estimate each k=4 XMERGE-minus-baseline margin and its 95% confidence interval.The resampling unit is the paired per-task score difference, and the procedure does not rerun reconstruction.
  • Confidence-Interval Results: The three largest CORE margins are separable from task noise, while narrow wins and the single Qwen3-1.7B loss have confidence intervals including zero.These intervals quantify task-sampling uncertainty rather than reconstruction-seed variance.
  • Confidence-Interval Results: Across backbones, XMERGE is strongest on six of seven, with a one-sided sign-test p-value of 0.06.This result summarizes the direction of the per-model k=4 comparisons.
  • Perplexity Context: WikiText-2 perplexity values expose operator collapses, with XMERGE and LaCo the only methods avoiding the 10^2–10^4 blow-ups at k=4.XMERGE has lower perplexity than LaCo on six of seven backbones.

F Calibration Under Compression (ECE)

XMERGE preserves calibration better than the evaluated alternatives in the first MMLU probe, although compression still leaves every method overconfident and the evidence covers only one backbone and metric.

  • Calibration metric: ECE measures the gap between stated confidence and actual accuracy, providing a trustworthiness measure distinct from task accuracy.The probe evaluates Llama-3-8B at k=4 under recovery-free Setting A.
  • Calibration results: XMERGE has the lowest ECE degradation at +0.010, roughly half the next-best operator’s and an order of magnitude below MKA’s +0.110.MKA is also the most miscalibrated by MCE.
  • Scope: The calibration finding is limited to a single backbone and metric and does not establish a general safety guarantee.The reported probe uses Llama-3-8B on MMLU at k=4.

G Baseline Implementation Audit

The audit standardizes baseline implementations and clarifies how selector signals, recovery, task regimes, and matched-depth inference costs shape the comparison.

  • Baseline audit: All five baselines use official implementations or faithfully audited ports with the same target depth and WikiText-2 calibration corpus, without weight training.Each follows the algorithm in its cited publication.
  • Baseline audit: The audit records method-specific implementation details, including ShortGPT’s BI-guided dropping, LaCo’s parameter folding, MKA’s corrected top-down rule, SWM’s weighted averaging, and CoMe’s head-group concatenation.CoMe’s concatenation degenerates on Qwen2.5-0.5B because its two KV heads produce groups [2, 0].
  • Selector validation: RM is the strongest CORE predictor on both tested backbones, while BI is close behind; other criteria predict perplexity more reliably than CORE.WS is strong on Llama but weak on Qwen, and P, PMR, and Taylor reach |ρ|=0.6–0.8 for PPL but |ρ|≤0.33 for CORE.
  • Selector validation: RM and BI are near-tied as CORE predictors, so XMERGE fuses them for robustness rather than claiming RM dominates BI.Max-fusion changes RM-only selections on only two of seven backbones, where the axes decorrelate.
  • Recovery protocol: The uniform Setting B healer uses rank-16 LoRA knowledge distillation, merges adapters into weights, and adds zero inference-time parameters.The protocol uses the original FP16 model as teacher and WikiText-2 train data.
  • Recovery protocol: SWM-native recovery is excluded from fair comparison because it uses instruction tuning rather than the shared WikiText-2 KD budget.Its comparable row uses the same uniform KD as the other methods.
  • Evaluation protocol: CORE separates seven zero-shot tasks from fifteen few-shot or in-context tasks, while matched k gives every method the same L−k decoder and inference cost.Measured batch-1 speedup is method-independent and averages 1.04/1.09/1.17× at k=1/2/4.

L Construction-Time Measurements

XMERGE incurs minutes-to-hours of one-time construction cost, scaling approximately with the number of merges, but recurring decode savings amortize that cost after thousands of requests.

  • Construction cost: XMERGE construction costs minutes to hours and scales approximately linearly with k, unlike training-free baselines that take seconds to minutes.Each merge uses approximately 300 boundary-optimization steps.
  • Construction cost: Relative to SWM at k=4, XMERGE’s one-time cost is 3.7–22.6× larger in wall-clock time.The absolute cost remains minutes to 4.4 hours.
  • Amortization: Break-even occurs after 0.48M–6.1M generated tokens, equivalent to approximately 1.9k–24k requests of 256 tokens.The threshold is computed from construction time, dense decode rate, and measured speedup.
  • Amortization: After break-even, each further token is faster than dense, although a production engine that accelerates decoding more than construction could raise the threshold somewhat.Even a stated 3× understatement keeps break-even in the low tens of thousands of requests.
  • Measurement caveat: Some baseline construction times are nonmonotonic in k because of one-off loading, paging, and warm-up effects rather than algorithmic scaling.These effects do not affect reported quality numbers.
Loading 2609.02083v1…