Elastic pretraining trains several nested sub-networks of different sizes inside one weight-shared supernet, so that a single run serves several inference budgets. ShellElastic makes this scalable: its hidden dimensions form a dense core followed by shells, so one full-width forward pass computes every sub-network exactly and each step costs about as much as training the widest sub-network alone. μ-Elastic, an extension of μP to ShellElastic and other elastic models, keeps all sub-networks maximally updated at once and transfers hyperparameters tuned on a small non-elastic proxy to elastic models and to larger supernets. At three scales from 119M to 1.69B parameters, one ShellElastic run produces sub-networks that outperform MatFormer and fully elastic baselines at matched parameter counts and approach separately trained models that together take five times the steps.

ShellElastic (Figure 1) partitions the residual stream into a dense core followed by ordered shells, each a multiple of a base width g, and each supported sub-network is the core plus a prefix of the shells. Every hidden matrix is block lower triangular: a shell reads the core, the earlier shells and itself, but never a later shell. Trunk normalization acts within each group, attention heads align with the group boundaries, and every supported width has its own readout. One full-width forward pass therefore contains the exact hidden states of every sub-network as prefixes (Proposition 1), and the per-step compute stays close to that of a single full-width pass however many sub-networks are trained, so every step has a fixed workload, as in non-elastic training.

μ-Elastic (Figure 2) extends maximal update parameterization (μP) (Yang et al., 2021) to elastic models with two factors: a μ-Transfer factor m = g/d₀ that carries HPs from a non-elastic proxy of width d₀ to base width g, and a μ-Internal factor for the different widths inside one shared model. Every group reads the same input coordinates in every sub-network that contains it, so we set its initialization, learning rate and weight decay from this fan-in, while each readout uses its own input width (Table 1). This keeps every group maximally updated in all sub-networks at once (Theorem 1), which uniform scaling across groups cannot do in general (Proposition 2). The same two factors apply to MatFormer and to fully elastic models (Appendix D). The standard parameterization (SP) sets both factors and the FFN expansion e to one in the scaling rules and fixes the readout scalars α_s = 1, preserving the architecture.

Recipe.

1. Tune (σ₀, η₀, λ₀) on a non-elastic proxy of width d₀.
2. Choose the base width g and the supported set 𝒮.
3. Compute m = g/d₀ and instantiate the parameterization in Table 1.
4. Train the supernet on the joint objective (Equation 8).

Table 2. Final validation loss and mean 5-shot accuracy over 13 downstream tasks at three scales, labeled by sub-network size. Entries at the 119M–285M and 375M–906M scales average three runs; those at 675M–1.69B are single runs. †Five separately trained models, together taking 5× the steps of any other row. Bold marks the best of the four methods that train all five sizes in one run.

Among the methods that train all five sizes in one run, ShellElastic has the lowest mean validation loss (0.028, 0.005 and 0.018 nats below the best baseline) and the highest mean 5-shot accuracy at every scale, and is best or tied-best at 12 of the 15 tiers in validation loss and 14 in 5-shot accuracy. With a fifth of their optimizer steps, ShellElastic comes within 0.039, 0.045 and 0.014 nats of the standalone models in mean validation loss and within 0.8 points in mean 5-shot accuracy.

_{Scroll sideways: Figure 3, then Figures 12–16 of Appendix H (each sub-network, and weight decay).}

Under μ-Elastic, the proxy’s optimum, η₀ = 2 × 10⁻³ and σ₀ = 3.3 × 10⁻³ (the base values of Section 4.1), transfers to all three elastic models at both levels of μ-Transfer (Figure 3, Corollary 1) and stays within 0.005 nats of the optimum of every sub-network (Figures 12–14). Under SP, the optimum departs from the proxy’s and shifts with scale, as Section 3.2 anticipates, and at the base values SP is 0.03–0.11 nats worse. Weight decay transfers as well (Figures 15 and 16): under μ-Elastic, the mean loss of ShellElastic at 119M–285M is flat over the same range as the proxy’s, from λ₀ = 10⁻³ to 10⁻¹ (within 0.003 nats), whereas SP is higher at every swept value.

Table 3. Effect of μ-Elastic on each elastic model at the 119M–285M and 375M–906M scales. For each model, the first row is the best of the 12 (13 at 375M–906M) SP configurations of Figure 3 at that scale (lowest mean validation loss), and the w/ μ-Elastic row is its result in Table 2, with the base values of Section 4.1 and no re-tuning. Table 28 gives the values with four decimals.

Without width-dependent scaling, re-tuning one global value only trades one width against another (Proposition 2). Table 3 gives SP the best of its configurations in Figure 3 and μ-Elastic only the base values; still, at the 119M–285M scale μ-Elastic gives lower training and validation losses at every tier for the fully elastic model and ShellElastic (by 0.011 and 0.013 nats in mean validation loss), as it keeps all widths maximally updated at once (Theorem 1), and at the 375M–906M scale tuned SP remains 0.006–0.021 nats worse for all three models. Under tuned SP, ShellElastic still beats both baselines by 0.026–0.027 and 0.020–0.024 nats at the two scales, so its lead is not due to μ-Elastic alone.

All ablations use the 119M–285M supernet of Section 4.1 and change one setting at a time; the default results are those of Table 2. Of the settings we ablate, sub-network selection, in-place distillation, ensembles of the sub-networks, stage-residual gradient scaling and the trunk gradient of the narrower widths can change the mean validation loss by 0.010 to 0.20 nats, whereas the loss weights, normalization, readouts, base width and head dimension, and optimization HPs change it by at most 0.007 nats (Section 4.4, Table 4; Appendix J, Tables 31 and 32).

Where the paper averages several runs, the tables and plots on the project page show its averages, and the curves, data files and logs come from one of the runs it averages for that entry.

Reference implementation of ShellElastic and the μ-Elastic parameterization from the paper, together with the baselines of the paper (fully elastic model, MatFormer, standalone and sampled standalone models, StairFormer) and everything needed to rerun its experiments. Section, appendix, equation, table and figure numbers in this README and in the code refer to the paper, arXiv:2610.15124.

Python ≥ 3.10 and PyTorch ≥ 2.6 with CUDA; run all commands from the repository root.

pip install -e ".[all,dev]"

The tokenized DCLM-Baseline training shards and C4 validation shards that the paper's runs read are on Hugging Face (Can111/ShellElastic-data):

python scripts/download_data.py --scale 1 --out data    # 119M–285M (20 GB)
python scripts/download_data.py --scale 2 --out data    # 375M–906M (72 GB)
python scripts/download_data.py --scale 3 --out data    # 675M–1.69B (141 GB), covers all three scales
python scripts/download_eval_data.py --out eval_data    # the 13 evaluation tasks

torchrun --nproc_per_node=8 -m elastic.train configs/paper/scale1/shellelastic.yaml
torchrun --nproc_per_node=8 -m elastic.train configs/paper/scale1/shellelastic.yaml optim.lr=1e-3 run.name=my_run

Any setting can be overridden as key=value. A run writes config.yaml, metrics.jsonl, checkpoints and the final weights (final/; final_<i>/ per model for sampled standalone) to outputs/<run.name>/ and resumes from the latest checkpoint there, so a new run needs a new run.name. The metrics of each sub-network are named by its width, e.g. val_loss_width_512 (by its FFN width for MatFormer). The paper configs also evaluate during training and need eval_data/; eval.probe=null train.final_eval=false turns this off.

The paper configs run on one node of 8 GPUs. On another number of GPUs, keep the global batch, data.batch_size × #GPUs, at 128 (119M–285M) or 256 (the two larger scales), and likewise data.eval_batch_size; the data order then differs from the paper's runs.

The paper configs end training with the 5-shot evaluation of every sub-network on the 13 tasks (Appendix F). To evaluate exported weights:

torchrun --nproc_per_node=8 -m elastic.eval outputs/<run.name>/final [widths=[4,8]] [out=results.json]

widths selects sub-networks by their entry of 𝒮 (default: every trained one); the results are icl/<task>_5shot_width_<n> and their mean icl/mean_width_<n>.

Configs are YAML files; inherit: builds on other files (the paper configs inherit configs/paper/common.yaml), and key=value arguments override any setting.

pytest tests

@article{jin2026shellelastic,
  title   = {{ShellElastic}: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer},
  author  = {Jin, Can and Kang, Yan and Gong, Yifan and Ge, Chongjian and Xu, Mingze and Xiong, Yuanjun and Metaxas, Dimitris N. and Peng, Hongwu},
  journal = {arXiv preprint arXiv:2610.15124},
  year    = {2026},
  url     = {https://arxiv.org/abs/2610.15124}
}

Apache-2.0. See LICENSE and THIRD_PARTY_NOTICES.md (xFormers, OpenLM, LLM Foundry, Composer). DCLM-Baseline is released under CC-BY-4.0, C4 under ODC-BY; the evaluation datasets keep their own licenses.