About ShellElastic
Foundation models need to support diverse inference budgets. Elastic models train several nested sub-networks inside one weight-shared supernet. ShellElastic is 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 the exact hidden states of every sub-network, so joint training keeps a fixed computation graph and costs about as much per step as training the widest sub-network alone.
µ-Elastic extends maximal update parameterization (µP) to ShellElastic and other elastic architectures. Its scaling rules keep all sub-networks of a supernet maximally updated simultaneously, and hyperparameters tuned on a small non-elastic proxy stay near-optimal for elastic families and larger supernets. A single ShellElastic 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.
Checkpoints
Final checkpoints of ShellElastic and four baselines at the three scales of the paper, with sub-networks of about
119M–285M, 375M–906M and 675M–1.69B parameters (Section 4.1). Each supernet repository holds all five sub-networks
(Low, Medium, High, Extra, Max) and its reference implementation; each Standalone and Sampled standalone repository
holds the five models in the folders low/ to max/. The repositories are named
Can111/ShellElastic-<scale> and
Can111/ShellElastic-baseline-<method>-<scale>.
| Method | Checkpoints (sub-network sizes) |
|---|---|
| ShellElastic Supernet with its five sub-networks (Low, Medium, High, Extra, Max) | 119M–285M 375M–906M 675M–1.69B |
| MatFormer Baseline trained for the paper: a supernet that nests only the FFN hidden width | 119M–285M 375M–906M 675M–1.69B |
| Fully elastic Baseline trained for the paper: a supernet that slices shared dense matrices | 119M–285M 375M–906M 675M–1.69B |
| Standalone Baseline: five separately trained models, each on the full budget | 119M–285M 375M–906M 675M–1.69B |
| Sampled standalone Baseline: the five Standalone models trained within the steps of one elastic run | 119M–285M 375M–906M 675M–1.69B |
Data
Can111/ShellElastic-data: the tokenized training and validation data that these runs read, in the layout that the code reads. Training: 3,136 shards of DCLM-Baseline (141.32 GB, 25,690,112 sequences of 2,049 tokens); validation: 11 shards of held-out C4 text (1.12 GB, 84,999 sequences).