Matryoshka Language Model Suites
Nathan Godey, Yoav Artzi
cs.AI, cs.CL
2026-08-10
A nested 500M/1.5B/3B Matryoshka suite matches independently trained models within 0.5 accuracy points, uses 36% less compute, and speeds 500M-draft speculative decoding by 14–26%.
Language-model releases are suites, not single checkpoints: 500M, 1.5B, 3B, and so on, each aimed at a different memory and latency budget. The usual recipe trains every size from scratch and then distills the large model into the small ones. The training bill is close to the sum of the sizes. Speculative decoding then adds a separate draft model with its own KV cache.
Godey and Artzi train the suite as one nested net: smaller parameters are a strict subset of larger ones, and one forward pass emits logits at every size. The goal is to keep quality and perplexity intact while cutting suite-level training compute, and to grow the draft model inside the verifier.
Three sub-models have increasing width, independent depth, and nested weights θ1⊂θ2⊂θ3. Unlike early-exit stacks, each exit has its own hidden size and LM head, so it detaches as a normal checkpoint. The main run is 500M / 1.5B / 3B with depth split (24, 10, 5) and 39 layers total, chosen to match a standalone 3B on KV cache and per-token FLOPs. The nested suite totals 3.2B parameters; three independent models sum to 5.2B (−38%).
Widths differ, so a small exit cannot feed the next block as-is. The junction adds no extra parameters: it rescales the m-th output to the L2 norm of a fresh embedding, concatenates along channels, and sends the result into the next Transformer stack. Dropping the norm match costs about +0.2 average PPL in the 200M ablation; replacing the fresh embedding with zeros costs +0.54.
Distillation is almost free. The largest model's softmax is the teacher; every smaller exit gets a token-level cross-entropy with αd=0.3, mixed with its own LM loss. Offline distillation would store logits or run a teacher in parallel. Here every forward pass already has all of them.
For speculative decoding, the 500M draft and 3B verifier share the first 24 layers of KV. The verifier reuses the draft cache and runs only the extra layers.
Training uses 35B FineWeb-Edu tokens, length 2048, AdamW, with αd chosen on a 200M proxy sweep. Independent Llama-style baselines share data and hyperparameters, plus a compute-matched 23B-token cooldown.
At matched tokens (35B), zero-shot average accuracy on seven benchmarks stays within 0.5 points at every size: 46.5 vs 47.0 at 500M, 52.2 vs 52.4 at 1.5B, 53.3 vs 53.5 at 3B. At matched FLOPs (independent suite stopped at 23B tokens), Matryoshka leads by 0.4 to 1.9 points. Out-of-domain byte PPL beats the token-matched baseline at 1.5B and 3B (2.121 vs 2.139, 2.067 vs 2.097) and ties at 500M. Validation PPL tracks the independent curves within ±1.4%.
Alignment comes from shared weights plus online distillation. Pairwise KL is lower through training, and top-1 next-token agreement ends 5.7% higher for the 1.5B/3B pair. On an A100, 500M draft and 3B verifier at draft length 6: greedy throughput is 2,650 tokens/s versus 2,100 for the independent pair (+26%). Under nucleus sampling the throughput gain is about 14%, with accepted length 5% higher. A vanilla 1:6 draft is slower than plain autoregression under nucleus; the nested pair at the same ratio speeds up. Shared cache also lifts max batch from 64 to 102.
At 200M versus MatFormer, MatFormer nests only FFN width and keeps one KV footprint for every granularity. Matryoshka nests depth, so KV shrinks with the sub-model. At validation PPL around 21, Matryoshka-100M matches MatFormer-M (139M) with fewer parameters and less KV.
This is an engineering paper for teams that ship a size ladder. A 3B budget almost gives 500M and 1.5B as by-products, plus a better-aligned draft. Speculative decoding is the sharpest use: the draft is a prefix of the verifier, KV is shared, and a 1:6 ratio that usually loses money starts to pay.
It is not a continuous width dial. Exit sizes are fixed at training time, and the detachable checkpoints are a few discrete rungs, not MatFormer's combinatorial granularities. The trade is that each rung has its own depth and KV and can be served like a normal model.
The main run stops at 3B and 35B tokens, an order of magnitude below current 70B suites. Per-sub-model losses are summed with equal weight; a 200M sweep shows reweighting closes much of the leftover gap, and the 3B suite was not reswept. There is no instruction tuning, alignment, or long-reasoning post-training, so it is unknown whether the nest survives SFT. Speculative decoding reports only the 500M/3B pair. The norm-matching junction is a stability patch; only two negative controls exist. Data is FineWeb-Edu. Code and multilingual generalization are untested.