模型架構與推論
Matryoshka Embeds Three Models in a Single Set of Weights, Cutting Training Compute by 36%
Matryoshka Language Model Suites links 500M, 1.5B, and 3B submodels in a nested architecture that allows each model to be extracted and run independently, avoiding duplicate training for every size. Sharing early layers and the KV cache also increases speculative decoding throughput by 14% to 26%, though the approach has so far been validated only on base models up to 3B parameters.
Researchers at Cornell University have introduced Matryoshka Language Model Suites, which turn language models of different sizes into a single set of progressively expanding, layer-by-layer weights rather than three unrelated checkpoints. The experimental architecture first runs a 24-layer, 500M-parameter submodel, then appends 10 wider layers to form a 1.5B model, and finally adds five layers to create a 3B model. Width transitions are connected by parameter-free junctions that first align the norm of the preceding submodel’s output with the new embeddings and then concatenate the additional channels.
This design requires just 3.2B parameters for the complete suite, 38% fewer than the combined 5.2B parameters of three independent models. Because a single forward pass simultaneously produces logits at every exit, the largest model can also distill directly into the smaller models at each step, without separately storing teacher outputs or rerunning the teacher. The team trained the suite from scratch on 35 billion tokens from FineWeb-Edu. Across seven zero-shot multiple-choice benchmarks, the average score at each size differed by no more than 0.5 points from independently trained models using the same number of tokens, while total training FLOPs fell by 36%. The 1.5B and 3B exits also achieved slightly lower out-of-domain perplexity.
For inference, the key advantage is not simply loading fewer weights, but allowing the draft model and verifier to share their early layers and corresponding KV cache. With conventional independent models, a 500M/3B pairing slows generation because the draft model is too large. The nested version, however, delivers roughly a 20% to 40% speedup over its own standard decoding, with overall throughput gains ranging from 14% to 26% across evaluation settings. The authors have released Transformers-compatible weights and checkpoints from multiple training stages, but the model is an unaligned English base model trained on FineWeb-Edu and requires `trust_remote_code`. It also remains unproven whether the savings extend to MoE models with tens or hundreds of billions of parameters, long-context workloads, and quantized deployments. Engineering teams should next watch whether vLLM and SGLang add native implementations of shared-layer speculative decoding, and whether multi-tenant serving can realize meaningful memory and latency gains in practice.