ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer

Can Jin; Yan Kang; Yifan Gong; Chongjian Ge; Mingze Xu; Yuanjun Xiong; Dimitris N. Metaxas; Hongwu Peng

cs.LGOct 10, 2026

ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer

Authors: Can Jin, Yan Kang, Yifan Gong, Chongjian Ge, Mingze Xu, Yuanjun Xiong, Dimitris N. Metaxas, Hongwu Peng

Organizations: Rutgers University · Adobe Research

Abstract

Foundation models need to support diverse inference budgets, yet training a separate model for each budget duplicates pretraining effort. Elastic models train several nested sub-networks inside one weight-shared supernet, but existing recipes face three obstacles: variable per-step workloads that hinder distributed scaling, limited hyperparameter (HP) transfer across model sizes, and no guarantee that all sub-networks are maximally updated at once. We present \Ours, 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 every sub-network's exact hidden states as prefixes, so joint training keeps a fixed computation graph and costs about as much per step as training the widest sub-network alone. We further introduce \muE, an extension of maximal update parameterization (\mup) to \Ours and other elastic architectures. Its scaling rules provably keep all sub-networks of a supernet maximally updated simultaneously (\muI), and keep doing so when the supernet is widened or given more shells; sweeps confirm that HPs tuned on a small non-elastic proxy stay near-optimal for elastic families and larger supernets (\mbox{\muT}). A single \Ours 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.

Figures & tables

Figure 1: Core and shells of ShellElastic preserve supported prefixes. (a) Each colored block row produces output group

\mathcal{G}_{k}

from the input prefix ending at

s_{k}g

. Group

\mathcal{G}_{1}

is the dense core; subsequent groups are shells. (b)

N_{k}

normalizes group

\mathcal{G}_{k}

using its own statistics.

Parameter	Active input width	Initialization	Learning rate	Weight decay	Multiplier
Token embedding $E$	$V$ (fixed)	$\sigma_{0}$	$\eta_{0}$	$\lambda_{0}$	$1$
Attention; FFN up/gate	$sg$	$\sigma_{0}/\sqrt{ms}$	$\eta_{0}/(ms)$	$\lambda_{0}ms$	$1$
FFN down	$esg$	$\sigma_{0}\sqrt{e}/\sqrt{ms}$	$\eta_{0}/(ms)$	$\lambda_{0}ms$	$1/e$
Readout $U_{s}$	$sg$	$\sigma_{0}/\sqrt{ms}$	$\eta_{0}/(ms)$	$\lambda_{0}ms$	$\alpha_{s}$
Norm gains; readout scalar $\alpha_{s}$	$1$ (elementwise)	$1$	$\eta_{0}$	$0$	$1$

Table 1:

\mu

-Elastic for ShellElastic . Scaling rules set by the

\mu

-Transfer factor

m

, the

\mu

-Internal factor

s

(Equation 7 ), and the FFN expansion

e=D_{\mathrm{ffn}}/D_{\mathrm{model}}

;

s

is the output group’s supported boundary for hidden weights and the sub-network width index for readouts. Initialization gives nominal weight standard deviations, or initial values for gains and scalars.

Figure 2: Overview of

\mu

-Elastic . The ratio

q=m\cdot s

of Equation 7 scales the HPs of every group and readout. (a)

\mu

-Transfer: the

(\sigma_{0},\eta_{0},\lambda_{0})

tuned on the proxy are reused (i) for an elastic supernet and (ii) for a wider one with more shells (Corollary 1 ). (b)

\mu

-Internal: block row

W_{k}

has fan-in

s_{k}g

in every sub-network that contains it, so

q=ms_{k}

keeps it maximally updated in all of them (Theorem 1 ). The bars show each group’s pre-activation size at initialization and update size: with the proxy’s

(\sigma_{0},\eta_{0})

, as under SP, all groups, even the core, exceed the proxy by

\sqrt{q}

and

q

, so no shared value fixes them all (Proposition 2 ); under

\mu

-Elastic , all match it.

		Val. loss $\downarrow$						5-shot acc. (%) $\uparrow$
	Method	Low	Medium	High	Extra	Max	Avg.	Low	Medium	High	Extra	Max	Avg.
$119$ M– $285$ M	Standalone †	3.369	3.308	3.264	3.219	3.191	3.270	35.6	35.8	36.6	37.1	37.1	36.5
	Sampled standalone	3.667	3.609	3.583	3.532	3.524	3.583	33.8	33.4	34.0	34.2	34.4	34.0
	MatFormer	3.425	3.357	3.325	3.312	3.314	3.347	35.1	35.2	35.8	35.9	35.9	35.6
	Fully elastic	3.429	3.354	3.315	3.295	3.290	3.336	34.9	35.5	35.7	35.9	36.0	35.6
	ShellElastic	3.409	3.338	3.294	3.264	3.240	3.309	35.5	36.6	36.6	36.9	36.9	36.5

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

119

M–

285

M and

375

M–

906

M scales average three runs; those at

675

M–

1.69

B are single runs. † Five separately trained models, together taking

5\times

the steps of any other row. Bold marks the best of the four methods that train all five sizes in one run.

		Train loss $\downarrow$						Val. loss $\downarrow$
	Method	Low	Medium	High	Extra	Max	Avg.	Low	Medium	High	Extra	Max	Avg.
$119$ M– $285$ M	MatFormer	3.261	3.184	3.157	3.140	3.139	3.176	3.434	3.360	3.327	3.314	3.311	3.349
	w/ $\mu$ -Elastic	3.254	3.182	3.156	3.141	3.144	3.175	3.425	3.357	3.325	3.312	3.314	3.347
	Fully elastic	3.272	3.192	3.154	3.130	3.126	3.175	3.445	3.368	3.325	3.303	3.297	3.348
	w/ $\mu$ -Elastic	3.257	3.178	3.145	3.123	3.120	3.165	3.429	3.354	3.315	3.295	3.290	3.336
	ShellElastic	3.247	3.180	3.136	3.104	3.079	3.149	3.423	3.354	3.308	3.275	3.249	3.322

Table 3: Effect of

\mu

-Elastic on each elastic model at the

119

M–

285

M and

375

M–

906

M scales. For each model, the first row is the best of the

12

(

13

375

M–

906

M) SP configurations of Figure 3 at that scale (lowest mean validation loss), and the w/

\mu

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

Selected sub-networks per step $\mathcal{A}_{\tau}$
	Val. loss $\downarrow$						5-shot acc. $\uparrow$
Setting	Low	Medium	High	Extra	Max	Avg.	Avg.
Sandwich: Low, Max and one other	3.4019	3.3617	3.3183	3.2816	3.2395	3.3206	36.1
Random: one of the five sub-networks	3.4311	3.3786	3.3542	3.3410	3.3356	3.3681	34.8
Max + random: Max and one other	3.4944	3.4079	3.3447	3.2956	3.2397	3.3564	35.5
All: all five sub-networks (default)	3.4086	3.3375	3.2944	3.2641	3.2403	3.3090	36.5

Table 4: Ablations at the

119

M–

285

M scale: final validation loss and mean 5-shot accuracy over

13

downstream tasks and the five tiers. Each block ends with the default, and bold marks the best value of each column within a block. Appendix J gives all ablations.

Appendix figures & tables33 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: The complete ShellElastic model. Embedding, one attention–FFN layer repeated

L

times, and the terminal norm and readout of the selected width. Group-local branch normalization and structured projections preserve every supported prefix through the trunk.

Figure 7: Prefix-preserving attention and FFN. Sub-network

s

uses the first

sg/d_{\mathrm{head}}

attention heads. The FFN’s up, gate, and down projections expand and contract the same groups.

Figure 8: Complete architecture comparison. MatFormer varies the FFN hidden dimension; fully elastic models vary the model dimension and the FFN hidden dimension together by slicing dense matrices; ShellElastic also varies both, but through block-lower-triangular projections in which each output group reads only a prefix of the input. Each layer frame repeats

L

times. Proportions are illustrative.

Parameter	Initialization	Learning rate	Weight decay	Forward multiplier
Token embedding (first $sg$ columns)	$\sigma_{0}$	$\eta_{0}$	$\lambda_{0}$	$1$
$W^{\mathrm{init}}$ of attention and FFN up/gate	$\sigma_{0}/\sqrt{m}$	frozen	–	$s^{-1/2}$
$W^{\mathrm{init}}$ of FFN down	$\sigma_{0}\sqrt{e}/\sqrt{m}$	frozen	–	$s^{-1/2}$
$\Delta$ of every hidden matrix	$0$	$\eta_{0}/m$	$\lambda_{0}m$	$s^{-1}$
$W^{\mathrm{init}}$ of the readout (first $sg$ columns)	$\sigma_{0}/\sqrt{m}$	frozen	–	$\alpha_{s}s^{-1/2}$
$\Delta$ of the readout (first $sg$ columns)	$0$	$\eta_{0}/m$	$\lambda_{0}m$	$\alpha_{s}s^{-1}$

Appendix

Table 5:

\mu

-Elastic for the fully elastic baseline.

s

is the sub-network width in units of

g

; multipliers act on the sliced tensor in the forward pass (Equation 23 ), and the FFN down projection multiplies its output, from both terms, by

1/e

Parameter	Initialization	Learning rate	Weight decay	Forward multiplier
Token embedding (full width)	$\sigma_{0}$	$\eta_{0}$	$\lambda_{0}$	$1$
$W^{\mathrm{init}}$ of attention and FFN up/gate	$\sigma_{0}/\sqrt{m}$	frozen	–	$1$
$\Delta$ of attention and FFN up/gate	$0$	$\eta_{0}/m$	$\lambda_{0}m$	$1$
$W^{\mathrm{init}}$ of FFN down (first $D^{(s)}_{\mathrm{ffn}}$ columns)	$\sigma_{0}\sqrt{e}/\sqrt{m}$	frozen	–	$s^{-1/2}$
$\Delta$ of FFN down	$0$	$\eta_{0}/m$	$\lambda_{0}m$	$s^{-1}$
Shared readout (fan-in $D_{\mathrm{model}}$ )	$\sigma_{0}/\sqrt{m}$	$\eta_{0}/m$	$\lambda_{0}m$	$\alpha$ (one scalar)

Appendix

Table 6:

\mu

-Elastic for MatFormer.

s=D^{(s)}_{\mathrm{ffn}}/h_{0}

is the FFN hidden width in units of the base FFN width

h_{0}

; the hidden matrices other than the FFN down projection and the readout have a fixed fan-in and take the non-elastic

\mu

P factor

m=D_{\mathrm{model}}/d_{0}

, and the embedding and gains keep the base values. The FFN down projection multiplies its output, from both terms, by

1/e

$119$ M– $285$ M scale
Model	$g$	$\mathcal{S}$	Sub-network widths	$d_{\mathrm{head}}$	$D_{\mathrm{ffn}}$	Parameters (M)
Proxy	$128$	$\{1\}$	$128$	$128$	$128$	$17.1$
Standalone	$64$	$\{M\}$ , $M=6,7,8,9,10$	$384$ – $640$	$64$	$4D_{\mathrm{model}}$	$119.0$ – $287.5$
MatFormer	$128$	$\{4,9,15,22,30\}$	FFN $512$ – $3840$	$128$	$3840$	$118.8$ – $282.4$
Fully elastic	$64$	$\{6,7,8,9,10\}$	$384$ – $640$	$64$	$2560$	$119.0$ – $287.5$
ShellElastic	$128$	$\{4,5,6,7,8\}$	$512$ – $1024$	$128$	$1024$	$118.8$ – $279.6$

Appendix

Table 7: Model configurations. Each model is defined by its base width

g

and supported set

\mathcal{S}

(Section 3.1 ): its sub-networks have widths

sg

for

s\in\mathcal{S}

, and a non-elastic model has

\mathcal{S}=\{M\}

. MatFormer nests only the FFN hidden width, so its

g

and

\mathcal{S}

refer to that width at a fixed

D_{\mathrm{model}}

(

512

, and

1024

at the two larger scales). All models have

32

layers and a

50{,}432

-token vocabulary; parameters count the non-zero weights of each sub-network (Appendix E.3 ).

Shared by all models
Layers	$32$ pre-norm layers with residual-branch multiplier $1$
Normalization	RMSNorm with learnable gains and $\epsilon_{N}=10^{-5}$ ; query–key RMSNorm per head
Attention	rotary embeddings with base $10{,}000$ , logit scale $d_{\mathrm{head}}^{-1/2}$ , elementwise sigmoid gate
FFN and biases	SwiGLU; no biases
Embeddings	untied, over the $50{,}432$ -token GPT-NeoX vocabulary
Initialization	truncated normal at $\pm 3$ standard deviations

Appendix

Table 8: Architecture settings at the

119

M–

285

M scale. The upper block is shared by all models, and the lower block lists the settings that differ between them (base widths, supported sets and sizes of all three scales in Table 7 ).

HP	Value
Sequence length	$2{,}048$
Batch size	$128$ sequences ( $262{,}144$ tokens)
Training steps	$25{,}000$ ( $6.55$ B tokens)
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$

Appendix

Table 9: Training HPs used unless otherwise specified (Appendix E.1 lists the changes at the two larger scales). The base values

\eta_{0}

\sigma_{0}

and

\lambda_{0}

are scaled per parameter group by

\mu

-Elastic (Table 1 and Appendix D ) or by

\mu

P; SP runs use them verbatim.

	Method	Low	Medium	High	Extra	Max	Avg.
$119$ M– $285$ M	Standalone	119.0	154.4	194.3	238.7	287.5	198.8
	Sampled standalone	119.0	154.4	194.3	238.7	287.5	198.8
	MatFormer	118.8	150.3	188.0	232.1	282.4	194.3
	Fully elastic	119.0	154.4	194.3	238.7	287.5	198.8
	ShellElastic	118.8	152.7	190.8	233.1	279.6	195.0
$375$ M– $906$ M	Standalone	377.5	487.7	611.8	749.8	901.6	625.7

Appendix

Table 10: Parameters (M) of the sub-networks of Table 2 at each scale and tier, counted as in Appendix E.3 ; Avg. averages the five tiers. Standalone, sampled standalone and the fully elastic model share the same five dense architectures.

Train latency (ms, forward + backward)
Method	Tier 1	Tier 2	Tier 3	Tier 4	Tier 5	Avg.
Method	Low	Medium	High	Extra	Max	Avg.
Standalone †	570 $\pm$ 2	644 $\pm$ 3	723 $\pm$ 6	834 $\pm$ 5	917 $\pm$ 10	737 $\pm$ 3
Sampled standalone	as Standalone
MatFormer	641 $\pm$ 1	677 $\pm$ 2	722 $\pm$ 4	776 $\pm$ 5	834 $\pm$ 6	730 $\pm$ 4
Fully elastic	597 $\pm$ 4	671 $\pm$ 3	748 $\pm$ 5	858 $\pm$ 7	933 $\pm$ 9	761 $\pm$ 6

Appendix

Table 11: Latency and peak training memory at the

119

M–

285

M scale on A100-40GB GPUs, for one GPU’s batch of

16

sequences of

2{,}048

tokens, with each tier run as its own sub-network (protocol in Appendix E.4 ). Entries are means over runs, with the standard deviation as a subscript; Avg. averages the tiers. † Five separately trained models.

Field	Dataset	Split	#Items	Metric	Random
World knowledge	MMLU	test	$14{,}042$	MC accuracy	$25\%$
	ARC-Easy	test	$2{,}376$	MC accuracy	$25\%$
	ARC-Challenge	test	$1{,}172$	MC accuracy	$25\%$
Symbolic problem solving	SVAMP	test	$300$	EM	$0\%$
Commonsense reasoning	COPA	val	$100$	MC accuracy	$50\%$
Commonsense reasoning	PIQA	val	$1{,}838$	MC accuracy	$50\%$

Appendix

Table 12: Downstream evaluation suite. All tasks are evaluated 5-shot; ARC-Easy, PIQA and HellaSwag are also probed during pretraining. MC: multiple-choice accuracy; EM: exact match.

5-shot accuracy (%) $\uparrow$
Tier 1 $\cdot$ Low	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	119.0	119.0	118.8	119.0	118.8
Train loss $\downarrow$	3.1951 $\pm$ .0035	3.4911 $\pm$ .0235	3.2536 $\pm$ .0032	3.2572 $\pm$ .0091	3.2305 $\pm$ .0015
Val. loss $\downarrow$	3.3695 $\pm$ .0034	3.6674 $\pm$ .0242	3.4248 $\pm$ .0036	3.4291 $\pm$ .0097	3.4086 $\pm$ .0009
MMLU	24.5 $\pm$ .3	24.4 $\pm$ .2	24.8 $\pm$ .5	24.7 $\pm$ .7	24.7 $\pm$ .6
ARC-Easy	44.8 $\pm$ .5	39.2 $\pm$ 1.0	43.5 $\pm$ .5	44.0 $\pm$ .9	45.0 $\pm$ .2

Appendix

Table 13: Per-task results at the

119

M–

285

M scale, Tier 1 (Low, about

119

M parameters). Sub-network width

384

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

512

D_{\mathrm{model}}=512

for MatFormer, and

512

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 2 $\cdot$ Medium	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	154.4	154.4	150.3	154.4	152.7
Train loss $\downarrow$	3.1344 $\pm$ .0019	3.4301 $\pm$ .0095	3.1821 $\pm$ .0034	3.1783 $\pm$ .0083	3.1641 $\pm$ .0031
Val. loss $\downarrow$	3.3081 $\pm$ .0016	3.6091 $\pm$ .0104	3.3566 $\pm$ .0037	3.3536 $\pm$ .0083	3.3375 $\pm$ .0027
MMLU	25.0 $\pm$ .4	24.9 $\pm$ .3	24.9 $\pm$ .2	24.7 $\pm$ .3	23.8 $\pm$ .7
ARC-Easy	46.7 $\pm$ .1	39.1 $\pm$ 1.1	45.2 $\pm$ 1.0	46.1 $\pm$ .8	45.8 $\pm$ .1

Appendix

Table 14: Per-task results at the

119

M–

285

M scale, Tier 2 (Medium, about

153

M parameters). Sub-network width

448

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

1152

D_{\mathrm{model}}=512

for MatFormer, and

640

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 3 $\cdot$ High	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	194.3	194.3	188.0	194.3	190.8
Train loss $\downarrow$	3.0916 $\pm$ .0086	3.4087 $\pm$ .0336	3.1561 $\pm$ .0036	3.1447 $\pm$ .0077	3.1229 $\pm$ .0018
Val. loss $\downarrow$	3.2635 $\pm$ .0092	3.5826 $\pm$ .0344	3.3250 $\pm$ .0036	3.3145 $\pm$ .0076	3.2944 $\pm$ .0016
MMLU	24.6 $\pm$ .4	25.0 $\pm$ .3	24.6 $\pm$ .4	24.6 $\pm$ .5	24.0 $\pm$ .5
ARC-Easy	47.9 $\pm$ .2	40.0 $\pm$ .6	45.9 $\pm$ .7	46.6 $\pm$ 1.0	47.1 $\pm$ .8

Appendix

Table 15: Per-task results at the

119

M–

285

M scale, Tier 3 (High, about

192

M parameters). Sub-network width

512

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

1920

D_{\mathrm{model}}=512

for MatFormer, and

768

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 4 $\cdot$ Extra	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	238.7	238.7	232.1	238.7	233.1
Train loss $\downarrow$	3.0474 $\pm$ .0006	3.3561 $\pm$ .0056	3.1407 $\pm$ .0036	3.1233 $\pm$ .0071	3.0941 $\pm$ .0015
Val. loss $\downarrow$	3.2186 $\pm$ .0012	3.5319 $\pm$ .0056	3.3123 $\pm$ .0037	3.2954 $\pm$ .0072	3.2641 $\pm$ .0017
MMLU	24.6 $\pm$ .5	24.7 $\pm$ .1	24.5 $\pm$ .4	24.6 $\pm$ .5	24.4 $\pm$ .7
ARC-Easy	50.7 $\pm$ 1.0	42.0 $\pm$ .7	46.3 $\pm$ .2	47.2 $\pm$ .9	47.6 $\pm$ .6

Appendix

Table 16: Per-task results at the

119

M–

285

M scale, Tier 4 (Extra, about

236

M parameters). Sub-network width

576

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

2816

D_{\mathrm{model}}=512

for MatFormer, and

896

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 5 $\cdot$ Max	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	287.5	287.5	282.4	287.5	279.6
Train loss $\downarrow$	3.0203 $\pm$ .0017	3.3505 $\pm$ .0321	3.1436 $\pm$ .0035	3.1196 $\pm$ .0070	3.0713 $\pm$ .0014
Val. loss $\downarrow$	3.1908 $\pm$ .0014	3.5239 $\pm$ .0340	3.3139 $\pm$ .0036	3.2899 $\pm$ .0071	3.2403 $\pm$ .0016
MMLU	24.8 $\pm$ .9	24.5 $\pm$ .2	24.8 $\pm$ .4	24.3 $\pm$ .2	24.2 $\pm$ .4
ARC-Easy	50.5 $\pm$ .5	41.1 $\pm$ .1	46.2 $\pm$ .5	47.1 $\pm$ .9	48.0 $\pm$ .7

Appendix

Table 17: Per-task results at the

119

M–

285

M scale, Tier 5 (Max, about

285

M parameters). Sub-network width

640

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

3840

D_{\mathrm{model}}=512

for MatFormer, and

1024

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 1 $\cdot$ Low	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	377.5	377.5	371.9	377.5	371.9
Train loss $\downarrow$	2.8302 $\pm$ .0016	3.0381 $\pm$ .0070	2.8993 $\pm$ .0014	2.8791 $\pm$ .0019	2.8746 $\pm$ .0035
Val. loss $\downarrow$	3.0069 $\pm$ .0016	3.2085 $\pm$ .0079	3.0724 $\pm$ .0015	3.0519 $\pm$ .0023	3.0546 $\pm$ .0039
MMLU	24.9 $\pm$ .5	25.1 $\pm$ .4	25.1 $\pm$ .2	23.9 $\pm$ .7	24.8 $\pm$ .0
ARC-Easy	57.3 $\pm$ .8	50.3 $\pm$ .5	54.9 $\pm$ .6	56.1 $\pm$ .4	55.8 $\pm$ .8

Appendix

Table 18: Per-task results at the

375

M–

906

M scale, Tier 1 (Low, about

375

M parameters). Sub-network width

832

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

1024

D_{\mathrm{model}}=1024

for MatFormer, and

1024

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 2 $\cdot$ Medium	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	487.7	487.7	485.1	487.7	481.6
Train loss $\downarrow$	2.7862 $\pm$ .0010	3.0000 $\pm$ .0032	2.8283 $\pm$ .0015	2.8172 $\pm$ .0013	2.8222 $\pm$ .0011
Val. loss $\downarrow$	2.9632 $\pm$ .0010	3.1750 $\pm$ .0039	3.0050 $\pm$ .0014	2.9948 $\pm$ .0015	2.9998 $\pm$ .0012
MMLU	25.2 $\pm$ .5	25.2 $\pm$ .7	24.7 $\pm$ .3	25.2 $\pm$ .1	24.8 $\pm$ .5
ARC-Easy	59.6 $\pm$ 1.1	50.8 $\pm$ .4	58.0 $\pm$ .2	58.2 $\pm$ .9	57.4 $\pm$ 1.1

Appendix

Table 19: Per-task results at the

375

M–

906

M scale, Tier 2 (Medium, about

486

M parameters). Sub-network width

960

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

2176

D_{\mathrm{model}}=1024

for MatFormer, and

1280

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 3 $\cdot$ High	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	611.8	611.8	610.9	611.8	608.1
Train loss $\downarrow$	2.7496 $\pm$ .0018	2.9719 $\pm$ .0037	2.8035 $\pm$ .0016	2.7933 $\pm$ .0011	2.7901 $\pm$ .0013
Val. loss $\downarrow$	2.9262 $\pm$ .0017	3.1414 $\pm$ .0040	2.9750 $\pm$ .0015	2.9662 $\pm$ .0016	2.9670 $\pm$ .0016
MMLU	24.1 $\pm$ .5	24.6 $\pm$ .9	24.9 $\pm$ .9	25.0 $\pm$ .8	24.8 $\pm$ .3
ARC-Easy	60.6 $\pm$ .5	52.4 $\pm$ .6	58.8 $\pm$ .3	59.8 $\pm$ .4	58.5 $\pm$ .6

Appendix

Table 20: Per-task results at the

375

M–

906

M scale, Tier 3 (High, about

611

M parameters). Sub-network width

1088

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

3456

D_{\mathrm{model}}=1024

for MatFormer, and

1536

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 4 $\cdot$ Extra	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	749.8	749.8	749.3	749.8	751.4
Train loss $\downarrow$	2.7201 $\pm$ .0016	2.9403 $\pm$ .0057	2.7857 $\pm$ .0015	2.7735 $\pm$ .0011	2.7671 $\pm$ .0010
Val. loss $\downarrow$	2.8972 $\pm$ .0016	3.1146 $\pm$ .0057	2.9628 $\pm$ .0016	2.9520 $\pm$ .0014	2.9435 $\pm$ .0016
MMLU	25.1 $\pm$ .5	23.9 $\pm$ .1	24.6 $\pm$ .5	24.4 $\pm$ .5	24.5 $\pm$ .3
ARC-Easy	62.4 $\pm$ 1.1	52.2 $\pm$ .6	59.9 $\pm$ .4	60.4 $\pm$ .1	58.9 $\pm$ .9

Appendix

Table 21: Per-task results at the

375

M–

906

M scale, Tier 4 (Extra, about

750

M parameters). Sub-network width

1216

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

4864

D_{\mathrm{model}}=1024

for MatFormer, and

1792

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 5 $\cdot$ Max	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	901.6	901.6	912.9	901.6	911.5
Train loss $\downarrow$	2.6922 $\pm$ .0012	2.9196 $\pm$ .0026	2.7871 $\pm$ .0015	2.7710 $\pm$ .0009	2.7488 $\pm$ .0008
Val. loss $\downarrow$	2.8696 $\pm$ .0010	3.0931 $\pm$ .0032	2.9631 $\pm$ .0016	2.9481 $\pm$ .0015	2.9250 $\pm$ .0012
MMLU	24.9 $\pm$ .5	24.8 $\pm$ .2	24.8 $\pm$ .0	24.9 $\pm$ .5	24.8 $\pm$ .1
ARC-Easy	63.4 $\pm$ .8	54.2 $\pm$ 1.1	60.0 $\pm$ .7	60.5 $\pm$ .1	59.9 $\pm$ .4

Appendix

Table 22: Per-task results at the

375

M–

906

M scale, Tier 5 (Max, about

906

M parameters). Sub-network width

1344

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

6528

D_{\mathrm{model}}=1024

for MatFormer, and

2048

for ShellElastic . Mean over three runs, with the standard deviation as a subscript.

5-shot accuracy (%) $\uparrow$
Tier 1 $\cdot$ Low	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	676.6	676.6	673.8	676.6	673.8
Train loss $\downarrow$	2.7335	2.8643	2.7365	2.7638	2.7347
Val. loss $\downarrow$	2.9118	3.0370	2.9140	2.9403	2.9135
MMLU	25.2	23.6	24.1	23.3	24.1
ARC-Easy	61.8	57.7	60.9	60.3	62.4

Appendix

Table 23: Per-task results at the

675

M–

1.69

B scale, Tier 1 (Low, about

675

M parameters). Sub-network width

832

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

4096

D_{\mathrm{model}}=1024

for MatFormer, and

1024

for ShellElastic . Single runs.

5-shot accuracy (%) $\uparrow$
Tier 2 $\cdot$ Medium	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	885.8	885.8	875.2	885.8	878.0
Train loss $\downarrow$	2.6791	2.8253	2.6938	2.7067	2.6865
Val. loss $\downarrow$	2.8587	2.9980	2.8720	2.8840	2.8657
MMLU	23.8	24.0	24.6	24.1	23.6
ARC-Easy	64.3	58.3	62.7	62.5	62.4

Appendix

Table 24: Per-task results at the

675

M–

1.69

B scale, Tier 2 (Medium, about

882

M parameters). Sub-network width

960

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

6144

D_{\mathrm{model}}=1024

for MatFormer, and

1280

for ShellElastic . Single runs.

5-shot accuracy (%) $\uparrow$
Tier 3 $\cdot$ High	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	1123.2	1123.2	1126.8	1123.2	1117.7
Train loss $\downarrow$	2.6424	2.7994	2.6730	2.6788	2.6550
Val. loss $\downarrow$	2.8212	2.9696	2.8487	2.8539	2.8351
MMLU	25.3	22.9	24.3	24.3	24.2
ARC-Easy	64.6	58.8	62.9	63.6	63.6

Appendix

Table 25: Per-task results at the

675

M–

1.69

B scale, Tier 3 (High, about

1.12

B parameters). Sub-network width

1088

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

8704

D_{\mathrm{model}}=1024

for MatFormer, and

1536

for ShellElastic . Single runs.

5-shot accuracy (%) $\uparrow$
Tier 4 $\cdot$ Extra	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	1388.5	1388.5	1403.7	1388.5	1393.1
Train loss $\downarrow$	2.6122	2.7838	2.6617	2.6621	2.6323
Val. loss $\downarrow$	2.7933	2.9557	2.8385	2.8384	2.8131
MMLU	23.7	24.6	24.2	24.8	24.2
ARC-Easy	66.2	59.7	63.0	64.3	65.3

Appendix

Table 26: Per-task results at the

675

M–

1.69

B scale, Tier 4 (Extra, about

1.39

B parameters). Sub-network width

1216

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

11520

D_{\mathrm{model}}=1024

for MatFormer, and

1792

for ShellElastic . Single runs.

5-shot accuracy (%) $\uparrow$
Tier 5 $\cdot$ Max	Standalone †	Sampled standalone	MatFormer	Fully elastic	ShellElastic
Params (M)	1682.0	1682.0	1705.6	1682.0	1704.2
Train loss $\downarrow$	2.5859	2.7509	2.6585	2.6549	2.6152
Val. loss $\downarrow$	2.7674	2.9242	2.8378	2.8334	2.7959
MMLU	25.2	24.7	24.2	24.3	24.6
ARC-Easy	67.3	61.7	62.8	64.6	65.9

Appendix

Table 27: Per-task results at the

675

M–

1.69

B scale, Tier 5 (Max, about

1.69

B parameters). Sub-network width

1344

for Standalone, sampled standalone and the fully elastic model, FFN hidden width

14592

D_{\mathrm{model}}=1024

for MatFormer, and

2048

for ShellElastic . Single runs.

$119$ M– $285$ M, train loss $\downarrow$
Method	Low	Medium	High	Extra	Max	Avg.
MatFormer	3.2614	3.1841	3.1566	3.1405	3.1393	3.1764
w/ $\mu$ -Elastic	3.2536	3.1821	3.1561	3.1407	3.1436	3.1752
Fully elastic	3.2717	3.1917	3.1544	3.1304	3.1263	3.1749
w/ $\mu$ -Elastic	3.2572	3.1783	3.1447	3.1233	3.1196	3.1646
ShellElastic	3.2472	3.1801	3.1356	3.1040	3.0794	3.1493

Appendix

Table 28: Effect of

\mu

-Elastic on each elastic model at the

119

M–

285

M and

375

M–

906

M scales , the values of Table 3 with four decimals.

5-shot accuracy (%) $\uparrow$
	Tier 1 $\cdot$ Low		Tier 2 $\cdot$ Medium		Tier 3 $\cdot$ High		Tier 4 $\cdot$ Extra		Tier 5 $\cdot$ Max
	Stair.	Ours	Stair.	Ours	Stair.	Ours	Stair.	Ours	Stair.	Ours
Params (M)	32.6	118.8	76.9	152.7	132.8	190.8	200.4	233.1	279.6	279.6
Train loss $\downarrow$	3.6874	3.2305	3.4224	3.1641	3.2764	3.1229	3.1871	3.0941	3.1258	3.0713
Val. loss $\downarrow$	3.8836	3.4086	3.6004	3.3375	3.4519	3.2944	3.3588	3.2641	3.2949	3.2403
MMLU	25.1	24.7	25.1	23.8	25.1	24.0	25.5	24.4	25.3	24.2

Appendix

Table 29: Per-task results of StairFormer and ShellElastic at every tier of the

119

M–

285

M scale, with the rows of Appendix G . Stair.: StairFormer. Ours: ShellElastic . Bold marks the better entry of each pair (ties both bold). Only Tier 5 is parameter-matched.

$B=42$ , 5 supported widths
$\mathcal{S}$	128	256	384	512	640	768	896	1024	Avg.
Params	17.1	38.4	63.9	118.8	152.7	190.8	233.1	279.6
$\{1,2,3,4,8\}$	4.1178	3.8183	3.6738	3.6023 †	–	–	–	3.2698	3.6964
$\{1,2,3,7,8\}$	4.1037	3.8310	3.7224	–	–	–	3.2910	3.2503	3.6397
$\{1,2,6,7,8\}$	4.0895	3.8832	–	–	–	3.3193	3.2721	3.2400	3.5608
$\{1,5,6,7,8\}$	4.1154	–	–	–	3.3657	3.3068	3.2687	3.2409	3.4595

Appendix

Table 30: Validation loss of every supported width in the architecture search , at step

25{,}000

, as the one group wider than

g

moves from the last shell (first row of each block) to the core (last row, bold). Params: non-zero parameters (M) of each width; † a width below the wide group, which has fewer parameters. Avg.: mean over the supported widths of each set. Bold marks the lowest value in each column.

(1) Selected sub-networks per step $\mathcal{A}_{\tau}$
	Val. loss $\downarrow$						5-shot acc. $\uparrow$
Setting	Low	Medium	High	Extra	Max	Avg.	Avg.
Sandwich: Low, Max and one other	3.4019	3.3617	3.3183	3.2816	3.2395	3.3206	36.1
Random: one of the five sub-networks	3.4311	3.3786	3.3542	3.3410	3.3356	3.3681	34.8
Max + random: Max and one other	3.4944	3.4079	3.3447	3.2956	3.2397	3.3564	35.5
All: all five sub-networks (default)	3.4086	3.3375	3.2944	3.2641	3.2403	3.3090	36.5

Appendix

Table 31: Ablations (1)–(6) of ShellElastic at the

119

M–

285

M scale. Each block changes one setting of the default recipe and ends with the default. Val. loss is taken at step

25{,}000

for each tier (Low to Max: widths

512

1024

), and 5-shot acc. averages the

13

downstream tasks over the five tiers. Bold marks the best value of each column within a block.

(7) Normalization type
	Val. loss $\downarrow$						5-shot acc. $\uparrow$
Setting	Low	Medium	High	Extra	Max	Avg.	Avg.
LayerNorm	3.4037	3.3379	3.2955	3.2655	3.2424	3.3090	35.8
RMSNorm (default)	3.4086	3.3375	3.2944	3.2641	3.2403	3.3090	36.5
(8) Ensembles of the five sub-networks
Learned per-token width router	3.7735	3.5611	3.4137	3.3055	3.2269	3.4562	35.1

Appendix

Table 32: Ablations (7)–(12) of ShellElastic at the

119

M–

285

M scale , in the format of Table 31 .

Predictors
	Loss $-$ Max loss
Ensemble	Train loss	Val. loss
Uniform weights	$+0.0313$	$+0.0316$
Confidence weights	$+0.0295$	$+0.0296$
Most confident readout	$+0.0648$	$+0.0678$
Best fixed ensemble	$0.0000$	$-0.0003$

Appendix

Table 33: Ensembles of the five sub-networks as predictors and as teachers. Train and validation loss minus those of Max. The row labeled Router is the ensemble weighted per token by the learned router, which is trained together with its own supernet, so this row subtracts the losses of that supernet’s Max. The loss-weighted ensemble and the soft oracle are the two ensembles used as distillation teachers in ablation (8).

Tier	Low	Medium	High	Extra	Max
Picked by the oracle (%)	22.7	16.3	15.1	16.4	29.4
Contribution to the gain (nats)	0.0878	0.0454	0.0309	0.0210	–
Beats Max (%)	39.0	42.7	44.9	46.7	–

Appendix

Table 34: The oracle on held-out C4 text (final checkpoints). Top: share of tokens on which each tier assigns the target the highest probability (the oracle’s pick), its contribution to the oracle’s mean gain over Max (the four sum to

0.185

), and the share of tokens on which it beats Max. Bottom: validation loss minus that of Max when the models in each row are ensembled with uniform weights, the best fixed weights or the oracle.

Elastic Neural Networks

See all 9 papers

Your Transformer Is Secretly Elastic
Oct 14, 2026Walter Simoncini, Dheeraj Varghese, Gertjan J. Burghouts +1Elastic Neural Networks Transformer Pruning
University of Amsterdam · TNO, Intelligent Imaging
Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control
May 8, 2026Ali Taghibakhshi, Ruisi Cai, Saurav Muralidharan +17LLM Compression LLM Post-Training
NVIDIA
Stepped MoE: Segment-Level Routing with Configurable Inference Complexity
Oct 5, 2026Arnav Kundu, Zhaoyang Xu, Bairu Hou +3On-Device Language Model Inference Dynamic Neural Networks
Apple

Oct 14, 2026J/K move · Enter open · S save

Your Transformer Is Secretly Elastic

Walter Simoncini, Dheeraj Varghese, Gertjan J. Burghouts, Cees G. M. Snoek

University of Amsterdam · TNO, Intelligent Imaging

Elastic Neural Networks Transformer Pruning Efficient Transformer Inference Neural Network Pruning Dynamic Neural Networks Structured Pruning+1

Open paper AskPDF arXiv

Scaling up transformers has produced models with broad capabilities, deployed across heterogeneous hardware and fluctuating workloads. Yet, model families offer only a few fixed sizes (e.g., S, B, L), forcing users to commit to a model based on hardware and latency constraints. This rigidity is often suboptimal under varying query loads. We introduce a retraining-free method that elastifies pre-trained transformers into a single model that adapts on-the-fly across a continuum of computational budgets. Our key contribution is a differentiable relaxation of the functional unit-ranking problem: unlike prior work, which optimizes a single importance weight per MLP block or attention head, we rank every individual MLP neuron and attention head jointly, on a common scale, by minimizing the model's loss under soft pruning across many sparsity targets, all at once. Across 27 vision and text transformer encoders ranging from 22M to 8B parameters, our method yields smooth degradation up to 60% sparsity and outperforms one-shot baselines under deep pruning. For example, pruning DINOv3 ViT-H+/16 to 50% sparsity removes 427M parameters while reducing ImageNet-1k linear-probing accuracy by only 0.8%, without the need for weight correction or retraining. At 60% sparsity, a pruned AugReg ViT-B/16 model achieves up to

1.80\times

higher throughput and

1.26\times

lower single-query latency on an A100 GPU. Elastic models are produced in minutes, with construction time scaling sub-linearly with parameter count.

Abstract

Elastic Neural Networks

Explore similar work

Your Transformer Is Secretly Elastic

Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control

Stepped MoE: Segment-Level Routing with Configurable Inference Complexity