Matryoshka Language Model Suites

2026-08-10Artificial Intelligence

Artificial IntelligenceComputation and Language
AI summary

The authors propose a new way to train language models by nesting smaller models inside bigger ones, like Russian Matryoshka dolls, so they can be trained together efficiently. This method reduces the total number of parameters needed and allows smaller models to learn from the largest one during training. Their approach matches the performance of separately trained models while using less computing power and speeding up certain types of text generation. They also studied which design choices work best for this nested model setup.

language modelnested architectureMatryoshka trainingparameter countdistillationspeculative decodingbenchmark performanceperplexitythroughputablation study
Authors
Nathan Godey, Yoav Artzi
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.