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 only
  • 1-5B_* → the 500M + 1.5B sub-models
  • 3B_* → the full 500M + 1.5B + 3B suite
  • main, 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