# One-stop Training of Multiple Capacity Models

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 37 of its 40 references that are in the map.
- First posted on arXiv 23 May 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Unsupervised Cross-lingual Representation Learning at Scale](https://mlsanity.com/paper/1911.02116/) (2019)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [mT5: A massively multilingual pre-trained text-to-text transformer](https://mlsanity.com/paper/2010.11934/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Improved Knowledge Distillation via Teacher Assistant](https://mlsanity.com/paper/1902.03393/) (2019)
- [Beyond English-Centric Multilingual Machine Translation](https://mlsanity.com/paper/2010.11125/) (2020)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [Q8BERT: Quantized 8Bit BERT](https://mlsanity.com/paper/1910.06188/) (2019)
- [Multilingual Machine Translation with Large Language Models: Empirical Results and Analysis](https://mlsanity.com/paper/2304.04675/) (2023)
- [On the Representation Collapse of Sparse Mixture of Experts](https://mlsanity.com/paper/2204.09179/) (2022)
- [Distilling Task-Specific Knowledge from BERT into Simple Neural Networks](https://mlsanity.com/paper/1903.12136/) (2019)
- [Taming Sparsely Activated Transformer with Stochastic Experts](https://mlsanity.com/paper/2110.04260/) (2021)
- [StableMoE: Stable Routing Strategy for Mixture of Experts](https://mlsanity.com/paper/2204.08396/) (2022)
- [MixKD: Towards Efficient Distillation of Large-scale Language Models](https://mlsanity.com/paper/2011.00593/) (2020)
- [Learning Student-Friendly Teacher Networks for Knowledge Distillation](https://mlsanity.com/paper/2102.07650/) (2021)
- [SpeechMoE: Scaling to Large Acoustic Models with Dynamic Routing Mixture of Experts](https://mlsanity.com/paper/2105.03036/) (2021)
- [Learning Language Specific Sub-network for Multilingual Machine Translation](https://mlsanity.com/paper/2105.09259/) (2021)
- [Parameter-Efficient and Student-Friendly Knowledge Distillation](https://mlsanity.com/paper/2205.15308/) (2022)

## Abstract (arXiv)

Training models with varying capacities can be advantageous for deploying them in different scenarios. While high-capacity models offer better performance, low-capacity models require fewer computing resources for training and inference. In this work, we propose a novel one-stop training framework to jointly train high-capacity and low-capactiy models. This framework consists of two composite model architectures and a joint training algorithm called Two-Stage Joint-Training (TSJT). Unlike knowledge distillation, where multiple capacity models are trained from scratch separately, our approach integrates supervisions from different capacity models simultaneously, leading to faster and more efficient convergence. Extensive experiments on the multilingual machine translation benchmark WMT10 show that our method outperforms low-capacity baseline models and achieves comparable or better performance on high-capacity models. Notably, the analysis demonstrates that our method significantly influences the initial training process, leading to more efficient convergence and superior solutions.