[2608.09703] Matryoshka Language Model Suites
Abstract: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.
View PDF HTML (experimental) Abstract: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…
saved by
related reading
- Fast Inference from Transformers via Speculative Decodingarxiv.org
- Fast Inference from Transformers via Speculative Decodingarxiv.org
- A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
- Matryoshka Sparse Autoencoders — LessWronglesswrong.com
- Towards Monosemanticity: Decomposing Language Models With Dictionary Learningtransformer-circuits.pub
- [2205.13147] Matryoshka Representation Learningarxiv.org
- [2509.14786] Pre-training under infinite computearxiv.org
- [2005.14165] Language Models are Few-Shot Learnersarxiv.org
- Towards Monosemanticity: Decomposing Language Models With Dictionary Learningtransformer-circuits.pub
- Mosaic LLMs: GPT-3 quality formosaicml.com
- Esoteric Language Modelsarxiv.org
- [2106.09685] LoRA: Low-Rank Adaptation of Large Language Modelsarxiv.org