PaperlayerView exact revision on arXiv

Source-linked AI summary

Train What You Deploy: Closing the MLP Reachability Gap in Low-Rank Clone Distillation

Wenhui Chen, Zhifeng Li, Jie Zhou, Navan Preet Singh, Madalina Ciobanu, Chenghua Wang, Qingqing Mao, Ritankar Das

arXiv:2609.02006v1cs.LGcs.CL

TL;DR

LRC deploys full-width MLP weights but restricts training to a teacher-induced slice, leaving deployed degrees of freedom unreachable. The paper trains the entire deployed matrix from the identical warm start through mergeable Dense-LRC and CORE-LRC realizations without changing inference structure. Across three teachers this recovers up to +10.45 Avg9, with controls supporting the enlarged reachable set as the source of the gain.

  • Problem

    LRC pays to deploy full-width MLP matrices while training can reach only a teacher-induced slice, leaving 62.5–81.4% of their independent degrees of freedom unreachable.

  • Method

    Dense-LRC and CORE-LRC train the entire deployed matrix from the identical LRC warm start, then merge to the same single deployed weight without changing deployed shape, parameter count, or inference FLOPs.

  • Results

    +2.36/+2.71/+10.45 Avg9 over matched-budget plain-LRC baselines across three teachers, with Qwen reaching the original ∼20B-token accuracy at 10B tokens.

  • Takeaways & Limitations

    Controls support attributing the gains to the enlarged reachable set rather than added parameters or the surrounding recipe, and the conclusion extends the audit beyond LRC.

  • Takeaways & Limitations

    Results are single-seed, MMLU drops, and the study varies dmodel under inherited dff rather than optimizing dff itself.

Abstract

from arXiv · show

A compressed student has two shapes that need not agree: the weight it deploys at inference and the weight family its training can reach. We show that a state-of-the-art weight-inheritance distiller, Low-Rank Clone (LRC), deploys a full-width student MLP but ties training to a teacher-induced slice, leaving 62.5-81.4% of each deployed matrix's independent linear degrees of freedom unreachable-paid for at inference, never trainable. Our principle is one line: train what you deploy. From the identical LRC warm start, we make the training object the entire deployed matrix, with no change in deployed shape, deployed parameter count, or inference FLOPs, via two mergeable realizations (Dense-LRC and CORE-LRC) that both collapse to one deployed weight. This recovers stranded capacity: taking the stronger realization per teacher, +2.36/+2.71/+10.45 Avg9 over matched-budget plain-LRC baselines across three teachers (Llama3.2-3B, Llama3.1-8B, Qwen2.5-3B), with the largest gain on the widest teacher (Qwen), where it reaches the original recipe's approx. 20B-token accuracy at 10B tokens (2x token efficiency); there the strictly same-lineage arm still recovers +6.39, the fully controlled figure. Controls strongly support attributing the gain to the enlarged reachable set, rather than to added parameters or the recipe. From approx. 10B distillation tokens plus a short SFT, a half-parameter 1.5B student matches its approx. 9T-token teacher's 9-task macro-average, within evaluation noise and with a residual MMLU deficit, and a 2.7B student beats Meta's own official compression of Llama3.1-8B at ~900x fewer compression tokens (a token count under unmatched recipes, not a compute claim). All results are from single-seed runs on the LRC backbone.

1 Introduction

LRC deploys full-width student MLP matrices while training reaches only a teacher-induced slice, leaving much deployed capacity inaccessible. Training that entire deployed matrix from the identical warm start recovers this capacity without changing deployed shape or inference cost.

  • Results: A 1.5B student matches its ∼9T-token teacher’s 9-task macro-average, while a 2.7B student beats Meta’s official same-lineage compression at ∼900× fewer compression tokens.The token comparison is under unmatched recipes and counts compression-stage tokens.
  • Train what you deploy: Training the entire deployed matrix from the identical LRC warm start preserves deployed shape, parameter count, and inference FLOPs.Dense-LRC and CORE-LRC are mergeable realizations that collapse to one deployed weight.
  • Controlled attribution: Controls support attributing the gains to the enlarged reachable set rather than added parameters, the surrounding recipe, or faster convergence.An equal-parameter slice-confined arm recovers none of the gain, while the gain persists under a stripped generic recipe.
  • The reachability gap: LRC trains only a teacher-induced subspace of each full-width deployed MLP matrix, leaving 62.5–81.4% of independent degrees of freedom unreachable.The gap incurs full inference cost without deployment savings.
  • Results: +2.36/+2.71/+10.45 Avg9 over matched-budget plain-LRC baselines across three teachers, with the strongest realization selected per teacher.The strictly same-lineage dense arm alone gives +2.23/+2.71/+6.39.

2 The Inherited Width and Its Stranded Complement

LRC inherits the teacher’s wide MLP intermediate dimension but trains only a student-dimensional teacher slice, creating a teacher-dependent stranded complement. The paper formalizes this as a deterministic reachability gap and tests whether opening it recovers useful capacity.

  • Intervention: Training the deployed matrix recovers +6.39 Avg9 for the same-lineage dense arm and reaches 63.44 on Qwen, matching the original ∼20B-token baseline at half the tokens.The table reports 10B PT, 0-shot Avg9, with stronger realizations selected per teacher.
  • Stranded complement: 62.5%, 71.4%, and 81.4% of the intermediate space are inaccessible at ρ = 2.67, 3.50, and 5.375, respectively.These are excluded matrix degrees of freedom for one non-square projection, not unused neurons.
  • Formalization: The LRC reachable set RLRC = {TZ⊤} is a strict subset of the deployed hypothesis space Hdeploy, so the gap is fixed by parameterization rather than initialization, objective, or optimizer.Finite-budget optimization within the reachable set can still depend on coordinates.

3 Related Work

The paper builds on compression-distillation and pruning while distinguishing its parameterization diagnosis from related activation-rank, adapter, and low-rank-sufficiency claims. Its central distinction is between restricting updates within a fixed family and opening the inherited width that LRC never trains.

  • Compression-distillation and pruning: LRC combines low-rank projection, teacher-weight soft pruning, and cloned activations to compress trillion-token-trained models with about 20B tokens.Structured-pruning methods such as Minitron and Sheared LLaMA instead remove depth or width, including MLP intermediate dimensions.
  • Concurrent diagnoses: RED diagnoses effective-rank collapse in hidden representations, whereas this paper removes LRC’s deterministic projected-family constraint.The reported activation-aware initialization gain is +0.68 Avg9, versus +6.39 to +10.45 from opening the reachable set.
  • Subspace and null-space adapters: Teacher-spectral updates resemble constrained SVD-defined subspace updates, but the reported gain does not depend on that mechanism because several full-matrix parameterizations recover it.The setting differs from PEFT because the goal is to add usable capacity rather than protect pretrained knowledge during finetuning.
  • Low-rank sufficiency: Claims that low-rank distillation is sufficient concern trainable-update rank within a fixed parameterization, while this paper studies inherited dff that LRC retains but never opens.The paper makes no claim that low-rank methods are inferior in general.

4 Train What You Deploy

Full-MLP LRC repairs LRC’s deployment–training reachable-set gap by training the entire deployed MLP matrix from the identical warm start, without changing inference shape or cost.

  • The reachable-set gap: LRC’s projection is confined to the teacher-occupied column space, leaving the deployed MLP’s structural complement unreachable during training.The excluded complement is opened by training the full deployed matrix.
  • What the construction changes: Both realizations preserve the deployed shape, parameter count, and FLOPs while training the full deployed MLP matrix from the LRC warm start.The intervention is therefore a structural reparameterization rather than an inference-time expansion.
  • Dense-LRC: Dense-LRC trains a full dff × r matrix initialized at the merged plain-LRC weight, so the student is exactly the LRC model at step zero.Its reachable column space is the full R^dff.
  • CORE-LRC: CORE-LRC opens the same complement through a zero-initialized teacher-spectral path and merges to the same single deployed dff × r weight.The teacher-aligned basis makes CORE-LRC functionally equivalent to Dense-LRC while changing the coordinates optimized by AdamW.
  • Claim scope: The paper frames the contribution as an auditable diagnosis and repair of LRC’s reachable-set gap, with teacher-aligned coordinates enabling measurement of optimization outside the inherited subspace.The claim is not that the approach beats unrestricted dense training.

5 Experiments

Across three teacher–student compressions, Full-MLP LRC improves matched-budget LRC, with the largest gain on Qwen’s widest MLP and evidence that continued training alone does not close the gap.

  • Cross-teacher results: +2.36/+2.71/+10.45 Avg9 are obtained over matched 10B-token plain-LRC baselines across Llama3.2-3B, Llama3.1-8B, and Qwen2.5-3B.The headline delta takes the stronger realization for each teacher.
  • Llama3.2-3B→1.5B: +2.36 Avg9 is achieved on Llama3.2-3B→1.5B, while short SFT reaches 66.21 Avg9 versus the teacher’s 66.18.The aggregate match is within the setting’s approximately 0.4 Avg9 evaluation standard error, although MMLU remains lower.
  • Llama3.1-8B→2.7B: +2.71 Avg9 lifts Llama3.1-8B→2.7B from 65.55 to 68.26 at 10B PT, and short SFT reaches 68.93.The student’s comparison with Meta’s compression is a same-lineage token-count comparison, not a controlled recipe or compute claim.
  • Controls: 62.86 Avg9 is the 50B-token no-complement LRC plateau, while the matched 10B baseline remains at 63.04, both below Full-MLP LRC’s 65.40 at 10B.The continued-training control supports a capacity gain rather than merely faster convergence.
  • Comparison boundaries: The 2.7B student scores 68.93 versus Meta’s 66.18, but the comparison uses different recipes, data, and post-training.Only compression-stage token budgets are counted, so the result is not a compute comparison.
  • Evaluation details: The Llama3.2-3B→1.5B Dense-LRC arm reaches 65.95 after SFT versus teacher 66.18, while the table reports 65.27 PT and +2.23.Table 2 also specifies a short 0.62B-token, 8-dataset instruction stage for +SFT rows.

6 Isolating the Reachable Set

Matched controls show that the performance gain comes from expanding LRC’s training-reachable set, not from a privileged basis, extra parameters, or the recipe. Full-matrix training also engages complement directions that are functionally important, especially in early layers.

  • Basis controls: Basis controls reach similar performance, showing that the gain does not require a privileged teacher-spectral coordinate system.Dense-LRC reaches 65.27 Avg9 versus 65.40 for the teacher-spectral arm; an independent data-shuffle seed again keeps them within noise.
  • Equal-parameter control: The equal-parameter teacher-slice control recovers essentially none of the gain, supporting reachable-set expansion rather than added coordinates or optimizer state.It reaches Avg9 64.02 versus 64.04 for plain LRC and approximately 1.1 points below the same-lineage full-matrix arm at 65.16.
  • Recipe-independence: +1.82 under LRC and +2.86 under GPD show that full-matrix training beats slice training across both recipes.The Qwen factorial matches 2B PT tokens and seed 1234 while crossing reachable set with recipe.
  • Functional engagement: Full-MLP training creates complement components that are functionally engaged, while plain-LRC weights remain almost entirely in the teacher column space.The full student places 6.2% of gate and 8.3% of up weight energy in the orthogonal complement on average, versus approximately 0.1% for plain LRC.
  • Functional engagement: Removing the complement component collapses performance, with the first three layers accounting for most of the −31.74 Avg9 drop.The collapse indicates co-adaptation with the complement, not that those directions independently encode capability.
  • Attribution scope: The core comparisons are strictly matched pairs, but the Qwen +10.45 headline delta and CORE-versus-Dense margin are cross-lineage.The controlled reachable-set claim rests on the same-lineage dense +6.39 result.

7 Limitations

The evaluation is limited by single-seed training, one compression backbone, a single hidden-dimension design axis, no MoE coverage, and a narrow task scope.

  • Attribution: Single-seed runs limit attribution strength, and evaluation standard errors capture evaluation noise rather than run-to-run variance.The Qwen headline +10.45 is cross-lineage; the strictly matched dense +6.39 supports the controlled claim.
  • Scope: The study uses only the LRC backbone, so whether other prune–distill methods contain the same reserve remains untested.External baselines provide context rather than a distillation-SOTA comparison.
  • Design axis: LRC compresses only hidden dimension, leaving the optimal intermediate width dff uncharacterized.The hidden-width sweep lacks per-width no-complement baselines, making its MMLU interpretation suggestive.
  • Behavioral scope: The paper makes no zero-forgetting claim because single-layer orthogonality does not control the end-to-end Jacobian and MMLU drops.
  • Architecture scope: There is no MoE coverage because per-expert expansion ratios are near or below 1, leaving little structural complement.

8 Conclusion

Training the entire deployed matrix recovers capacity stranded by LRC without increasing inference cost, and matched controls attribute the improvement to the enlarged reachable set. The paper proposes auditing this deployment–training mismatch in other compression backbones.

  • Conclusion: +2.36/+2.71/+10.45 Avg9 over matched-budget plain-LRC baselines are obtained at zero added inference cost.The stronger realization per teacher is reported, alongside 2× token efficiency and a half-parameter student matching its teacher.
  • Conclusion: Controls attribute the gain to reachable-set expansion rather than added parameters or the training recipe.
  • Broader implication: The paper generalizes its audit: compare the dimension training can reach with the dimension deployment pays for, then justify or close any gap.It expects analogous reserves wherever a width is inherited rather than deleted, while identifying other backbones as future work.

A CORE-LRC: The Teacher-Spectral Realization

CORE-LRC reparameterizes the full deployed MLP weight in a teacher-aligned orthonormal basis, adding a zero-initialized complement path that merges back to the original deployed shape. Its basis matches Dense-LRC on narrow teachers but performs better on the wide Qwen target at a fixed token budget.

  • Construction: CORE-LRC augments LRC with a complement path whose zero initialization preserves the exact plain-LRC model at the start.For down projections, the dual construction fills the right complement.
  • Deployment: At inference, the two paths merge into one dff × r weight with the plain-LRC shape, adding no parameters or FLOPs.The construction therefore changes training reachability without changing the deployed object.
  • Reachable space: Because the teacher and complement bases are complete and orthonormal, CORE-LRC spans the entire deployed weight space.It is functionally equivalent to Dense-LRC, differing only in the coordinate system optimized by AdamW.
  • Basis conditioning: On matched narrow-teacher settings, Dense-LRC and CORE-LRC tie at 65.27 versus 65.40 Avg9, indicating basis-independence there.
  • Basis conditioning: Avg9 63.44 for CORE-LRC (+10.45 over the matched plain-LRC baseline) versus 59.38 for Dense-LRC (+6.39) appears on the wide Qwen target at 10B tokens.The CORE result matches the original recipe’s approximately 20B-token baseline, 63.43, at half the tokens; the cross-realization margin is cross-lineage and qualitative.

B Additional Result Tables

Additional sweeps show that opening the full MLP makes Avg9 rise with student hidden width, while MMLU remains width-limited and reaches the teacher only near native width. Because widening changes several capacity axes, the width result is suggestive rather than cleanly isolated.

  • Hidden-dimension sweep: MMLU remains at 54.67 at d(S)model=1536 despite full-matrix training.It reaches the teacher’s 60.56 only near the native width.
  • Interpretation: The residual MMLU gap is associated with overall student width, not the already-open MLP complement.More MLP-tail capacity does not close the gap once the complement is open.
  • Caveat: Widening d(S)model moves several capacity axes simultaneously, so the result is suggestive rather than a clean isolation.Matched no-complement baselines at each width are identified as future work.

C Implementation Details

The implementation trains mergeable full-width MLP realizations from LRC-compatible initialization and evaluates them under a fixed distillation and evaluation setup. Training adds substantial optimization-time parameters and some spectral-reconstruction cost, while deployment remains unchanged.

  • Objective: The objective combines hidden-state MSE, attention Q/K/V/O cosine alignment, and MLP-intermediate MSE, with identical reported PT weights across LRC variants.The fixed settings are wKL=wNTP=1.0, T=40, and λaux=0.2.
  • Reparameterization: CORE-LRC decomposes each non-square MLP weight into the teacher-occupied range and structural complement, updating the zero-initialized complement.The realization is mergeable and is intended to reclaim hidden-dimension capacity.
  • Training and evaluation: Training uses ≈10B FineWeb-Edu tokens plus ≈0.35B OpenHermes tokens for PT and ≈0.62B tokens for SFT, evaluated with standard 0-shot accuracy and Avg9.Avg9 averages nine named tasks, including MMLU.
  • Training-time cost: Dense-LRC and CORE-LRC roughly double trainable parameters during training, but merging restores the plain-LRC deployed parameter count and FLOPs.The 3B setting rises from ≈0.40B to ≈1.06B trainable parameters; CORE-LRC adds ≈12% step time over Dense-LRC and 7–28 GB of frozen bases.
  • Result tables: +2.71 Avg9 / +4.55 MMLU is reported for the matched Dense-LRC Llama3.1-8B→2.7B arm at 10B PT.The 2.7B student exceeds the cited official reference at ≈900× less PT under unmatched recipes.
  • Result tables: +10.45 Avg9 is reported for teacher-spectral CORE-LRC on Qwen2.5-3B→1.7B, matching the original ∼20B-token baseline at half the tokens.Canonical Dense-LRC recovers +6.39 in the strictly matched same-lineage comparison.
  • Functional engagement: The first three layers contain 12–16% complement energy and their removal causes −31.74 Avg9, while an identity pass is lossless.Plain-LRC measures ≈0.1% complement energy, indicating functional engagement of the opened coordinates.
  • External context: Published external comparisons are contextual because teachers, budgets, and sizes differ, and PT-token counts exclude the short SFT.The cited table reports 10B-PT+SFT students and adds ≈0.62B SFT tokens.
Loading 2609.02006v1…