ShellElastic

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

Overview

Elastic pretraining trains a family of nested models, from small to large, in a single pretraining run. This overview explains what stands in its way at foundation-model scale and how ShellElastic and μ-Elastic address each obstacle. The Method and Results sections then give the details, each with its place in the paper.

Elastic pretraining: every model size from one run

Foundation models need to support different compute budgets for different tasks and deployment platforms. Pretraining a separate model for each budget multiplies the cost: the Standalone baseline of our experiments, five separately trained models, takes five times the optimizer steps of one elastic run (Table 2).

Elastic training offers a way out. It optimizes several nested sub-networks inside one shared parameter set: each smaller sub-network is contained in the larger ones, all of them are trained in the same run, and afterwards any of them can be extracted and deployed as a model of its own. Slimmable networks and Once-for-All established the idea for vision models; for transformers, MatFormer nests the FFN width and underlies released on-device models, while Flextron and Nemotron Elastic make already pretrained LLMs elastic by post-training (paper §1–2). Elastic pretraining asks for more: the nested family is trained from scratch, within the pretraining run itself.

Three problems at pretraining scale

Carried into foundation-model pretraining, the elastic training recipe meets three problems (paper §1).

  1. Problem 1 · Scalability

    The workload changes from step to step

    Standard recipes sample one or a few sub-networks per step, so the computation graph, the parameter access pattern and the workload change from step to step. Under parameter sharding, this variability complicates activation memory, recomputation and distributed scheduling, the very machinery that large-scale pretraining depends on. Training every sub-network at every step fixes the workload, but MatFormer and fully elastic models (a transformer adaptation of slimmable networks) do not nest their hidden states, so each width needs its own forward and backward pass, and the cost grows with the size of the family.

    Addressed by ShellElastic

  2. Problem 2 · Transferability

    Hyperparameters do not carry over

    Hyperparameters (HPs) tuned at one model size need not remain effective at another, and repeating the search over learning rate, initialization scale and weight decay at every target scale cancels the compute that elastic training saves. Maximal update parameterization (μP) transfers HPs across widths, but its rules assume separately trained models, in which each weight has one fan-in. Elastic models raise two questions: can HPs tuned on a non-elastic proxy serve an entire elastic family, and can those tuned on a smaller elastic family serve a larger one? We call these two levels of transfer μ-Transfer.

    Addressed by μ-Elastic, through its μ-Transfer factor

  3. Problem 3 · Simultaneous maximal updates

    One setting cannot suit every width

    A parameter group is maximally updated when the pre-activations it produces are Θ(1) at initialization and move by Θ(η₀) after one optimizer step, so that its contribution neither vanishes nor blows up as the width grows. In an elastic model the same weights serve sub-networks of different widths, and existing recipes do not generally keep all of them maximally updated at once: without width-dependent scaling, HPs that give well-scaled updates in one sub-network make updates vanish or blow up in another, and re-tuning one global value only trades one width against another. We call the missing property, every width maximally updated simultaneously, μ-Internal.

    Addressed by μ-Elastic, through its μ-Internal factor

Our approach: nest the hidden states, scale by the fan-in

We address the first problem with a new architecture, ShellElastic, and the other two with a parameterization, μ-Elastic.

Answers problem 1

ShellElastic: one full-width pass contains every sub-network

ShellElastic divides the residual stream into a dense core followed by ordered shells, and each sub-network is the core plus a prefix of the shells. Our main model has a 512-dimensional core and four 128-dimensional shells, giving five sub-networks of widths 512 to 1024. Every hidden matrix is block lower triangular, so a shell reads the core, the earlier shells and itself, but never a later shell; normalization acts within each group, attention heads align with the group boundaries, and every supported width has its own readout (Figure 1).

Every operation of the trunk therefore preserves the prefixes that end at group boundaries: the hidden states of each sub-network are exactly the first dimensions of the full-width ones (Proposition 1). One full-width pass thus yields the same joint loss and gradients as separate passes of all sub-networks. The per-step compute stays close to that of a single full-width pass, every step has the same computation graph, parameter access pattern and workload, and after training each sub-network is extracted as a standalone transformer that reproduces its supernet logits exactly.

Paper §3.1, §3.3 · proof in App. B

Answers problems 2 and 3

μ-Elastic: scale every group by the width it reads

μP scales a weight’s initialization and learning rate with its fan-in, the width of the input it reads; in an elastic model, the same weights serve several widths. In ShellElastic, however, the block row that produces a group reads the same input prefix in every sub-network that contains it, so its fan-in is well defined. μ-Elastic scales each group by this fan-in relative to a non-elastic proxy of width d₀, on which σ₀, η₀ and λ₀ are tuned: the ratio q = m · s combines a μ-Transfer factor m = g/d₀ with a μ-Internal factor s, the supported boundary of a hidden weight’s output group or a readout’s own width index. Hidden weights and readouts then take initialization σ₀/√q, learning rate η₀/q and weight decay λ₀q, with extra factors for the FFN down projection (Figure 2, Table 1).

Under the standard heuristic assumptions of μP, this keeps every group maximally updated in every sub-network at once (Theorem 1), which no rule shared by all groups can do in general (Proposition 2), and it keeps doing so when the supernet is widened or given more shells, so the proxy’s HPs can be reused (Corollary 1). The theory fixes the scales of initialization and updates, not the optimal values; whether the proxy’s values stay near-optimal is tested by experiment (Figure 3).

Paper §3.2 · proofs in App. C · other elastic architectures in App. D

What the experiments show

We pretrain on DCLM-Baseline at three scales, with sub-networks of about 119M–285M, 375M–906M and 675M–1.69B parameters. Each baseline provides five sizes, Low, Medium, High, Extra and Max, matched to the sub-networks of ShellElastic in parameter count to within 3%, and every elastic model uses μ-Elastic (paper §4.1).

  • One run, five competitive models. Among the methods that train all five sizes in one run, ShellElastic has the lowest mean validation loss and the highest mean 5-shot accuracy at every scale. With a fifth of their optimizer steps, it comes within 0.039, 0.045 and 0.014 nats of the five Standalone models in mean validation loss at the three scales, and within 0.8 points in mean 5-shot accuracy. Table 2 · paper §4.2
  • HPs tuned on the proxy transfer. Under μ-Elastic, the proxy’s optimum, η₀ = 2 × 10⁻³ and σ₀ = 3.3 × 10⁻³, transfers to all three elastic models at both levels of μ-Transfer: it is within 0.003 nats of the best point of every μ-Elastic sweep. Under SP, the optimum departs from the proxy’s and shifts with scale, and at these values SP is 0.03–0.11 nats worse. Figure 3 · paper §4.2
  • Re-tuning SP does not close the gap. Given the best of its 12 (13 at 375M–906M) swept configurations, SP still trails μ-Elastic with the untuned base values: by 0.011 and 0.013 nats in mean validation loss for the fully elastic model and ShellElastic at 119M–285M, and by 0.006–0.021 nats for all three elastic models at 375M–906M. Table 3 · paper §4.2
  • The architecture matters. At the parameter count of each ShellElastic sub-network, StairFormer, which also nests sub-networks as prefixes of one forward pass but uses equal-width groups, cumulative prefix normalization and one shared readout, has a validation loss 0.055 to 0.083 nats higher. At equal parameters and width count, the core is the best place for extra width: it has the lowest loss, or is within 0.006 nats of it, at every shared width. Figure 4 · paper §4.3
  • Five settings matter more than the rest. Of the settings we ablate at 119M–285M, one at a time, sub-network selection, in-place distillation, ensembles of the sub-networks, stage-residual gradient scaling and the trunk gradient of the narrower widths can change the mean validation loss by 0.010 to 0.20 nats, whereas the loss weights, normalization, readouts, base width and head dimension, and optimization HPs change it by at most 0.007 nats. Table 4 · paper §4.4 · App. J
  • Train every sub-network at every step. Of the selection rules we test, training all five sub-networks at every step gives the lowest mean loss, 0.012 to 0.059 nats below the others, and the highest mean 5-shot accuracy. Table 4 · paper §4.4
  • In-place distillation helps Low, not Max, and slows each step. Distilling Low to Extra from Max lowers the loss of Low by up to 0.034 nats and the mean loss by up to 0.010 nats, the largest gain of any ablation. Yet it raises the loss of Max by up to 0.008 nats, improves the mean 5-shot accuracy by at most 0.1 points and makes a step 2.2 to 2.6 times as long, so we do not use it. Table 4 · paper §4.4
  • Ensembles of the sub-networks do not yet beat Max. One trunk pass yields the predictions of all five sub-networks, but none of the five ensembles we test as predictors beats Max, and neither of the two we test as teachers helps Max. An oracle that uses the target, picking per token the sub-network that assigns it the highest probability, is 0.19 nats below Max, and we believe a better ensemble could close part of this gap, a promising way to improve ShellElastic at almost no extra cost. Figure 5 · paper §4.4
  • Stage-residual gradient scaling helps Low but hurts Extra and Max. We scale by γ the gradient that each sub-network’s loss sends to the trunk through the narrower sub-network it contains. A γ of 0.3 to 0.8 lowers the loss of Low by up to 0.028 nats but raises those of Extra and Max by up to 0.012 nats, leaving the mean loss within 0.002 nats, whereas γ = 0 raises it by 0.047 nats, so we keep γ = 1. Its stage view, the core giving a base prediction and each shell a correction, still points to a property that MatFormer and fully elastic models lack: a new shell can be appended to a trained supernet and trained with all trained weights frozen, so every trained sub-network keeps its exact outputs. Table 4 · paper §4.4

The rest of the page follows the paper. Method details the architecture, μ-Elastic with its guarantees, and training and inference (paper §3). Results & logs presents each table and figure of paper §4 with the paper’s reading of it and the runs behind it, and All runs lists every published run.

Method

Results & logs

Each result opens with the paper’s reading of it and its place in the paper, then shows the paper’s table or figure, followed by interactive views of the runs behind it: training curves, run files and training logs.

All runs

Every run published with this page. Search and filter the list, select up to 12 runs to compare their curves and final values, and open their data files or training logs.

BibTeX