Five Model Widths, One Shared Trunk

TL;DR for operators

A model platform offering five inference-capacity tiers ordinarily faces a choice between separately training five models and sharing parameters while varying the width trained at each step. ShellElastic¹ introduces a third arrangement: all five supported widths receive gradients through one shared full-width transformer trunk computation.

Its architectural guarantee is exact, provided the supported widths and operations satisfy specific constraints. Its companion optimizer parameterization, mu-Elastic, addresses the different update scales encountered by nested widths. Across three tested model scales, the combined method improves mean validation loss over the evaluated elastic alternatives. Independently trained models still achieve better average loss, however, at five times the aggregate optimizer steps. For training-platform operators, the central question is whether the shared architecture delivers acceptable quality across all service tiers at lower total GPU cost, not merely fewer training steps.

Multiple model sizes need not repeat the same trunk computation

Consider a foundation-model platform offering five capacity tiers. Smaller models serve requests with tighter inference budgets, while larger models handle workloads requiring more capacity. Training each model independently gives the platform considerable flexibility, but repeats much of the pretraining process.

Elastic models reduce that duplication by sharing parameters among nested sub-networks. Yet parameter sharing does not automatically mean computation sharing. Evaluating each width separately still requires multiple transformer executions. Alternatively, sampling different widths at successive training steps changes the computational workload and parameter-access pattern, complicating distributed scheduling.

ShellElastic tackles this problem by arranging the transformer so that a smaller supported model’s intermediate hidden states are exactly the leading dimensions of a larger model’s hidden states.

One full-width forward pass consequently provides the hidden representations required for every supported sub-network. The system still calculates separate terminal outputs and training losses for each width, but it does not need to execute the transformer trunk independently five times.

This distinction matters because it converts an elastic model’s shared parameters into reusable computation.

Exact reuse requires a deliberately restricted transformer

ShellElastic partitions hidden dimensions into a dense core followed by successive shells. In one configuration, a 512-dimensional core and four 128-dimensional shells support hidden widths of 512, 640, 768, 896, and 1,024.

The critical restriction is directional: an earlier group cannot depend on a later shell.

The architecture enforces that relationship through block-lower-triangular projections. Each output group reads its own and preceding input groups, never subsequent groups. Group-local RMSNorm, aligned attention heads, and structured feed-forward projections preserve the same dependency restriction across transformer layers.

The result is prefix preservation: running a smaller supported sub-network independently produces the same trunk hidden states as selecting the corresponding prefix from a larger supported sub-network.

Proposition 1 establishes this equivalence formally, including equivalent logits when the correct width-specific terminal operators are applied. Appendix B also reports a reduced two-layer numerical check with a maximum parameter-gradient difference of \(1.68\times10^{-8}\).

The computational consequence appears in the paper’s Equation 9:

Joint evaluation pays for one widest trunk and the terminal computations of the supported widths. Separate evaluation pays for every sub-network’s trunk as well as its output.

That is an analytical computation-reuse advantage, not a guarantee that training five widths costs exactly the same as training the widest model alone. Readouts, losses, and their gradients still require work.

Nor can arbitrary intermediate widths be extracted. Exact preservation applies at the designated shell boundaries, with compatible operators and stochastic behavior.

Shared computation still needs width-aware optimization

Even when all five models reuse the same hidden states, their parameters do not necessarily experience comparable optimization dynamics.

The problem is active fan-in: different parameter groups receive inputs from different numbers of hidden dimensions. A group in the core and a group incorporating later shells therefore operate at different effective input widths.

As the width ratios increase, one uniform initialization and learning-rate rule cannot generally maintain comparable feature-update scales across every group. Proposition 2 formalizes this incompatibility under the paper’s scaling assumptions.

ShellElastic addresses it with mu-Elastic, an extension of maximal update parameterization (muP).

The method combines two scaling dimensions:

Here, \(d_0\) is the width of a small non-elastic tuning model, \(g\) is the supernet’s base width, and \(s\) identifies the supported group boundary.

The factor \(m=g/d_0\) handles transfer from the small proxy to a larger model, while \(s\) handles differences between nested widths within the supernet.

For the relevant hidden-weight groups, mu-Elastic scales initialization standard deviation inversely with \(\sqrt{ms}\), AdamW learning rate inversely with \(ms\), and weight decay proportionally to \(ms\). Other parameter categories have their own rules.

Under stated heuristic assumptions about activations, gradients, and Adam updates, Theorem 1 shows that the supported groups maintain order-one pre-activations and consistently scaled one-step feature updates.

That theorem concerns optimization scaling. It does not prove that hyperparameters selected on a small model will be numerically optimal at every larger width.

The empirical transfer tests provide a separate result: hyperparameters tuned on a 128-wide non-elastic proxy remain within 0.003 nats of the best tested mu-Elastic configurations in mean validation loss across the evaluated families and sizes.

For an optimization team repeatedly changing supported model widths, this suggests that proxy tuning may replace some expensive full-scale hyperparameter searches. Its demonstrated scope remains the tested transformer and AdamW configurations.

Three model scales favor the combined training recipe

The researchers pretrained 32-layer decoder-only transformers on DCLM-Baseline, evaluated next-token loss on held-out C4 text, and measured 5-shot accuracy across 13 downstream datasets.

The main comparison covers five model tiers at each of three parameter scales, with approximately matched sub-network parameter counts across elastic baselines.

Mean C4 validation loss, in nats, is reported below. Lower is better.

Source: Section 4.2, Table 2. Bold identifies the lowest mean loss among the three jointly trained elastic methods, not the standalone family.

ShellElastic has the lowest mean validation loss among the evaluated elastic methods at all three scales. Its corresponding mean downstream 5-shot accuracies are 36.5%, 41.9%, and 44.0%, also exceeding the two evaluated elastic baselines at each scale.

The margins are relatively modest, particularly at the intermediate scale, where the mean-loss difference from the fully elastic model is 0.005 nats. These are measured quality differences rather than percentages of training computation saved.

The standalone comparison establishes another boundary. Five independently trained models achieve better average validation loss at every scale. They also collectively receive five times as many optimizer steps as a single elastic training run.

Thus, ShellElastic narrows the quality difference between jointly trained and independently trained families while reducing repeated optimization work. It does not establish a fivefold reduction in FLOPs or GPU-hours.

Training every width is more effective than sampling in the tested setting

The paper’s ablations help explain which recipe choices contribute to the result.

At the smallest scale, training all five supported widths at every step produces mean validation loss of 3.3090. Sandwich sampling, which trains the smallest, largest, and another selected width, produces 3.3206. Training Max plus a random width gives 3.3564, while randomly training one width per step gives 3.3681.

This is an ablation of width-selection strategy, not another independent demonstration of architectural superiority. It indicates that, for this configuration, consistent optimization of every tier improves average quality.

A separate architecture search favors allocating additional capacity to the dense core rather than placing the enlarged group in later shells. The wide-core designs are best or within 0.006 nats of the best at shared evaluation widths. That result informs capacity allocation under the tested parameter and block-count constraints.

Other experiments reveal less favorable additions. In-place distillation can improve narrow-tier validation loss, but increases step latency by approximately 2.2–2.6 times and produces little consistent downstream accuracy improvement. Tested non-oracle width ensembles do not outperform the widest model as predictors, despite a target-aware oracle showing potential complementary information.

Those experiments concern auxiliary training objectives and exploratory routing strategies. They do not establish additional deployable savings.

The remaining cost question belongs to the training platform

For a model-platform operator, theoretical trunk reuse must translate into measured execution efficiency.

Appendix E.4 reports favorable per-tier latency on A100-40GB GPUs. ShellElastic averages 601 ms in per-tier training latency versus 730 ms for MatFormer and 761 ms for the fully elastic model. Corresponding average inference latencies are 182 ms, 202 ms, and 209 ms.

However, these are per-tier measurements, not measurements of five independently trained models against one complete all-width training run. The authors also attribute part of ShellElastic’s latency and memory performance to fused Triton kernels. The results therefore reflect both the architecture and its implementation.

Several further boundaries matter. The smallest and intermediate scales average three runs using the same seed, while the largest-scale comparison uses one run. The maximal-update result relies on heuristic scaling assumptions. Training is restricted to one pretraining corpus and model sizes below contemporary frontier-scale systems. The reported comparisons also combine architectural and optimization-recipe differences.

Cognaptus inference: A foundation-model platform offering a fixed set of capacity tiers has a credible reason to test ShellElastic when those tiers align with supported shell boundaries. Its fixed shared-trunk computation could simplify multi-width training schedules, while proxy-derived optimizer settings may reduce repeated tuning.

The adoption case should compare aggregate GPU-hours at matched per-tier quality, training-memory requirements, extraction and serving costs, and latency under the platform’s actual traffic distribution. It should also account for the engineering cost of maintaining the required prefix-preserving operators and optimized kernels.

The paper supplies components of that assessment, but not a complete end-to-end cost comparison.

A more disciplined way to train a model family

ShellElastic’s principal contribution is the combination of two properties that parameter sharing alone does not guarantee: exact computation reuse across supported widths and appropriately scaled parameter updates across those widths.

Its experiments show that the combination can improve a jointly trained model family’s average quality over evaluated elastic alternatives, while independently trained models retain a measurable quality advantage at a larger aggregate optimizer-step budget.

For organizations maintaining multiple inference-capacity offerings, the resulting choice is more specific than whether elastic models are generally efficient. It is whether a defined portfolio of supported widths can meet its quality and serving requirements through one shared training process at a lower verified total cost.

That is where ShellElastic moves the engineering question—and where further evidence at the intended deployment scale is still needed.

Cognaptus: Automate the Present, Incubate the Future.