Matryoshka Language Model Suites

The paper introduces Matryoshka Language Model Suites, a pretraining framework that nests sub-models of increasing size into a single architecture trained end-to-end, rather than training each model in a suite separately. The authors state: We improve both training and inference efficiency by stacking sub-models of increasing size into a single nested architecture trained end-to-end.

The key insight is that the natural structure of a model suite, a set of models sharing the same training distribution, can be exploited by nesting at lower cost than independent training. Instead of training separate models, the approach progressively stacks the outputs of smaller sub-models as inputs to larger ones. For example, rather than training a 1B, 8B, 30B, 70B suite as four independent runs, the method trains a single architecture combining a 1B sub-model with additional 7B, 22B, and 40B parameter blocks, resulting in 70B trained parameters instead of 109B.

Architecture: The framework enforces strict parameter nesting: θ1 ⊂ θ2 ⊂ ⋯ ⊂ θ M = θ. Each sub-model T θm is a separate Transformer stack with its own hidden dimension D m, number of layers n m, and a dedicated LM head W m. The authors require strictly increasing hidden dimensions D1 ≤ D2 ≤ ⋯ ≤ D M.

Inter-model Junction: To pass outputs between sub-models of different widths, the authors design a parameter-free junction mechanism. The output of sub-model m is norm-rescaled to match the norm of the input embedding of sub-model m+1 before concatenation: õ m θ(x) = o m θ(x) · (‖e m+1 θ(x)‖2 / ‖o m θ(x)‖2). This addresses the magnitude mismatch where Transformer outputs typically have much larger norms than input embeddings.

Integrated Distillation: The framework makes distillation from the largest model to all smaller ones a natural by-product, as logits from all submodels are available during training. The distillation loss is defined as the cross-entropy between teacher and student next-token distributions: L M→m d = −Σ v stop grad(σ(l M t) v) log σ(l m t) v. The total loss for each sub-model is a convex combination: L m = (1 − α d) L m ce + α d L M→m d, with α d = 0.3 for main runs.

Speculative Decoding: The nested structure benefits speculative decoding because the draft model is nested in the verifier model, and both already share KV cache, eliminating the usual memory overhead of maintaining a separate draft model. Additionally, cross-model prediction agreement is also higher than in independently trained suites, directly increasing token acceptance rates.

The authors train a 3B-parameter Matryoshka suite with 500M, 1.5B, and 3B sub-models on 35B tokens from FineWeb-Edu. The architecture uses depth triplet (n1, n2, n3) = (24, 10, 5) with a total of 39 layers, chosen to match the KV cache and per-token FLOPs of the Vanilla 3B baseline. The suite totals 3.2B parameters versus 5.2B for the Vanilla suite (−38%).

Training Compute: The Matryoshka suite uses 36% less training compute than independently trained baselines while achieving near-parity on benchmark performance.

Benchmark Performance: "The Matryoshka suite reaches near-parity with the Token-matched Vanilla baseline on average accuracy at every size (within 0.5 points) while using 36% less training compute, and outperforms the FLOPs-matched Vanilla baseline at every size by +0.4 to +1.9 points."

Out-of-Distribution Perplexity: On average, OOD byte perplexity, the Matryoshka suite outperforms the Token-matched Vanilla baseline at the 1.5B and 3B scales (and ties at 500M), with the largest single-corpus advantage on arXiv.

Speculative Decoding Throughput: At a draft length of 6, Matryoshka reaches 2,650 tokens/s in greedy mode against 2,100 for Vanilla (a 26% speedup), and the gain is preserved under nucleus sampling. The authors note that the Vanilla 500M–3B pair barely beats plain autoregressive inference and is actually slower under nucleus sampling, whereas the Matryoshka 500M / 3B pair turns the same configuration into a 20–40% speedup over its own standard decoding.

Cross-Model Agreement: Matryoshka pairs exhibit noticeably lower KL divergence and higher token-level agreement than independently trained Vanilla pairs of matched sizes, with the effect strongest for the (1.5B, 3B) pair (+5.7% final agreement).

The authors compare against MatFormer and find: Matryoshka achieves a better size–performance frontier. MatFormer's shared attention backbone fixes its KV cache at 31.5 KB/token across all sub-models, whereas Matryoshka's KV cache scales with the sub-model. Additionally, at matched validation perplexity (∼21), Matryoshka-100M matches MatFormer-M (139M) with roughly half the KV cache and fewer parameters.

Distillation: Removing distillation causes a +0.15 PPL average gap more strongly noticeable on the smaller models.

Width/Depth Configuration: The authors sweep depth allocations and find several variants are comparable to the vanilla suite both in terms of total KV cache growth and of average performance gap.

Distillation Coefficient: α d = 0.3 minimizes the average gap; higher values harm the larger sub-models.

The paper's contributions are summarized as:

1. Introducing Matryoshka Language Model Suites, a pretraining framework that nests sub-models of increasing width and depth into a single jointly trained architecture

2. Deriving a low-cost distillation objective that leverages the availability of all sub-model logits at every forward pass

3. Training a 3B-parameter Matryoshka suite showing near-parity with independently trained baselines while using 36% less training compute

4. Demonstrating that the shared-weight structure improves cross-model agreement and makes larger drafts practical for speculative decoding

5. Providing extensive ablations at smaller scale offering insights for architectural hyperparameters and training objectives

The authors conclude: "At the 3B scale, our Matryoshka suite reaches near-parity with independently trained baselines on standard benchmarks while using 36% less training compute, allowing low-cost online distillation and enabling speculative decoding with higher throughput and no memory overhead for the draft model. They note that our ablations produce a solid basis for scaling up Matryoshka suites to larger sizes."

Improvements to AI systems:

1. Nested model architectures for efficient multi-size deployment – Instead of training separate models for different compute budgets (e.g., edge vs. cloud), AI systems can use a single nested architecture where smaller sub-models are extracted from larger ones. This reduces total training compute by 36% while maintaining near-parity performance, enabling one system to serve users with varying latency/accuracy needs without redundant training.

2. Parameter-free inter-model junctions for width scaling – The norm-rescaling and concatenation mechanism allows seamless information flow between sub-models of different hidden dimensions without learned parameters. This can be applied to any transformer-based system (e.g., vision, multimodal) to build elastic models that adapt their width dynamically during inference, improving efficiency on heterogeneous hardware.

3. Integrated online distillation for continual improvement – Since all sub-model logits are available during every forward pass, the system can perform low-cost distillation from the largest model to smaller ones without extra forward passes. This enables AI systems to continuously improve smaller deployed models using the largest model's knowledge, with a distillation coefficient (α d=0.3) that balances performance across sizes.

4. Speculative decoding with shared KV cache – The nested structure eliminates the memory overhead of separate draft models in speculative decoding. AI systems can use the smallest sub-model as a draft for the largest, achieving 20–40% throughput gains over standard decoding, with higher token acceptance rates due to improved cross-model agreement. This is directly applicable to production LLM serving systems to reduce latency.

5. Compute-aware model suite design – The framework provides a method to design model suites where total parameter count and KV cache scale sub-linearly with the largest model size (e.g., 3.2B vs. 5.2B parameters for a 3B suite). AI systems can use this to pack more capability into fixed memory budgets, particularly for on-device or edge deployments where multiple model sizes are needed.

6. Improved out-of-distribution robustness – The nested training with distillation yields better OOD perplexity on domains like arXiv compared to independently trained models. This suggests AI systems can use this approach to improve generalization to unseen data distributions, particularly for scientific or technical text processing.

7. Flexible depth-width trade-offs – The ablations show that multiple depth/width configurations achieve comparable performance, giving system designers freedom to optimize for specific hardware constraints (e.g., KV cache limits) without sacrificing accuracy. This enables AI systems to be tailored to specific deployment scenarios (e.g., long-context tasks) while maintaining a single training run.

What the improved AI system can do:

• Deploy a single trained model that serves 500M, 1.5B, and 3B inference requests with near-parity to individually trained models, using 36% less training compute.

• Run speculative decoding with a 500M draft and 3B verifier, achieving 2,650 tokens/s (26% faster than vanilla pairs) without extra memory for a separate draft model.

• Continuously distill knowledge from the largest sub-model to smaller ones during training, improving small-model accuracy without additional forward passes.

• Adapt to different hardware (e.g., mobile, laptop, server) by extracting the appropriate sub-model on-the-fly, with minimal performance loss.

• Maintain better performance on out-of-distribution text (e.g., scientific papers) compared to independently trained models of the same size.

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.