README.md

metadata

license: cc-by-4.0
language:
  - en
datasets:
  - mlfoundations/dclm-baseline-1.0
  - Can111/ShellElastic-data
library_name: pytorch
pipeline_tag: text-generation
tags:
  - elastic-models
  - foundation-models
  - hyperparameter-transfer
  - maximal-update-parameterization
  - pretraining

Scalable Elastic Pretraining
with Simultaneous Maximal Updates and Hyperparameter Transfer

Can Jin^1* · Yan Kang² · Yifan Gong² · Chongjian Ge² · Mingze Xu² · Yuanjun Xiong² · Dimitris N. Metaxas¹ · Hongwu Peng²

¹Rutgers University   ²Adobe 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.<j> (j = 1, …, 4) is W_(j+1), the block row of the shell G_(j+1), of shape [|G_(j+1)|, s_(j+1)·g] (for the FFN, e times as many rows in gate_proj and up_proj and e times as many columns in down_proj). With i = 0, …, 31 numbering the layers and <tier> one of low, medium, high, extra, max:

Tensor	Shape
embed_tokens.weight	[50432, 1024]; sub-network s uses the first s·g columns
layers.<i>.{attn_norm,mlp_norm}.weight, layers.<i>.attn.{q,k}_norm.weight	[1024]; sub-network s uses the first s·g entries
layers.<i>.attn.{q,k,v,gate,o}_proj.core.weight	[512, 512]
layers.<i>.attn.{q,k,v,gate,o}_proj.shell.<j>.weight	[128, 640], [128, 768], [128, 896], [128, 1024] for j = 1, …, 4
layers.<i>.mlp.{gate,up}_proj.core.weight	[512, 512]
layers.<i>.mlp.{gate,up}_proj.shell.<j>.weight	[128, 640], [128, 768], [128, 896], [128, 1024]
layers.<i>.mlp.down_proj.core.weight	[512, 512]
layers.<i>.mlp.down_proj.shell.<j>.weight	[128, 640], [128, 768], [128, 896], [128, 1024]
norm.<tier>.weight	[s·g]: the terminal norm of sub-network s
lm_head.<tier>.weight	[50432, s·g]: the readout U_s multiplied by its learned logit scale α_s

mlp.down_proj includes the forward multiplier 1/e of Table 1 (here e = 1), so the model needs no further scalars; query, key and value are separate tensors.

Training

Data. DCLM-Baseline (Li et al., 2024), a filtered Common Crawl corpus, with the GPT-NeoX-20B tokenizer (Black et al., 2022) and a sequence length of 2,048 (Section 4.1). A run at this scale takes 25,000 steps of 128 sequences, 6.55B tokens (Section 4.1). The tokenized training and validation data of these runs are in the dataset repository Can111/ShellElastic-data.

Recipe. Trained as described in the paper (Sections 3.3 and 4.1, Appendices E.1 and E.2, and Table 9).

Hyperparameters (Section 4.1, Appendix E.1 and Table 9):

Setting	Value
Parameterization	µ-Elastic (Section 3.2, Table 1) with d0 = 128 and m = g/d0 = 1
Optimizer	AdamW with decoupled weight decay
Base learning rate η0	2×10⁻³
Base initialization σ0	3.3×10⁻³, truncated normal at ±3 standard deviations
Base weight decay λ0	0.01; none on norm gains and α_s
β1, β2	0.95, 0.95
AdamW ε	10⁻¹⁴
Learning-rate schedule	warmup-stable-decay
Learning-rate warmup	linear, first 1,000 steps
Learning-rate decay	linear to 0, from step 20,000 to step 25,000
Batch size	128 sequences (262,144 tokens)
Sequence length	2,048
Training steps	25,000 (6.55B tokens)
Gradient clipping	global norm 4
Step skipping	after step 5,000, a step is skipped if its gradient norm exceeds 20× the clip value
z-loss coefficient λz	10⁻⁴ (Equation 8); all reported losses are the cross-entropy alone (Appendix E.2)
Loss weights ω_s, Low to Max	(1, 1, 1, 1, 4) (Equation 8)
Precision	bf16 mixed precision with fp32 master weights and fp32 gradient reduction

Usage

The repository is self-contained: the reference implementation needs only PyTorch, plus safetensors to read the weights (pip install -r requirements.txt), and runs on CPU and GPU. Download the repository (for example with huggingface_hub.snapshot_download("Can111/ShellElastic-119M-285M", local_dir="ShellElastic-119M-285M")) and run, from that directory:

import torch
from modeling_shellelastic import ShellElasticForCausalLM

model = ShellElasticForCausalLM.from_pretrained(".")  # float32 on the CPU
# model = ShellElasticForCausalLM.from_pretrained(".", device="cuda", dtype=torch.bfloat16)
ids = torch.tensor([[510, 5347, 273, 6181, 310]])     # GPT-NeoX-20B ids of "The capital of France is"
logits = model(ids, tier="medium")                    # [1, 5, 50432]; tiers: low, medium, high, extra, max
per_tier = model(ids, tiers=model.config.tiers)       # all five sub-networks from one trunk pass
logits, loss = model(ids[:, :-1], tier="low", labels=ids[:, 1:])   # mean next-token cross-entropy
out = model.generate(ids, max_new_tokens=20, tier="max")           # greedy decoding (no key/value cache)
low = model.extract("low")                            # the Low sub-network as a standalone model
low.save_pretrained("low_model")                      # its config.json and model.safetensors

extract(tier) keeps the embedding columns, the core and the shells of that sub-network, the leading entries of every norm gain, and its own terminal norm and readout (Section 3.3); by Proposition 1 its logits equal the supernet's for that sub-network, and from_pretrained loads the saved folder. dense_state_dict(tier) returns the weights of a sub-network with dense matrices (zeros above the block diagonal), for inspection and porting; running them still requires the group-local norms. python example.py --tier medium runs a sub-network on the CPU and checks its extraction (--save DIR writes the extracted model).

Inputs are token ids of the GPT-NeoX-20B tokenizer, for example from tokenizers.Tokenizer.from_pretrained("EleutherAI/gpt-neox-20b"); the tokenizer is not bundled here. The vocabulary has 50,432 entries (Table 7), of which the tokenizer uses the first 50,277.

Validation

Loaded in bfloat16 (the paper uses bf16 mixed precision, Table 9), the weights and code in this repository reproduce the final validation loss in the this checkpoint column of Results for every sub-network to within 2×10⁻⁴ nats.

Files

• README.md: this card
• LICENSE: the Creative Commons Attribution 4.0 International license (CC BY 4.0)
• requirements.txt: torch and safetensors
• modeling_shellelastic.py: the reference implementation (PyTorch only)
• example.py: loads the supernet on the CPU and runs one sub-network and extracts it
• config.json: the architecture of the supernet
• model.safetensors: the weights of the supernet (float32)

Limitations

These are research checkpoints of language models pretrained on DCLM-Baseline, a filtered Common Crawl corpus, for a study of elastic pretraining. They have not been instruction-tuned, aligned or evaluated for safety, are not intended for deployment, and can produce incorrect, biased or offensive text. Use them for research only.

License

Everything in this repository (the weights, the configurations and the code) is released under CC-BY-4.0 (LICENSE).

Third-party terms: the training data, DCLM-Baseline, are distributed under CC-BY-4.0 and derived from Common Crawl; the validation text is from C4, distributed under ODC-BY together with the Common Crawl terms of use. Neither is included in this repository; the tokenized data of these runs are in the dataset repository Can111/ShellElastic-data. The GPT-NeoX-20B tokenizer (Apache-2.0) is not bundled here.

Citation

@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}
}