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:

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

† 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):

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.