# Matryoshka Model Learning for Improved Elastic Student Models

arXiv: https://arxiv.org/abs/2505.23337 · HTML: https://mlsanity.com/paper/2505.23337/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 57 references that are in the map.
- First posted on arXiv 29 May 2025.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Deep Mutual Learning](https://mlsanity.com/paper/1706.00384/) (2017)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation](https://mlsanity.com/paper/1905.08094/) (2019)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Large scale distributed neural network training through online distillation](https://mlsanity.com/paper/1804.03235/) (2018)
- [Soft Weight-Sharing for Neural Network Compression](https://mlsanity.com/paper/1702.04008/) (2017)
- [Less is More: Task-aware Layer-wise Distillation for Language Model Compression](https://mlsanity.com/paper/2210.01351/) (2022)
- [MatFormer: Nested Transformer for Elastic Inference](https://mlsanity.com/paper/2310.07707/) (2023)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [End-to-End Speech Translation with Knowledge Distillation](https://mlsanity.com/paper/1904.08075/) (2019)
- [Does Graph Distillation See Like Vision Dataset Counterpart?](https://mlsanity.com/paper/2310.09192/) (2023)
- [Three Pillars improving Vision Foundation Model Distillation for Lidar](https://mlsanity.com/paper/2310.17504/) (2023)
- [One Teacher is Enough? Pre-trained Language Model Distillation from Multiple Teachers](https://mlsanity.com/paper/2106.01023/) (2021)
- [On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models](https://mlsanity.com/paper/2209.05310/) (2022)
- [Incorporating User Micro-behaviors and Item Knowledge into Multi-task Learning for Session-based Recommendation](https://mlsanity.com/paper/2006.06922/) (2020)
- [Train Once, Deploy Anywhere: Matryoshka Representation Learning for Multimodal Recommendation](https://mlsanity.com/paper/2409.16627/) (2024)

## Built on by (most cited first, 2 of 2)

- [Unlocking Multimodal Document Intelligence: From Current Triumphs to Future Frontiers of Visual Document Retrieval](https://mlsanity.com/paper/2602.19961/) (2026)
- [Objective-Specific Privileged Bases via Full-Prefix Matryoshka Learning](https://mlsanity.com/paper/2605.09160/) (2026)

## Abstract (arXiv)

Industry-grade ML models are carefully designed to meet rapidly evolving serving constraints, which requires significant resources for model development. In this paper, we propose MatTA, a framework for training multiple accurate Student models using a novel Teacher-TA-Student recipe. TA models are larger versions of the Student models with higher capacity, and thus allow Student models to better relate to the Teacher model and also bring in more domain-specific expertise. Furthermore, multiple accurate Student models can be extracted from the TA model. Therefore, despite only one training run, our methodology provides multiple servable options to trade off accuracy for lower serving cost. We demonstrate the proposed method, MatTA, on proprietary datasets and models. Its practical efficacy is underscored by live A/B tests within a production ML system, demonstrating 20% improvement on a key metric. We also demonstrate our method on GPT-2 Medium, a public model, and achieve relative improvements of over 24% on SAT Math and over 10% on the LAMBADA benchmark.