Research map / Architectures and efficiency / RNNs and normalization
MatFormer: Nested Transformer for Elastic Inference
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MatFormer, a Transformer with nested FFN widths trained by random granularity sampling so one model yields many accurate submodels (via Mix'n'Match) for elastic inference across decoders (MatLM) and ViT encoders (MatViT).
Problem
Foundation models are typically released as a few independently trained sizes, forcing practitioners to pick a model that does not match their latency or cost budget; training many sizes is expensive.
Approach
Nest FFN hidden neurons Matryoshka-style so smaller blocks are prefixes of larger ones, train g=4 granularities by sampling one per step, and extract many intermediate submodels post hoc by mixing granularities across layers with a low-slope increasing heuristic (Mix'n'Match).
Key results
- MatLM-850M submodels (582M to 850M) achieve better validation loss and one-shot downstream scores than independently trained baselines of the same size.
- MatViT-L/16 reduces compute by 40% with under 0.5% 1-NN retrieval accuracy drop for adaptive query encoding.
- Speculative decoding with MatLM S draft and XL verifier gives 1.14x (LAMBADA) and 1.11x (TriviaQA) speedups vs 1.10x and 1.08x for baselines.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MatLM instead of independently trained baseline models for speculative decoding | 850M draft and verifier models · LAMBADA | speculative decoding speedup: $1.10\times$ -> $1.14\times$ | Helped |
| Use MatLM instead of independently trained baseline models for speculative decoding | 850M draft and verifier models · TriviaQA | speculative decoding speedup: $1.08\times$ -> $1.11\times$ | Helped |
| Add shared attention cache (ablation) | MatLM speculative decoding · LAMBADA | speculative decoding speedup: $1.14\times$ -> $1.16\times$ | Helped |
| Add shared attention cache (ablation) | MatLM speculative decoding · TriviaQA | speculative decoding speedup: $1.11\times$ -> $1.14\times$ | Helped |
| Use MatFormer for fine-tuning while retaining ViT pretraining (ablation) | ViT-L/16 pretrained on ImageNet-21K · ImageNet-1K, 306M parameters | top-1 accuracy (%): 85.26 -> 85.57 | Helped |
| Use MatFormer for fine-tuning while retaining ViT pretraining (ablation) | ViT-L/16 pretrained on ImageNet-21K · ImageNet-1K, 206M parameters | top-1 accuracy (%): 85.12 -> 84.27 | Hurt |
| Use MatFormer for fine-tuning while retaining ViT pretraining (ablation) | ViT-L/16 pretrained on ImageNet-21K · ImageNet-1K, 156M parameters | top-1 accuracy (%): 85.02 -> 82.79 | Hurt |
| Use MatFormer for fine-tuning while retaining ViT pretraining (ablation) | ViT-L/16 pretrained on ImageNet-21K · ImageNet-1K, 131M parameters | top-1 accuracy (%): 84.42 -> 82.1 | Hurt |
| Use MatFormer for both pretraining and fine-tuning (ablation) | ViT-L/16 · ImageNet-1K, 306M parameters | top-1 accuracy (%): 85.26 -> 85.61 | Helped |
| Use MatFormer for both pretraining and fine-tuning (ablation) | ViT-L/16 · ImageNet-1K, 206M parameters | top-1 accuracy (%): 85.40 | Helped |
| Use MatFormer for both pretraining and fine-tuning (ablation) | ViT-L/16 · ImageNet-1K, 156M parameters | top-1 accuracy (%): 85.02 | Helped |
| Use MatFormer for both pretraining and fine-tuning (ablation) | ViT-L/16 · ImageNet-1K, 131M parameters | top-1 accuracy (%): 84.41 | Helped |
Surprising findings
- DynaBERT-style joint optimization of all granularities underperforms MatFormer's per-step sampling by about 0.01 log perplexity at 850M, even with 15% more compute.
- Mix'n'Match configurations match or beat NAS-predicted configurations at equal parameter budgets.
Limitations
- Granularities are fixed to g=4 exponentially spaced FFN ratios; sampling and mix-and-match rely on heuristics without theoretical guarantees.
- Scaling laws do not capture multi-submodel optimization effects, and NAS at scale is costly.
- Larger-scale results are limited to 850M LMs and ViT-B/L; the paper notes that submodels are less consistent at very large scales.
Experimental setup
Datasets: ImageNet-1K ImageNet-21K 25 English downstream tasks (TriviaQA, Natural Questions, LAMBADA, HellaSwag, etc.)
Benchmarks: ImageNet-1K 25 English LM tasks (1-shot) LAMBADA TriviaQA
Metrics: validation loss (log perplexity) top-1 accuracy 1-nearest neighbor accuracy speedup consistency (token match / KL divergence)
Baselines: Independently trained Baseline (S/M/L/XL) OFA DynaBERT Sorted Network
Builds on: Matryoshka Representation Learning Once-for-All DynaBERT Vision Transformer speculative decoding
Closest papers by meaning
- MoDeGPT: Modular Decomposition for Large Language Model Compression (ICLR 2025 · similarity 0.87)
- Tandem Transformers for Inference Efficient LLMs (ICML 2024 · similarity 0.87)
- Trainable Transformer in Transformer (ICML 2024 · similarity 0.87)
- Block Transformer: Global-to-Local Language Modeling for Fast Inference (NeurIPS 2024 · similarity 0.86)
- Building on Efficient Foundations: Effective Training of LLMs with Structured Feedforward Layers (NeurIPS 2024 · similarity 0.86)
- Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.