Scalable Elastic Pretraining
with Simultaneous Maximal Updates and Hyperparameter Transfer
Can Jin1* · Yan Kang2 · Yifan Gong2 · Chongjian Ge2 · Mingze Xu2 · Yuanjun Xiong2 · Dimitris N. Metaxas1 · Hongwu Peng2
1Rutgers University 2Adobe Research
*Corresponding author: can.jin@rutgers.edu. Work done during an internship at Adobe Research.
## 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.
## ShellElastic supernet, 119M–285M scale
Final weights of the ShellElastic supernet at the 119M–285M scale of the paper (Section 4.1). `model.safetensors` holds the supernet with its five sub-networks (Low, Medium, High, Extra, Max); `modeling_shellelastic.py` runs any of them and extracts it as a standalone model (Section 3.3).
## Sub-networks
| Sub-network | Model width | Attention heads | FFN width | Parameters |
|---|---|---|---|---|
| Low | 512 | 4 | 512 | 118,817,285 (118.8M) |
| Medium | 640 | 5 | 640 | 152,715,909 (152.7M) |
| High | 768 | 6 | 768 | 190,808,837 (190.8M) |
| Extra | 896 | 7 | 896 | 233,096,069 (233.1M) |
| Max | 1,024 | 8 | 1,024 | 279,577,605 (279.6M) |
Parameter counts include the five learned readout scalars α_s; rounded to 0.1M, they equal the entries of Table 10 of the paper. `model.num_parameters(tier)` counts the extracted sub-network, whose readout weights already contain α_s, and reports 5 fewer. `model.safetensors` stores the supernet: 421,596,928 parameters (1.69 GB in float32).
## Results
| Sub-network | Val. loss, this checkpoint | Val. loss, paper | 5-shot acc. (%), this checkpoint | 5-shot acc. (%), paper |
|---|---|---|---|---|
| Low | 3.4079 | 3.409 | 35.6 | 35.5 |
| Medium | 3.3377 | 3.338 | 36.4 | 36.6 |
| High | 3.2938 | 3.294 | 36.8 | 36.6 |
| Extra | 3.2632 | 3.264 | 37.1 | 36.9 |
| Max | 3.2395 | 3.240 | 37.3 | 36.9 |
| Avg. | 3.3084 | 3.309 | 36.6 | 36.5 |
*This checkpoint*: the final validation loss of each sub-network of these weights (per-token cross-entropy in nats on held-out C4 text, at step 25,000) and its mean 5-shot accuracy over the 13 downstream tasks of Appendix F. *Paper*: the main results of the paper for ShellElastic at the 119M–285M scale (Table 2). *Avg.*: the mean over the five sub-networks. Appendix G of the paper gives the per-task accuracies.
## Architecture
ShellElastic (Section 3.1 of the paper) organizes the hidden dimensions into a dense core followed by successive shells (Figure 1). This supernet has the settings of Section 4.1 and Appendix E.1 (Tables 7 and 8):
| | |
|---|---|
| Base width, supported set | g = 128, S = {4, 5, 6, 7, 8}, M = 8: a 512-dimensional core followed by four 128-dimensional shells |
| Groups | G_1 (the core) to G_5, of widths (512, 128, 128, 128, 128); sub-network s_k contains the first k groups, so Low to Max have widths 512 to 1,024 |
| Projections | every hidden matrix is block lower triangular: the block row W_k, which produces output group G_k, reads only the first s_k·g input dimensions (Equation 2) |
| Normalization | the two branch norms of every layer are group-local over G_1, …, G_5 (Equation 3); every sub-network has its own terminal norm over all of its s·g dimensions (Equation 6) |
| Attention | d_head = 128, one head per base width; sub-network s uses the first s·g/d_head heads (Section 3.1) |
| FFN | SwiGLU with expansion e = 1 (Equation 5) |
| Readout | a readout U_s and a learned logit scale α_s per sub-network (Equation 6) |
| Shared settings | 32 pre-norm layers with residual-branch multiplier 1; RMSNorm with learnable gains and ε_N = 10⁻⁵; query–key RMSNorm per head; rotary embeddings with base 10,000 (the code uses the rotate-half form); attention logit scale d_head^−1/2 and an elementwise sigmoid gate on the attention output, sigmoid(x W_G^⊤) (Equation 4); SwiGLU FFN; no biases; untied embeddings over the 50,432-token GPT-NeoX vocabulary; sequence length 2,048 (Table 8 and Section 4.1) |
`model.safetensors` stores every hidden matrix by block row, without the zero blocks above the diagonal: `core` is W_1, the block row of the core G_1, and `shell.
