Matryoshka Language Model Suites

2026-08-10

AI Summary

This work addresses the inefficiency of conventional language models, which require separate training and deployment of submodels at different scales, leading to suboptimal training and inference efficiency. To overcome this limitation, the authors propose a nested Matryoshka training framework that unifies multi-scale submodels into a single architecture and enables end-to-end joint training. This design naturally integrates cross-scale knowledge distillation and speculative decoding within a cohesive pipeline. The proposed approach achieves comparable baseline performance and perplexity while substantially reducing total parameter count and training overhead—cutting training compute requirements by 36%—and simultaneously improving speculative decoding throughput by 14–26%.

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.

Problem

• language model suite
• training efficiency
• inference efficiency
• model distillation
• speculative decoding

Innovation

• Matryoshka Language Models
• Nested Architecture
• End-to-End Training
• Speculative Decoding
• Knowledge Distillation