Matryoshka-3B — a nested 500M / 1.5B / 3B LM suite
A Matryoshka language-model suite: three sub-models of increasing size (500M ⊂ 1.5B ⊂ 3B) stacked into a single nested architecture and trained end-to-end in one run. Running the full model produces all three sub-models' outputs in a single forward pass, while activating fewer parameters and writing less KV cache than three independent models.
This reduces the total parameter count of the suite (3.20B vs. 5.20B for an equivalent set of independent models, −38%), enables free online distillation from the largest sub-model to the smaller ones at every step, and is well-suited to speculative decoding — the draft model is contained within the verifier and shares its early layers and KV cache.
⚠️ Custom architecture — load with
trust_remote_code=True.
Suite architecture
All sub-models share the SmolLM2 tokenizer (V = 49,152) and use RoPE
(θ = 1e5). Sub-model m+1 consumes sub-model m's output through a
parameter-free norm-rescaled junction: o^m is rescaled to match the norm of
a fresh embedding covering the new channels, then concatenated.
| Sub-model | Incr. params | Cumul. params | Width D | Layers (cumul.) | Heads | Head dim | Intermediate |
|---|---|---|---|---|---|---|---|
| 500M | 0.50B | 0.50B | 1024 | 24 (24) | 16 | 64 | 4096 |
| 1.5B | 0.98B | 1.48B | 2304 | 10 (34) | 24 | 96 | 9216 |
| 3B | 1.72B | 3.20B | 4352 | 5 (39) | 34 | 128 | 17408 |
Depth triplet (n₁,n₂,n₃) = (24,10,5) (39 layers total) was chosen to match the
KV-cache and per-token FLOPs of a Vanilla 3B baseline. The 500M sub-model uses the
same width/depth as the Vanilla 500M for a strictly comparable data point.
Checkpoints & revisions
Every checkpoint is a separate git branch, selected via revision=. Branch
names follow:
{exit}_{tokens}B[_cd|_distill] e.g. 3B_35B, 1-5B_42B, 3B_60B_cd
{exit}— how many nested exits the checkpoint exposes:500M_*→ the 500M sub-model only1-5B_*→ the 500M + 1.5B sub-models3B_*→ the full 500M + 1.5B + 3B suitemain,main_{tokens}B→ pointers to the full suite (main= default/latest)
{tokens}B— training tokens seen:5, 10, 16, 21, 26, 31, 35, 42, 52, 60._cd— WSD cooldown applied (learning-rate decay), released at 60B._distill— distillation-ablation variant (small exits only). Will be discussed in a next version of the paper.
The paper's main results use the 35B-token checkpoints; the repo additionally
ships a longer 60B run with cooldown (*_60B_cd) and intermediate checkpoints
for studying training dynamics.
Because the suite is nested, 500M_35B, 1-5B_35B and 3B_35B are slices of
the same trained weights — load the largest exit you need.
Usage
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
"nthngdy/matryoshka-3B",
revision="3B_35B", # full suite; use "1-5B_35B" or "500M_35B" for smaller
trust_remote_code=True,
dtype=torch.bfloat16,
attn_implementation="sdpa",
).eval()
tokenizer = AutoTokenizer.from_pretrained("nthngdy/matryoshka-3B", revision="3B_35B")
# The nested sub-models are exposed as a dict, ordered small → large:
for tag, submodel in model.lm_model_dict.items():
n = sum(p.numel() for p in submodel.parameters())
print(f"{tag}: {n/1e6:.0f}M params")
A single forward pass through the full suite yields logits for all contained exits. For per-exit generation and for speculative decoding with a small exit drafting for a larger one (sharing early layers + KV cache).
Training
| Data | FineWeb-Edu, sequences packed to length 2048 |
| Tokens | 35B (main); extended 60B run also released |
| Optimizer | AdamW (β₁=0.9, β₂=0.95, ε=1e-8) |
| Peak LR | 4e-4, WSD schedule (3,000 cooldown steps / 33,000 total) |
| Batch size | 512 sequences |
| Weight decay | 0.01 · Grad clip |
| Precision | bf16-mixed |
| Distillation | online, from the 3B exit to smaller ones, α_d = 0.3 |
| Hardware | NVIDIA B200 (~52 GPU-days across both suites) |
Results (from the paper)
- On par with independently-trained baselines on validation PPL, out-of-domain PPL, and benchmark accuracy, while using 36% less training compute.
- 14–26% higher throughput via speculative decoding. A 500M/3B pair (1:6 ratio) that degrades latency for independent models becomes a 20–40% speedup here, thanks to shared KV cache and early layers.
Intended use & limitations
Research artifact for studying nested LM suites, elastic-size deployment, and draft-in-verifier speculative decoding. Base (non-instruct) models trained on English FineWeb-Edu (~35–60B tokens); not aligned or safety-filtered. Outputs may be low-quality, biased, or factually wrong. Not intended for production use.
Citation
TBD