README.md

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

<div align="center">
<div align="center" class="dark:hidden"><img alt="ShellElastic" src="https://huggingface.co/spaces/Can111/ShellElastic/resolve/main/shellelastic-logo-light.svg" height="42"></div>
<div align="center" class="hidden dark:block"><img alt="ShellElastic" src="https://huggingface.co/spaces/Can111/ShellElastic/resolve/main/shellelastic-logo-dark.svg" height="42"></div>
<h3 align="center">Scalable Elastic Pretraining<br>with Simultaneous Maximal Updates and Hyperparameter Transfer</h3>
<p align="center"><small><b>Can&nbsp;Jin</b><sup>1*</sup> · <b>Yan&nbsp;Kang</b><sup>2</sup> · <b>Yifan&nbsp;Gong</b><sup>2</sup> · <b>Chongjian&nbsp;Ge</b><sup>2</sup> · <b>Mingze&nbsp;Xu</b><sup>2</sup> · <b>Yuanjun&nbsp;Xiong</b><sup>2</sup> · <b>Dimitris&nbsp;N.&nbsp;Metaxas</b><sup>1</sup> · <b>Hongwu&nbsp;Peng</b><sup>2</sup></small></p>
<p align="center"><small><sup>1</sup>Rutgers University &emsp; <sup>2</sup>Adobe Research</small></p>
<p align="center"><small><sup>*</sup>Corresponding author: <a href="mailto:can.jin@rutgers.edu">can.jin@rutgers.edu</a>. Work done during an internship at Adobe Research.</small></p>
<p align="center">
<a href="https://arxiv.org/pdf/2610.15124"><img alt="Paper" src="https://img.shields.io/badge/Paper-PDF-b31b1b?style=for-the-badge"></a>
<a href="https://jincan333.github.io/ShellElastic_page"><img alt="Project Page" src="https://img.shields.io/badge/Project-Page-2f6db5?style=for-the-badge"></a>
<a href="https://github.com/jincan333/ShellElastic"><img alt="Code" src="https://img.shields.io/badge/Code-GitHub-181717?style=for-the-badge&logo=github&logoColor=white"></a>
<a href="https://huggingface.co/spaces/Can111/ShellElastic"><img alt="Checkpoints" src="https://img.shields.io/badge/Checkpoints-Hugging%20Face-ffd21e?style=for-the-badge&logo=huggingface&logoColor=white"></a>
<a href="https://huggingface.co/spaces/Can111/ShellElastic"><img alt="Datasets" src="https://img.shields.io/badge/Datasets-Hugging%20Face-ffd21e?style=for-the-badge&logo=huggingface&logoColor=white"></a>
</p>
</div>

## 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](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0) (Li et al., 2024), a filtered Common Crawl corpus, with the [GPT-NeoX-20B tokenizer](https://huggingface.co/EleutherAI/gpt-neox-20b) (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](https://huggingface.co/datasets/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:

```python
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](https://huggingface.co/EleutherAI/gpt-neox-20b), 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](https://creativecommons.org/licenses/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](https://huggingface.co/datasets/Can111/ShellElastic-data). The GPT-NeoX-20B tokenizer (Apache-2.0) is not bundled here.

## Citation

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