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

> ShellElastic's block lower-triangular architecture enables one forward pass to train all nested sub-networks simultaneously, matching separately trained models that use 5x more compute.

- **Source:** [arXiv](https://arxiv.org/abs/2610.15124)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/X9AGgf
- **Whiteboard:** https://picx.dev/p/X9AGgf/image

## Summary

## Summary (Overview)

- **ShellElastic** introduces a novel elastic transformer architecture where hidden dimensions are organized into a dense core followed by successive shells, with every hidden matrix being block lower triangular. This enables one full-width forward pass to contain the exact hidden states of every sub-network as prefixes, making joint training cost comparable to training the widest sub-network alone.
- **μ-Elastic** extends maximal update parameterization (μP) to elastic models with two scaling factors: a μ-Transfer factor $m = g/d_0$ for hyperparameter transfer from non-elastic proxies, and a μ-Internal factor for maintaining simultaneous maximal updates across all sub-networks.
- **Theoretical guarantees**: The paper proves (Theorem 1) that μ-Elastic keeps all sub-networks maximally updated simultaneously, and (Corollary 1) that hyperparameters tuned on a small non-elastic proxy transfer to elastic families and larger supernets.
- **Empirical results**: A single ShellElastic run produces sub-networks that outperform MatFormer and fully elastic baselines at matched parameter counts, approaching separately trained models that together take 5× the steps.
- **Comprehensive ablations** reveal that training all sub-networks at every step works best, and that a wide core with thin shells is the optimal architecture configuration.

## Introduction and Theoretical Foundation

### Background and Motivation

Foundation models must support diverse inference budgets across different tasks and deployment platforms. Elastic training addresses this by optimizing several nested sub-networks within one shared parameter set. However, existing elastic training recipes face three critical obstacles when scaled to foundation-model pretraining:

1. **Scalability**: Standard recipes sample one or a few sub-networks per step, causing variable computation graphs and workloads that complicate distributed training.
2. **Transferability**: Hyperparameters tuned at one model size may not remain effective at another, requiring expensive re-tuning.
3. **Simultaneous maximal updates**: Existing recipes don't ensure all sub-networks receive maximal updates simultaneously, causing feature updates to vanish or blow up across different widths.

### Theoretical Foundation

The work builds on **maximal update parameterization (μP)** (Yang et al., 2021), which ensures that a parameter group is *maximally updated* when its pre-activations are $\Theta(1)$ at initialization and move by $\Theta(\eta_0)$ after one optimizer step. The paper extends this to elastic models with two key concepts:

- **μ-Transfer**: HPs tuned on a non-elastic proxy serve an entire elastic family, and those tuned on a smaller elastic family serve larger ones.
- **μ-Internal**: A shared parameterization keeps every width maximally updated simultaneously despite parameter sharing.

The key insight is the composition of two scaling factors:

$$q(s) = \underbrace{m}_{\text{μ-Transfer}} \cdot \underbrace{s}_{\text{μ-Internal}} = \frac{sg}{d_0}, \quad s \in S$$

where $m = g/d_0$ is the ratio of base width to proxy width, and $s$ is the supported boundary index.

## Methodology

### ShellElastic Architecture

The architecture organizes hidden dimensions into a **dense core** followed by **successive shells**. For base width $g$, the supported sub-network widths are:

$$D^{(s)}_{\text{model}} = sg, \quad D_{\text{model}} = Mg$$

with $S = \{s_1 < \cdots < s_K = M\}$ being the supported width indices. Each hidden matrix is **block lower triangular**—a shell reads the core, earlier shells, and itself, but never a later shell.

**Group-local normalization** uses separate statistics within each group:

$$[N_k(x_{G_k})]_j = a_j \frac{x_j}{\sqrt{|G_k|^{-1} \sum_{u \in G_k} x_u^2 + \epsilon_N}}, \quad j \in G_k$$

**Proposition 1 (Exact reuse)**: One full-width trunk evaluation yields the same joint loss and parameter gradients as separate evaluations of all supported sub-networks:

$$T_s(E_s[u]) = P_s(T_t(E_t[u])), \quad s, t \in S, s \leq t$$

### μ-Elastic Parameterization

**Table 1: μ-Elastic scaling rules** for ShellElastic:

| Parameter | Active input width | Initialization | Learning rate | Weight decay | Multiplier |
|---|---|---|---|---|---|
| Token embedding | $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 | 1 (elementwise) | 1 | $\eta_0$ | 0 | 1 |

### Theoretical Guarantees

**Theorem 1 (μ-Internal)**: Under the rule of Table 1, every hidden group's pre-activations have $\Theta(1)$ coordinates at initialization and move by $\Theta(\eta_0)$ after one AdamW step, with constants independent of $s$, $t$, $M$, and $m$.

**Proposition 2**: If all groups share one initialization variance and learning rate, groups with boundaries $s < s'$ differ by factors $\Theta(\sqrt{s'/s})$ in pre-activation size and $\Theta(s'/s)$ in update size—no shared rule keeps every group maximal.

### Training and Inference

The joint objective at step $\tau$ with selected widths $A_\tau \subseteq S$:

$$\mathcal{L}_\tau = \frac{\sum_{s \in A_\tau} \omega_s \ell_s}{\sum_{s \in A_\tau} \omega_s}, \quad \ell_s = \text{CE}(z_s, y) + \lambda_z \mathbb{E}\left[\left(\log \sum_v e^{z_{s,v}}\right)^2\right]$$

**Cost comparison**:
- Joint training: $C_{\text{joint}} = C_{\text{trunk}}(M) + \sum_{s \in S} C_{\text{out}}(s)$
- Separate training: $C_{\text{separate}} = \sum_{s \in S} [C_{\text{trunk}}(s) + C_{\text{out}}(s)]$

## Empirical Validation / Results

### Main Results Across Three Scales

**Table 2: Final validation loss and mean 5-shot accuracy** at three scales (119M–285M, 375M–906M, 675M–1.69B):

| Scale | Method | Val. loss Avg. | 5-shot acc. Avg. (%) |
|---|---|---|---|
| 119M–285M | Standalone† | 3.270 | 36.5 |
| | Sampled standalone | 3.583 | 34.0 |
| | MatFormer | 3.347 | 35.6 |
| | Fully elastic | 3.336 | 35.6 |
| | **ShellElastic** | **3.309** | **36.5** |
| 375M–906M | Standalone† | 2.933 | 42.2 |
| | Sampled standalone | 3.147 | 38.2 |
| | MatFormer | 2.996 | 40.6 |
| | Fully elastic | 2.983 | 41.4 |
| | **ShellElastic** | **2.978** | **41.9** |
| 675M–1.69B | Standalone† | 2.830 | 44.8 |
| | Sampled standalone | 2.977 | 41.4 |
| | MatFormer | 2.862 | 43.4 |
| | Fully elastic | 2.870 | 43.6 |
| | **ShellElastic** | **2.845** | **44.0** |

† Standalone models take 5× the steps of any other row.

### Hyperparameter Transfer Results

Under μ-Elastic, the proxy's optimal hyperparameters ($\eta_0 = 2 \times 10^{-3}$, $\sigma_0 = 3.3 \times 10^{-3}$) transfer to all elastic models at both levels of μ-Transfer. The base values remain within 0.003 nats of the best point of every μ-Elastic sweep. Under standard parameterization (SP), the optimum departs from the proxy's and shifts with scale.

### Comparison with StairFormer

At matched parameter counts, StairFormer's validation loss is 0.055–0.083 nats higher than ShellElastic. Even StairFormer's Max (with 1.5× parameters) is 0.001 nats above ShellElastic's High tier.

### Architecture Search

With fixed block budget $B(S) = \sum_k (s_k - s_{k-1})s_k$, the **core is the best place for the wide group**: core placement has the lowest loss at every shared width, whereas shell placement makes the first width containing the wide group 0.025–0.052 nats worse.

### Key Ablations

**Table 4: Selected ablations** (validation loss and 5-shot accuracy):

| Setting | Low | Medium | High | Extra | Max | Avg. | 5-shot Avg. |
|---|---|---|---|---|---|---|---|
| **Selected sub-networks** | | | | | | | |
| Sandwich | 3.4019 | 3.3617 | 3.3183 | 3.2816 | 3.2395 | 3.3206 | 36.1 |
| Random | 3.4311 | 3.3786 | 3.3542 | 3.3410 | 3.3356 | 3.3681 | 34.8 |
| Max + random | 3.4944 | 3.4079 | 3.3447 | 3.2956 | 3.2397 | 3.3564 | 35.5 |
| **All (default)** | **3.4086** | **3.3375** | **3.2944** | **3.2641** | **3.2403** | **3.3090** | **36.5** |
| **In-place distillation** | | | | | | | |
| Forward KL, weight 1 | 3.3748 | 3.3184 | 3.2873 | 3.2653 | 3.2483 | 3.2988 | 36.1 |
| None (default) | 3.4086 | 3.3375 | 3.2944 | 3.2641 | 3.2403 | 3.3090 | 36.5 |

Key findings:
- **Training all sub-networks at every step** gives the lowest mean loss (0.012–0.059 nats below alternatives).
- **In-place distillation** improves Low by up to 0.034 nats but hurts Max and makes steps 2.2–2.6× longer.
- **Ensembles** don't beat Max alone; the oracle (using targets) is 0.19 nats below Max, showing potential for better label-free ensembles.
- **Stage-residual gradient scaling** ($\gamma = 0.3$–$0.8$) lowers Low's loss but doesn't improve the mean.

## Theoretical and Practical Implications

### Theoretical Contributions

1. **μ-Elastic provides provable guarantees**: Theorem 1 establishes that all sub-networks remain maximally updated simultaneously, with constants independent of scale. This extends μP theory to the shared-tensor setting of elastic supernets.

2. **Two-level hyperparameter transfer**: Corollary 1 proves that proxy-tuned HPs work for (i) elastic families and (ii) larger supernets (wider or with more shells), eliminating the need for repeated HP searches.

3. **Architecture-agnostic parameterization**: μ-Elastic extends to MatFormer, fully elastic models, and StairFormer with per-architecture definitions of $m$ and $s$.

### Practical Implications

1. **Fixed computation graph**: ShellElastic maintains a constant workload per step regardless of the number of sub-networks, simplifying distributed training infrastructure.

2. **Compute efficiency**: One ShellElastic run (25K steps) produces five model sizes that approach separately trained models taking 5× the steps, with the gap as small as 0.014 nats at the largest scale.

3. **Memory efficiency**: All readout gradients combine in one trunk backward pass, so memory requirements match the widest sub-network alone.

4. **Extensible architecture**: New shells can be appended to a trained supernet as new stages with frozen weights, enabling incremental model growth—a property MatFormer and fully elastic models lack.

## Conclusion

ShellElastic demonstrates that elastic pretraining can be both scalable and effective by ensuring that one full-width forward pass yields the exact hidden states of every sub-network. Combined with μ-Elastic parameterization, this approach:
- Makes joint training cost comparable to training the widest sub-network alone
- Transfers hyperparameters from small proxies to elastic families and larger supernets
- Keeps all sub-networks maximally updated simultaneously

**Future directions** identified by the authors:
1. **Better sub-network ensembles**: The oracle is 0.19 nats below Max, suggesting a label-free ensemble could close part of this gap at almost no extra cost.
2. **Shell appending**: Training new shells on a frozen trained supernet could enable continual model growth.
3. **Scaling to larger models and longer training**: The method's theoretical guarantees suggest it should scale favorably, warranting validation at larger scales.

The experiments favor training all sub-networks at every step and a wide core with thin shells, providing practical guidance for future elastic pretraining recipes.

---

_Markdown view of https://picx.dev/p/X9AGgf, served by PicX — AI-generated visual whiteboard summaries of research papers._