Nested and Elastic Transformers

A product team needs a 2B model for the phone, a 4B for the laptop, a 9B for the cheap cloud tier and a 12B for the expensive one. The standard answer is four pretraining runs, four data mixes, four sets of evals, four sets of bugs. The Llama 3 family consumed 39.3M H100 hours across its three sizes (Meta, 2024, The Llama 3 Herd of Models, arXiv:2407.21783), and nothing about that number says the 8B and the 70B had to be separate projects.

Elastic architectures take the other route. Train one parent model whose sub-networks are nested inside it by construction, then slice a child out of the parent's weights at deployment time with no fine-tuning. The child is not an approximation of the parent; it is a region of the parent that was explicitly optimised during training.

Nesting the FFN

MatFormer is the cleanest version of the idea. Each transformer block's feed-forward network is defined with a nested granularity structure: the first \(m/8\) hidden units form the smallest sub-block, the first \(m/4\) contain those, and so on up to the full \(m\). Training samples across granularities, so every sub-block is optimised as a functioning FFN in its own right rather than as a truncation (Devvrit, Kudugunta et al., 2023, MatFormer: Nested Transformer for Elastic Inference, arXiv:2310.07707).

The payoff is combinatorial. If a 32-layer model has four trained granularities per layer, you do not get four models, you get \(4^{32}\) possible configurations, because layer \(7\) can run at the smallest granularity while layer \(8\) runs at the largest. The paper calls picking a good one Mix'n'Match, and it is a heuristic search with no training cost. A 2.6B MatLM yields extracted models from 1.5B to 2.6B whose validation loss and one-shot evaluations are comparable to independently trained models of the same size, and at the 850M scale the extracted sub-models beat their independently trained counterparts.

Two things make the nested structure work that are easy to miss. The sub-blocks must be prefixes, not arbitrary subsets, or there is no single set of weights that serves every budget. And the training loss has to include the small granularities from the start; bolting elasticity onto a finished model is a different and harder problem, which is what Flextron solves by converting a trained LLM into an elastic one using 7.63% of the original pretraining tokens (Cai et al., 2024, Flextron: Many-in-One Flexible Large Language Model, arXiv:2406.10260).

Elastic versus compressed

Pruning-and-distillation pipelines also produce a family from one big model, and they produce very good ones: Minitron derives 8B and 4B models from a pretrained 15B using up to 40x fewer training tokens per model than training from scratch, for a 1.8x compute saving across the family and up to 16% better MMLU than the from-scratch 4B (Muralidharan et al., 2024, Compact Language Models via Pruning and Knowledge Distillation, arXiv:2407.14679). The distinction is what you hold at the end. Compression gives you N independent checkpoints, each with its own weights, its own storage and its own serving process. Elasticity gives you one checkpoint whose children are views into it.

That difference is what makes multi-budget serving cheap. Nemotron Elastic embeds 6B and 9B models inside a 12B parent using 110B training tokens, reports roughly 360x lower cost than training the family from scratch and about 7x lower than iterative compression, and lets a server switch budgets zero-shot because every size shares the same resident weights (Nemotron Elastic, 2025, arXiv:2511.16664). By 2026 the same machinery was being applied post-training to hybrid Mamba-Transformer-MoE reasoning models, with different sub-models used for the thinking and answering phases of a single request (Star Elastic, 2026, arXiv:2605.07182).

When it breaks

Memory does not shrink the way compute does. Running the 2B child of a 12B parent still requires the parent's weights to be loadable, because the child's weights are a slice of them. Elasticity buys FLOPs and latency, not footprint, unless you materialise and ship the slice as its own artefact, at which point you are back to N checkpoints for storage purposes and have only saved on training.

Granularities interfere. Gradient signal from the small sub-blocks and the full block lands on the same parameters, and the small ones act as a regulariser on the large one. That is sometimes a gift (the 850M MatLM sub-models beating their baselines) and sometimes a tax at the top end. Once-for-All hit the same wall in the CNN era and answered it with progressive shrinking, training large first and admitting smaller configurations gradually (Cai et al., 2019, Once-for-All, arXiv:1908.09791).

The evaluation surface explodes. Hundreds of extractable models means hundreds of models you have not evaluated. Mix'n'Match finds a configuration under a parameter budget; it does not tell you that configuration is safe, unbiased or correctly calibrated. In practice teams certify a handful of named configurations and treat the rest as unsupported, which quietly returns the fixed ladder the architecture was supposed to replace.

Slicing the FFN leaves the KV cache alone. FFN granularity controls a large share of the parameters and the per-token FLOPs, and none of the attention state. For long-context or high-concurrency serving, where the KV cache dominates memory, an elastic FFN changes very little; see cross-layer KV sharing for the lever that does move that number.