Scalable elastic transformer trains multiple models at once efficiently
ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer
Machine LearningArtificial Intelligence
Summary
Training large AI models for different sizes or speeds usually means doing the work many times over. The authors introduce a way to train many nested models inside one big model without extra cost per training step. Their method organizes the model so all smaller versions are exact parts of the big one and get updated equally during training. They also provide guidelines to keep the training efficient and reuse tuning settings across different model sizes. This approach produces models that perform well compared to training each size individually from scratch.
What this means in practice
- •For machine learning engineers: Train one large transformer model to support multiple inference speeds without duplicating training costs for each size.
- •For cloud ai service operators: Deploy flexible models that adjust to diverse client hardware limits while sharing a single pretrained model instance.
Authors
Can Jin, Yan Kang, Yifan Gong, Chongjian Ge, Mingze Xu, Yuanjun Xiong, Dimitris N. Metaxas, Hongwu Peng
Abstract
Foundation models need to support diverse inference budgets, yet training a separate model for each budget duplicates pretraining effort. Elastic models train several nested sub-networks inside one weight-shared supernet, but existing recipes face three obstacles: variable per-step workloads that hinder distributed scaling, limited hyperparameter (HP) transfer across model sizes, and no guarantee that all sub-networks are maximally updated at once. We present \Ours, an elastic transformer that organizes hidden dimensions into a dense core and successive shells, so that every hidden matrix is block lower triangular. One full-width forward pass contains every sub-network's exact hidden states as prefixes, so joint training keeps a fixed computation graph and costs about as much per step as training the widest sub-network alone. We further introduce \muE, an extension of maximal update parameterization (\mup) to \Ours and other elastic architectures. Its scaling rules provably keep all sub-networks of a supernet maximally updated simultaneously (\muI), and keep doing so when the supernet is widened or given more shells; sweeps confirm that HPs tuned on a small non-elastic proxy stay near-optimal for elastic families and larger supernets (\mbox{\muT}). A single \Ours run produces sub-networks that outperform MatFormer and fully elastic baselines at matched parameter counts, and approach separately trained models that together take several times the steps.