Matryoshka Language Model Suites
TL;DR
This paper improves both training and inference efficiency by stacking sub-models of increasing size into a single nested architecture trained end-to-end in a Matryoshka training framework that reduces the total parameter count, enables low-cost distillation from the largest to all smaller sub-models at every training step, and is well-suited for speculative decoding.
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.