PaperlayerView exact revision on arXiv

Source-linked AI summary

Matryoshka Language Model Suites

Nathan Godey, Yoav Artzi

arXiv:2608.09703v1cs.AIcs.CL

TL;DR

Training model suites independently is costly, and joint efficiency in training and serving them remains rarely studied. This paper introduces nested, jointly trained Matryoshka suites that match independent baselines while reducing training compute and improving speculative-decoding throughput.

  • Problem

    The efficiency of jointly training and serving language-model suites, whose models are typically trained separately, remains rarely studied.

  • Method

    Matryoshka training nests increasing-width and depth sub-models in one end-to-end architecture, exposing all sub-model logits for low-cost distillation.

  • Results

    36% less training compute achieves near-parity within 0.5 average points at every size with independently trained baselines on standard benchmarks.

  • Takeaways & Limitations

    The nested structure supports higher-throughput speculative decoding without memory overhead for the draft model and provides guidance for scaling Matryoshka suites.

  • Takeaways & Limitations

    Open questions remain about scaling to more sub-models and larger budgets and about interactions with instruction tuning, alignment, and reasoning-oriented finetuning.

Abstract

from arXiv · show

Training a language model suite classically requires training each model separately and serving them independently. We improve both training and inference efficiency by stacking sub-models of increasing size into a single nested architecture trained end-to-end. This Matryoshka training framework reduces the total parameter count of the suite, enables low-cost distillation from the largest to all smaller sub-models at every training step, and is well-suited for speculative decoding as the draft model is contained within the verifier. We validate our approach by training a Matryoshka suite comprising 500M, 1.5B, and 3B sub-models. Our suite is on par with independently trained baselines on benchmark performance and validation and out-of-domain perplexities, while using 36% less training compute and improving the throughput of speculative decoding by 14-26%. We also ablate key architectural choices, offering guidance for building strong Matryoshka LM suites.

1 Introduction

Matryoshka Language Model Suites nest independently detachable sub-models into one jointly trained architecture, reducing training cost while preserving baseline quality. The shared structure also benefits speculative decoding by sharing the KV cache and improving cross-model agreement.

  • Introduction: The framework provides models across a broad size spectrum for deployment tradeoffs because smaller models are more memory- and compute-efficient than larger ones.This demand spans models from 100M [Allal et al., 2025] to trillion parameters.
  • Introduction: 14% higher speculative-decoding throughput is enabled by nesting the draft model within the verifier and sharing their KV cache.The architecture also improves cross-model prediction agreement, increasing token acceptance rates and making larger draft models practical.
  • Introduction: Nested sub-models are progressively stacked into one architecture, reducing a {1B, 8B, 30B, 70B} suite from 109B to 70B trained parameters.Each sub-model can be detached as a standalone language model, while the suite is trained jointly end-to-end.
  • Introduction: 36% less total training compute achieves near-parity with independently trained baselines across validation, out-of-domain perplexity, and standard benchmarks.The 3B suite contains 500M, 1.5B, and 3B sub-models trained on 35B FineWeb-Edu tokens; benchmark performance is within 0.5 average points at every size.
  • Introduction: A low-cost distillation objective uses all sub-model logits available at every forward pass to distill smaller models from the largest.Sub-models of increasing width and depth are nested into a single jointly trained architecture.

2 Method

Matryoshka suites strictly nest sub-model parameters while allowing increasing widths, connect exits through a parameter-free junction, and train all sub-models with cross-entropy plus largest-model distillation. The shared structure also enables free online distillation and efficient speculative decoding with KV-cache reuse between draft and verifier models.

  • Integrated Distillation: The training objective combines each exit’s cross-entropy loss with distillation from the largest sub-model, summing the resulting losses across exits.Online distillation is free because one forward pass through the largest model produces logits for every sub-model; α_d controls teacher-signal strength and should be lower than in offline setups.
  • Architecture: Strictly nested Transformer sub-models use increasing hidden dimensions tailored to each model size, unlike fixed-width early-exit designs.Each sub-model has its own depth, width, and LM head, with D_1 ≤ D_2 ≤ ··· ≤ D_M.
  • Inter-model Junction: A parameter-free junction transfers smaller-model representations into larger widths by retaining existing channels and initializing added channels from fresh embeddings.Norm rescaling prevents magnitude mismatch; omitting rescaling or zeroing fresh embeddings degrades perplexity across all sub-model sizes (Table 3).
  • Speculative Decoding: Any smaller–larger sub-model pair forms a speculative-decoding draft–verifier pair, reusing the draft KV cache because their initial layers are shared.This removes the separate draft model’s memory cost and makes larger drafts practical, including a 500M / 3B pair with a 1:6 size ratio.
  • Architecture: 38% fewer parameters are required at matched exit sizes: the 3B Matryoshka suite totals 3.2B parameters versus 5.2B for Vanilla.Table 1 reports incremental and cumulative parameter totals for both suites.

3 Experiments

Experiments show that Matryoshka suites match or exceed independently trained Vanilla suites across accuracy, perplexity, and compute efficiency, while improving cross-model alignment and speculative decoding throughput. The evaluation also examines architecture choices, training dynamics, and junction and depth configurations.

  • Architecture: The selected (24, 10, 5) depth allocation closely matches Vanilla 3B memory and per-token compute while nesting 500M, 1.5B, and 3B models in a single 39-layer suite.The experiment matches memory and compute footprints across suite types and sweeps feasible depth triplets under a fixed 39-layer budget.
  • Benchmark accuracy and perplexity: 36% less training compute yields near-parity accuracy with Token-matched Vanilla models at every size and +0.4 to +1.9 average-accuracy points over FLOPs-matched baselines.Matryoshka also outperforms Token-matched Vanilla on average out-of-distribution byte perplexity at 1.5B and 3B, and ties it at 500M.
  • Training dynamics and cross-model alignment: Matryoshka maintains Vanilla-level validation perplexity while producing lower pairwise KL divergence and higher token agreement, including +5.7% final agreement for the 1.5B–3B pair.The alignment advantage follows from shared weights and online distillation during training.
  • Speculative decoding: 26% higher speculative-decoding throughput reaches 2,650 tokens/s versus 2,100 for Vanilla at draft length 6 in greedy mode, with gains preserved under nucleus sampling.Matryoshka also achieves 5% higher average accepted length in the nucleus setting, whereas Vanilla barely improves over autoregressive decoding.
  • Inference memory profile: Matryoshka’s lighter inference memory profile comes from depth nesting: KV cache shrinks with sub-model size to 6.0 KB/token, unlike MatFormer’s fixed 31.5 KB/token footprint.The comparison also reports a better quality–size trade-off at matched validation perplexity of approximately 21.
  • Inter-model Junction: The norm-rescaled concatenation junction matches Vanilla within 0.02 average PPL, while removing norm matching creates a +0.2 PPL gap and a zero junction performs poorly at all sizes.These comparisons are reported in the junction ablation at the smaller proxy scale.

4 Related Work

Related work connects Matryoshka language model suites to early exiting, speculative decoding with distillation, elastic compression, and nested Transformer architectures. The approach reduces online distillation costs while complementing methods that compress pretrained models or vary submodel capacity.

  • Early Exiting: Early-exit methods make predictions at intermediate layers, including language-modeling, sparse-routing, and self-speculative-decoding variants [Teerapittayanon et al., 2016; Schuster et al., 2022; Elbayad et al., 2020; Raposo et al., 2024; Elhoushi et al., 2024].These methods avoid a full forward pass when the model is confident.
  • Speculative Decoding and Distillation: Speculative decoding uses a small draft model and larger verifier, while distillation improves their alignment and token acceptance rate [Leviathan et al., 2023; Zhou et al., 2024].The approach avoids offline distillation and reduces the cost of online distillation; offline teacher logits can reduce data requirements by 2.4×.
  • Elastic Compression of Pretrained LLMs: Elastic compression methods convert pretrained LLMs into smaller-model families through continued training, recovering them cheaply after the pretraining cost is already sunk [Cai et al., 2024a; Taghibakhshi et al., 2025].They build on slimmable networks and Once-for-All and are complementary to the paper’s setting.
  • Matryoshka LLMs: Matryoshka Representation Learning and MatFormer nest smaller models by training informative embedding prefixes or varying FFN hidden size across granularities [Kusupati et al., 2022; Devvrit et al., 2023].MatFormer keeps depth, attention heads, and embeddings fixed, yielding FLOPs and parameter savings but identical per-token KV-cache footprints across submodels.

5 Conclusion

The paper introduces Matryoshka Language Model Suites, which nest increasing-size sub-models in one end-to-end-trained architecture and achieve near-parity with independently trained baselines at 3B scale at lower training cost. It also identifies scaling, capacity allocation, loss weighting, and post-training integration as open research directions.

  • Conclusion: Matryoshka suites stack increasing-size sub-models into one nested architecture trained end-to-end, reaching near-parity with independent baselines at 3B scale at lower cost.The framework exploits the shared training distribution of a model suite rather than training each model in isolation.
  • Conclusion: Open questions concern scaling to more sub-models and larger budgets, allocating capacity and loss weight, and integrating instruction tuning, alignment, and reasoning-oriented finetuning.These directions extend joint pretraining from co-designed suites toward broader post-training and scaling settings.

A Distillation Ablation

The distillation ablation on a 200M proxy suite shows that a moderate coefficient, αd = 0.3, best improves sub-model validation performance, while larger coefficients harm it despite increasing teacher-student alignment.

  • A Distillation Ablation: αd = 0.3 closes most of the 50M sub-model’s gap to Vanilla and yields the minimum average gap across the four settings.The sweep evaluates αd ∈ {0, 0.3, 0.5, 0.7} on 50M, 100M, and 200M sub-models.
  • A Distillation Ablation: Higher distillation coefficients are counter-productive: αd = 0.5 makes the largest sub-model trail Vanilla, while αd = 0.7 degrades every size.Without distillation, the 50M and 200M models remain above Vanilla, although the 100M model is already below it.
  • A Distillation Ablation: KL divergence to the 200M teacher decreases monotonically as αd increases, indicating progressively stronger student-teacher alignment.Figure 8(a) reports validation cross-entropy relative to Vanilla, while Figure 8(b) reports each smaller model’s KL divergence to the teacher.

B Compute, Memory, and Training-FLOPs Accounting

This appendix defines the accounting conventions for parameter counts, KV-cache footprint, inference FLOPs, and training FLOPs used throughout the paper. It applies these formulas to standalone Transformers and nested Matryoshka suites, with training cost based on measured throughput and nominal GPU peak performance.

  • Notation: The appendix establishes notation for vocabulary size V, layer count L, hidden dimension H, and per-token dataset size D used in the accounting formulas.H equals the product of the attention-head count and head dimension.
  • Parameter count: Parameter accounting separates embedding and LM-head costs from per-layer attention and FFN costs, then accumulates parameters across the nested Matryoshka sub-models.The standalone expression assigns 2VD to input embeddings and the LM head, while 16LD^2 captures per-layer attention and FFN parameters; nested suites additionally account for smaller-submodel LM heads.
  • KV Cache: KV-cache accounting stores keys and values of width D for each Transformer layer, with memory scaled by b bytes per element.For bf16, b = 2.
  • Theoretical Inference FLOPs: Inference FLOPs count two FLOPs per multiply-accumulate, add the LM-head cost for next-token logits, and sum stack costs across Matryoshka exits before one deepest-exit LM head.The standalone LM-head application costs 2VH FLOPs, while the suite expression uses the stack costs of all sub-models plus 2VDM at the deepest exit.
  • Training FLOPs: Training FLOPs combine wall-clock GPU-hours, logged tokens-per-GPU-hour throughput, and nominal BF16 Tensor-Core peak throughput, summing the runs that compose each suite.The nominal peaks are 4.5 × 10^15 FLOP/s for an NVIDIA B200 and 3.12 × 10^14 FLOP/s for an NVIDIA A100; the convention incorporates kernel differences through measured throughput.

C Choice of Total Layer Budget · D Reading the Ternary Plots

The paper selects the total layer budget by balancing the 3B sub-model’s per-token FLOPs against its KV-cache footprint, then uses ternary plots to represent normalized depth allocations and architectural feasibility. The ternary diagrams also show how reallocating layers changes sub-model width requirements under fixed parameter targets.

  • C Choice of Total Layer Budget: The total budget L directly determines the Matryoshka 3B sub-model’s KV-cache footprint and, with per-block widths, its per-token FLOPs.The 3B sub-model uses every layer across the three blocks, so changing L changes both architectural costs.
  • C Choice of Total Layer Budget: The budget sweep evaluates L ∈{28, 39, 50} using roughly 500 valid depth allocations per budget, with widths solved to match the 500M, 1.5B, and 3B targets.Head dimensions are fixed at 64, 96, and 128, and each allocation’s full-3B forward-pass FLOPs and bf16 KV cache are computed.
  • C Choice of Total Layer Budget: L = 39 offers the key trade-off between the lowest-FLOPs L = 50 designs and the lowest-KV-cache L = 28 designs, while L = 50 costs ∼25−40% more KV cache than Vanilla.L = 28 generally has the lowest KV cache but mostly higher FLOPs than Vanilla, whereas L = 50 reaches the lowest FLOPs at a steep cache cost.
  • C Choice of Total Layer Budget: The width-depth trade-off can affect language-model quality, but prior work finds scale matters more than shape over a reasonable ratio range; Matryoshka performance additionally depends on matching KV-cache footprint.The cited literature is Kaplan et al., 2020, Tay et al., 2023, and Petty et al., 2024.
  • D Reading the Ternary Plots: Ternary plots encode normalized depth triplets (L0/Ltotal, L1/Ltotal, L2/Ltotal) whose components sum to one, so every triangle point is a valid allocation before feasibility constraints.Vertices assign all layers to one sub-model, while edges assign zero layers to one sub-model.
  • D Reading the Ternary Plots: Moving toward the top vertex increases the 500M model’s layer share, producing a deeper, narrower model; moving toward the bottom-left shifts layers toward the 1.5B model and makes the 500M model shallower and wider.These width changes maintain the prescribed 500M parameter target.
  • D Reading the Ternary Plots: The dashed unfeasible region contains depth triplets that cannot realize the prescribed parameter targets with head dimensions (64, 96, 128).Invalidity can arise when a model with too few layers would require fewer than one head, or when nesting D0 ≤D1 ≤D2 is violated.

E Resource Comparison with MatFormer · F Out-of-Distribution Perplexity

Matryoshka provides substantially smaller deployable sub-models and scalable KV-cache costs than MatFormer under a shared 3B shape, while retaining an out-of-domain perplexity advantage on held-out corpora. Its 3B sub-model improves arXiv byte perplexity by 0.13 over the Token-matched Vanilla baseline.

  • E Resource Comparison with MatFormer: MatFormer nests only the FFN block, keeping attention shape, hidden width, and depth fixed at the universal-model values.Its sub-models truncate the FFN intermediate dimension to a fraction r of the universal 4H width.
  • E Resource Comparison with MatFormer: The resource comparison fixes both methods to the Vanilla 3B shape, using H = 2560 and L = 28, and evaluates Matryoshka at its 500M, 1.5B, and 3B exits.MatFormer is evaluated by sweeping FFN ratio r ∈(0, 4], whereas Matryoshka uses cumulative formulas at the three suite exit points.
  • E Resource Comparison with MatFormer: Matryoshka exposes a 500M sub-model, whereas a 3B universal MatFormer cannot expose any sub-500M model in parameter count or FLOPs/tok.MatFormer retains a smallest-model floor of roughly 1.26B parameters at r = 0.5 and approximately 1.05B as r →0; Matryoshka scales down from its smallest cumulative exit point.
  • E Resource Comparison with MatFormer: MatFormer’s KV cache remains 280 KB/tok across extracted sub-models, while Matryoshka reduces KV-cache cost at its smaller exit points.The methods follow essentially the same FLOPs-per-token curve as a function of parameters, but their KV-cache costs diverge because MatFormer fixes attention shape and depth.
  • F Out-of-Distribution Perplexity: The out-of-distribution evaluation measures byte perplexity on WikiText-103, C4, PG-19, arXiv, and PubMed Central using a 2 048-token rolling window, stride 1 024, and the same SmolLM2 tokenizer.The comparisons include FLOPs-matched and Token-matched Vanilla baselines.
  • F Out-of-Distribution Perplexity: 0.13 byte-PPL improvement: the 3B Matryoshka sub-model lowers arXiv byte perplexity by 0.13 versus the Token-matched Vanilla baseline.The 1.5B and 500M Matryoshka sub-models inherit part of this advantage through the shared early layers.

G Closing the Per-Size Gap with Loss Weighting

Loss weighting can reduce the residual per-size perplexity gap to Vanilla: several weighted variants outperform uniform weighting, with the best cutting aggregate excess by roughly 15%. The evidence indicates that uniform weighting over-serves the 100M model while the 200M model remains the bottleneck, though the sweep was limited to the 200M suite.

  • Motivation: The default equal weighting can over-serve some sizes because the 50M, 100M, and 200M models contribute unequally to the gradient.The study tests whether reweighting can close the residual gap to Vanilla using the 200M proxy suite.
  • Experimental scope: The sweep re-runs only the final 1,000 cooldown steps with nine normalized weight vectors, using uniform weighting (1/3, 1/3, 1/3) as the reference.The 3B suite was not tested because its compute cost was prohibitive; systematic per-size weighting rules remain future work.
  • Aggregate result: Roughly 15%: the best loss-weighting configuration cuts the residual positive perplexity excess over Vanilla compared with uniform weighting.The metric averages only positive per-size PPL excesses, so it measures the remaining gap for sub-models still above Vanilla.
  • Per-size effects: Under uniform weighting, 100M is already substantially below Vanilla while 200M is the bottleneck; shifting weight from 100M toward 200M closes the latter gap.The sweep also shifts some budget mildly toward 50M, with per-size effects visualized in Figure 13.

H Additional Size-and-Shape Ablation Views

This section extends the 200M size-and-shape ablation with additional NLL, theoretical-FLOPs, and KV-cache views, which support the same conclusion as the main figure. It also examines how loss-weight configurations affect performance across sub-model sizes.

  • Additional Size-and-Shape Views: Figure 14 adds mean ∆NLL, per-token theoretical-FLOPs, and mean NLL-gap-versus-KV-cache views that support the same conclusion as Figure 7.The views characterize the 200M sweep through alternative quality and efficiency measures.
  • Loss-Weight Ablation: Several loss-weighted variants reduce residual PPL excess relative to uniform weighting, while shifting weight toward the 200M sub-model closes its gap without breaking the 100M sub-model’s advantage.Uniform weighting is (1/3, 1/3, 1/3); the medium and large sub-model weights vary while w50 is residual.
  • Ablation Scope: Each ablation point is a Matryoshka variant with a shared 50M base, evaluated by ∆ values averaged across 50M, 100M, and 200M sub-models against a matching Vanilla 200M baseline.The comparison uses total KV cache per token and mean ∆ppl relative to the Vanilla baseline in the main sweep.

I 200M Proxy Suite Benchmarks

On the 200M proxy suite, Matryoshka matches Vanilla at 50M and outperforms it at 100M and 200M on downstream benchmarks, consistent with perplexity parity. Both suites use the same recipe and tokenizer and are trained for 20B tokens.

  • 200M Proxy Suite Benchmarks: The 200M proxy evaluation uses the same downstream benchmarks as Table 2 and compares the three independently trained Vanilla models with the chosen v9 Matryoshka configuration.Vanilla comprises independently trained 50M, 100M, and 200M models sharing the Matryoshka run’s recipe and tokenizer.
  • 200M Proxy Suite Benchmarks: Matryoshka ties Vanilla at 50M (−0.1 avg) and beats it at 100M and 200M (+0.4 and +0.1 avg) on downstream benchmarks.Table 5 reports per-benchmark zero-shot accuracy across seven benchmarks for the 50M, 100M, and 200M submodels.
  • 200M Proxy Suite Benchmarks: Both Vanilla and Matryoshka suites are trained for 20B tokens, and the downstream results confirm the perplexity parity reported in Figure 14.The comparison’s average accuracy aggregates seven benchmarks, including Winogrande.
Loading 2608.09703v1…