{
 "title": "ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer",
 "authors": [
  {
   "name": "Can Jin",
   "affiliations": [
    1
   ],
   "marks": "*"
  },
  {
   "name": "Yan Kang",
   "affiliations": [
    2
   ]
  },
  {
   "name": "Yifan Gong",
   "affiliations": [
    2
   ]
  },
  {
   "name": "Chongjian Ge",
   "affiliations": [
    2
   ]
  },
  {
   "name": "Mingze Xu",
   "affiliations": [
    2
   ]
  },
  {
   "name": "Yuanjun Xiong",
   "affiliations": [
    2
   ]
  },
  {
   "name": "Dimitris N. Metaxas",
   "affiliations": [
    1
   ]
  },
  {
   "name": "Hongwu Peng",
   "affiliations": [
    2
   ]
  }
 ],
 "affiliations": [
  "Rutgers University",
  "Adobe Research"
 ],
 "footnotes": [
  "* Corresponding author: can.jin@rutgers.edu. Work done during an internship at Adobe Research."
 ],
 "links": [
  {
   "label": "Paper",
   "message": "PDF",
   "color": "#b31b1b",
   "href": "https://arxiv.org/pdf/2610.15124"
  },
  {
   "label": "Project",
   "message": "Page",
   "color": "#2f6db5",
   "href": "./"
  },
  {
   "label": "Code",
   "message": "GitHub",
   "color": "#181717",
   "logo": "github",
   "href": "https://github.com/jincan333/ShellElastic"
  },
  {
   "label": "Checkpoints",
   "message": "Hugging Face",
   "color": "#ffd21e",
   "logo": "huggingface",
   "href": "https://huggingface.co/spaces/Can111/ShellElastic"
  },
  {
   "label": "Datasets",
   "message": "Hugging Face",
   "color": "#ffd21e",
   "logo": "huggingface",
   "href": "https://huggingface.co/spaces/Can111/ShellElastic"
  },
  {
   "label": "License",
   "message": "Apache-2.0",
   "color": "#3da639",
   "logo": "apache",
   "href": "https://github.com/jincan333/ShellElastic/blob/main/LICENSE"
  }
 ],
 "method": [
  {
   "type": "h",
   "id": "architecture",
   "where": "§3.1 · App. A–B",
   "text": "The ShellElastic architecture"
  },
  {
   "type": "p",
   "text": "ShellElastic (Figure 1) partitions the residual stream into a dense core followed by ordered shells, each a multiple of a base width g, and each supported sub-network is the core plus a prefix of the shells."
  },
  {
   "type": "p",
   "text": "Every hidden matrix is block lower triangular: a shell reads the core, the earlier shells and itself, but never a later shell. Trunk normalization acts within each group, attention heads align with the group boundaries, and every supported width has its own readout."
  },
  {
   "type": "figure",
   "key": "figure1",
   "src": "figures/architecture.png",
   "ref": "Figure 1"
  },
  {
   "type": "p",
   "text": "Group-local normalization uses a separate operator N<sub>k</sub> and statistics within each group (Figure 1b). Both operations preserve prefixes ending at the boundaries in 𝒮. Cuts within a group generally change the retained outputs, so sub-networks can be extracted only at these boundaries. Attention is elastic in its number of heads: sub-network s uses the first sg/d<sub>head</sub> heads, keeping the head dimension d<sub>head</sub> fixed."
  },
  {
   "type": "explainer",
   "id": "structure"
  },
  {
   "type": "p",
   "text": "One full-width forward pass therefore contains the exact hidden states of every sub-network as prefixes (Proposition 1), and the per-step compute stays close to that of a single full-width pass however many sub-networks are trained, so every step has a fixed workload, as in non-elastic training."
  },
  {
   "type": "p",
   "style": "statement",
   "id": "proposition1",
   "where": "App. B",
   "text": "<b>Proposition 1 (Exact reuse at supported widths).</b> For fixed model parameters, let E<sub>s</sub> be the first sg columns of the full-width embedding table, T<sub>s</sub> the L-layer trunk restricted to width sg, and P<sub>s</sub> the operator selecting the first sg hidden dimensions. For any token sequence u, evaluating sub-network s directly gives the same hidden states as extracting its prefix from sub-network t: T<sub>s</sub>(E<sub>s</sub>[u]) = P<sub>s</sub>(T<sub>t</sub>(E<sub>t</sub>[u])), s, t ∈ 𝒮, s ≤ t."
  },
  {
   "type": "p",
   "text": "Applying sub-network s’s normalization and readout to either result gives identical logits. Thus, one full-width trunk evaluation yields the same joint loss and parameter gradients as separate evaluations of all supported sub-networks."
  },
  {
   "type": "h",
   "id": "parameterization",
   "where": "§3.2 · App. C–D",
   "text": "μ-Elastic: parameterizing elastic models"
  },
  {
   "type": "p",
   "text": "μ-Elastic (Figure 2) extends maximal update parameterization (μP) to elastic models with two factors: a μ-Transfer factor m = g/d₀ that carries HPs from a non-elastic proxy of width d₀ to base width g, and a μ-Internal factor for the different widths inside one shared model."
  },
  {
   "type": "figure",
   "key": "figure2",
   "src": "figures/parameterization-overview.png",
   "ref": "Figure 2"
  },
  {
   "type": "p",
   "text": "Let a non-elastic proxy of width d₀ supply the base initialization scale σ₀, base learning rate η₀, and weight decay λ₀. For base width g, define m = g/d₀ and q(s) = m · s = sg/d₀ for s ∈ 𝒮. For a hidden weight, s is the supported boundary of its output group, whose input width stays sg (esg for an FFN down projection) in every sub-network that contains it, so one rule based on q(s) serves all of them. The readout of sub-network s uses the same ratio."
  },
  {
   "type": "p",
   "text": "The standard parameterization (SP) sets m = s = e = 1 only in the scaling rules and fixes α<sub>s</sub> = 1, preserving the architecture."
  },
  {
   "type": "explainer",
   "id": "parameterization"
  },
  {
   "type": "p",
   "text": "Every group reads the same input coordinates in every sub-network that contains it, so we set its initialization, learning rate and weight decay from this fan-in, while each readout uses its own input width (Table 1). This keeps every group maximally updated in all sub-networks at once (Theorem 1), which uniform scaling across groups cannot do in general (Proposition 2)."
  },
  {
   "type": "figure",
   "key": "table1",
   "src": "figures/parameterization-rules.png",
   "ref": "Table 1"
  },
  {
   "type": "p",
   "text": "A parameter group is maximally updated when the pre-activations it produces are Θ(1) at initialization and move by Θ(η₀) after one optimizer step, with η₀ the base learning rate before width scaling, so that its contribution neither vanishes nor blows up as the width grows."
  },
  {
   "type": "p",
   "text": "<b>Theoretical guarantees.</b> The statements below are proved in the paper’s Appendix C under the standard heuristic assumptions of μP: before the scaling of Table 1 is applied, the pre-activations of a group with active input width sg are Θ(√(sg)) at initialization and move by Θ(sg) after one step, which is why the initialization and the learning rate carry different exponents."
  },
  {
   "type": "p",
   "style": "statement",
   "id": "theorem1",
   "where": "App. C.2",
   "text": "<b>Theorem 1 (μ-Internal: simultaneous maximal updates).</b> Under the assumptions of the paper’s Appendix C and the rule of Table 1, consider a hidden group with supported boundary s and any sub-network t ∈ 𝒮 with t ≥ s. In every layer, its pre-activations inside sub-network t (i) have Θ(1) coordinates at initialization and (ii) move by Θ(η₀) after one AdamW step on that group, with constants independent of s, t, M, and m. The logits of every readout U<sub>s</sub> have the same two properties."
  },
  {
   "type": "p",
   "text": "The ratio q(s) is what makes this possible: a group with boundary s sums over exactly sg input coordinates in every sub-network containing it, so one initialization scale and one learning rate satisfy both conditions for all t ≥ s at once. Training several sub-networks jointly does not change this, since they all feed a group the same input and Adam is invariant to the gradient scale."
  },
  {
   "type": "p",
   "style": "statement",
   "id": "proposition2",
   "where": "App. C.3",
   "text": "<b>Proposition 2 (A shared rule leaves groups mismatched).</b> If all groups of a hidden matrix share one initialization variance, learning rate and multiplier, the groups with boundaries s &lt; s′ differ by factors Θ(√(s′/s)) in pre-activation size and Θ(s′/s) in update size, whatever the shared values. For unbounded M/s₁, no shared rule thus keeps every group maximal in the sense of Theorem 1."
  },
  {
   "type": "p",
   "text": "Under SP, the groups of a hidden matrix share (σ₀, η₀), so by Proposition 2 re-tuning only trades the core against the outermost group, which stay √(M/s₁) and M/s₁ apart in initialization and update; the best HPs are compromises depending on m, M and 𝒮. Figure 3 and Table 3 show both effects: the SP optimum departs from the proxy’s, and SP trails μ-Elastic even after re-tuning."
  },
  {
   "type": "p",
   "style": "statement",
   "id": "corollary1",
   "where": "App. C.4",
   "text": "<b>Corollary 1 (μ-Transfer: two-level transfer).</b> Theorem 1 holds for every base width g = md₀ and every supported set 𝒮 ⊆ {1, …, M}. Hence the values (σ₀, η₀, λ₀) tuned on the width-d₀ proxy (i) keep every group and readout of a ShellElastic supernet maximally updated, and (ii) keep doing so when the supernet is widened (larger m) or given more shells (larger M)."
  },
  {
   "type": "p",
   "text": "As with μP, the corollary fixes the scales of the initialization and the updates, not the optimal HP values; Figure 3 verifies that the proxy’s tuned HP values remain near-optimal at both levels."
  },
  {
   "type": "p",
   "text": "<b>μ-Elastic for other elastic architectures.</b> μ-Elastic carries over to other elastic architectures, but m and s are defined per architecture. MatFormer keeps the model width fixed and nests only the FFN hidden width, so only its FFN down projection has a fan-in that varies with the sub-network. It uses m = D<sub>model</sub>/d₀, with s the sub-network’s FFN hidden width in units of a base FFN width; the down projection takes the factor ms, and the other hidden matrices and the readout take m (the paper’s Table 6). Fully elastic models, a transformer adaptation of slimmable networks, slice shared dense matrices along the residual, attention and FFN widths, so every hidden matrix has a varying fan-in; their m and s follow the definitions above, and every hidden matrix takes the factor ms (the paper’s Table 5). StairFormer, with equal-width groups, takes Table 1 with 𝒮 = {1, …, M} in its hidden matrices. Theorem 1 and Corollary 1 extend to all three (the paper’s Appendix C.5)."
  },
  {
   "type": "h",
   "id": "training",
   "where": "§3.3 · App. E",
   "text": "Training and inference"
  },
  {
   "type": "p",
   "text": "<b>Recipe.</b> (1) Tune (σ₀, η₀, λ₀) on a non-elastic proxy of width d₀. (2) Choose the base width g and the supported set 𝒮. (3) Compute m = g/d₀ and instantiate the parameterization in Table 1. (4) Train the supernet on the joint objective below."
  },
  {
   "type": "p",
   "text": "<b>Joint objective.</b> At step τ, with selected widths 𝒜<sub>τ</sub> ⊆ 𝒮 and loss weights ω<sub>s</sub> ≥ 0,"
  },
  {
   "type": "p",
   "style": "equation",
   "text": "ℒ<sub>τ</sub> = (Σ<sub>s∈𝒜<sub>τ</sub></sub> ω<sub>s</sub>ℓ<sub>s</sub>) / (Σ<sub>s∈𝒜<sub>τ</sub></sub> ω<sub>s</sub>),  ℓ<sub>s</sub> = CE(z<sub>s</sub>, y) + λ<sub>z</sub> 𝔼[(log Σ<sub>v</sub> e<sup>z<sub>s,v</sub></sup>)²],"
  },
  {
   "type": "p",
   "text": "where z<sub>s</sub> are the logits of sub-network s, λ<sub>z</sub> is the z-loss coefficient, the expectation runs over the non-padded target positions, and every ℓ<sub>s</sub> backpropagates into the shared trunk. The narrower widths may be selected by any rule, but every 𝒜<sub>τ</sub> contains the full width M, so the trunk always runs at width D<sub>model</sub> and every trunk parameter receives a gradient at every step."
  },
  {
   "type": "p",
   "text": "<b>Cost and scalability.</b> Let C<sub>trunk</sub>(s) and C<sub>out</sub>(s) be the forward costs of the trunk and of one terminal norm and readout at width sg. Training the family jointly and training each sub-network separately cost"
  },
  {
   "type": "p",
   "style": "equation",
   "text": "C<sub>joint</sub> = C<sub>trunk</sub>(M) + Σ<sub>s∈𝒮</sub> C<sub>out</sub>(s),  C<sub>separate</sub> = Σ<sub>s∈𝒮</sub> [C<sub>trunk</sub>(s) + C<sub>out</sub>(s)]."
  },
  {
   "type": "p",
   "text": "By Proposition 1, one trunk term serves every width, whereas C<sub>separate</sub> pays one per width, so the gap grows with |𝒮| and C<sub>joint</sub> stays comparable to training the widest sub-network alone. The same holds for memory: all readout gradients combine in one trunk backward pass, so the trunk’s activations, parameters and optimizer state are those of the widest one. Training all sub-networks at every step works best (Table 4 and the paper’s Appendix J) and gives every step the same computation graph, parameter access pattern and workload. In MatFormer and fully elastic models, hidden states are not nested, so each width needs its own trunk pass: sampling widths per step varies the workload, while training all widths at every step pays C<sub>separate</sub>."
  },
  {
   "type": "p",
   "text": "<b>Inference.</b> For a chosen width s ∈ 𝒮, we retain the embedding prefix, the core and shells up to s, and the matching terminal norm and readout, keeping the forward multipliers of Table 1. By Proposition 1, the extracted model reproduces its supernet logits exactly and runs as a standalone transformer of width sg without further training."
  }
 ],
 "results": {
  "main": {
   "where": "§4.2 · App. G",
   "text": [
    "Table 2 compares ShellElastic with MatFormer, the fully elastic model and the two standalone baselines at three scales; all three elastic models use μ-Elastic. ShellElastic is the supernet of the paper’s Section 4.1 at 119M–285M, doubles its base width to g = 256 at 375M–906M, and also raises the FFN expansion from e = 1 to 4 at 675M–1.69B; these two scales train for 50,000 and 100,000 steps of 256 sequences.",
    "Among the methods that train all five sizes in one run, ShellElastic has the lowest mean validation loss (0.028, 0.005 and 0.018 nats below the best baseline) and the highest mean 5-shot accuracy at every scale. It is best or tied-best at 12 of the 15 tiers in validation loss and 14 in 5-shot accuracy. Its lead grows with size, from at most 0.016 nats at Low to 0.023–0.050 at Max, likely because the Max-only parameters of the elastic baselines receive a gradient on only one step in five on average, whereas every parameter of ShellElastic receives one at every step.",
    "At a fixed step budget, weight sharing also helps: the fully elastic model beats sampled standalone, its unshared counterpart, by 0.11 to 0.25 nats on average. With a fifth of their optimizer steps, ShellElastic comes within 0.039, 0.045 and 0.014 nats of the standalone models in mean validation loss and within 0.8 points in mean 5-shot accuracy. The paper’s Appendix E.4 analyzes the training and inference cost of all methods at 119M–285M."
   ]
  },
  "transfer": {
   "where": "§4.2 · App. H",
   "text": [
    "Under μ-Elastic, the proxy’s optimum, η₀ = 2 × 10⁻³ and σ₀ = 3.3 × 10⁻³ (the base values of the paper’s Section 4.1), transfers to all three elastic models at both levels of μ-Transfer (Figure 3, Corollary 1). At the 119M–285M scale, both values are optimal or adjacent to the optimum (η₀: exact for MatFormer and the fully elastic model, adjacent for ShellElastic; σ₀: exact for the fully elastic model and ShellElastic, adjacent for MatFormer). At the 375M–906M scale, which enlarges m, M or both and doubles the batch size and steps, the same holds, and σ₀ becomes exact for MatFormer too.",
    "The base values are thus within 0.003 nats of the best point of every μ-Elastic sweep and within 0.005 nats for every sub-network (Figures 12 to 14). Under SP, the optimum departs from the proxy’s and shifts with scale, as the paper’s Section 3.2 anticipates. At the base values SP is 0.03–0.11 nats worse. Weight decay transfers as well (Figure 15): under μ-Elastic, the mean loss of ShellElastic is flat over the same range as the proxy’s, from λ₀ = 10⁻³ to 10⁻¹ (within 0.003 nats), whereas SP is higher at every swept value."
   ]
  },
  "maximal-updates": {
   "where": "§4.2 · App. H",
   "text": [
    "Without width-dependent scaling, re-tuning one global value only trades one width against another (Proposition 2). Table 3 tests this limitation at the 119M–285M and 375M–906M scales, giving SP the best of its 12 (13 at 375M–906M) configurations in Figure 3 and μ-Elastic only the base values.",
    "At the 119M–285M scale, for the fully elastic model and ShellElastic, μ-Elastic still gives lower training and validation losses at every tier, by 0.011 and 0.013 nats in mean validation loss, as it keeps all widths maximally updated at once (Theorem 1). At the 375M–906M scale, tuned SP remains 0.006–0.021 nats worse than μ-Elastic for all three models (the paper’s Appendix H). Under tuned SP, ShellElastic still beats both baselines by 0.026–0.027 and 0.020–0.024 nats at the two scales, so its lead is not due to μ-Elastic alone."
   ]
  },
  "arch": {
   "where": "§4.3 · App. I",
   "text": [
    "<b>Comparison with StairFormer.</b> StairFormer nests sub-networks as prefixes of one forward pass; it uses equal-width groups, cumulative prefix normalization and one shared readout. Here it has g = 208 and 𝒮 = {1, 2, 3, 4, 5} in our notation, whereas ShellElastic has a 512-wide core and four 128-wide shells (𝒮 = {4, 5, 6, 7, 8}), group-local normalization and per-width readouts. Like ShellElastic, it trains all sub-networks at every step; we give it the same μ-Elastic and HPs. Its narrower sub-networks are much smaller (32.6M to 200.4M); only its Max matches ShellElastic in size. At the parameter count of each ShellElastic sub-network, its validation loss is 0.055 to 0.083 nats higher (Figure 4a; curves in the paper’s Figure 17), and even its Max, with 1.5× the parameters, is 0.001 nats above the High of ShellElastic. At Max, ShellElastic also leads in average 5-shot accuracy, 36.9 against 35.7 (Figure 4b and the paper’s Table 29).",
    "<b>Architecture search at equal parameters and width count.</b> Since 𝒮 fixes both the extractable widths and the parameter count (a hidden matrix keeps B(𝒮) = Σ<sub>k</sub>(s<sub>k</sub> − s<sub>k−1</sub>)s<sub>k</sub> of its M² blocks of size g × g), we ask where extra width should go when B and the number of supported widths are fixed. At B = 42 with five supported widths, each candidate set has four groups of width g and one of width 4g and differs only in where this wide group sits (block masks in the paper’s Figure 18). We train every supported width at every step and compare the sets at each width they share with equal parameter counts (Figure 4c). The core is the best place for the wide group: the core placement has the lowest loss, or is within 0.006 nats of it, at every shared width, whereas a shell placement makes the first width containing the wide group 0.025 to 0.052 nats worse. This motivates the wide core and thin shells of ShellElastic."
   ]
  },
  "ablations": {
   "where": "§4.4 · App. J",
   "text": [
    "All ablations use the 119M–285M supernet of the paper’s Section 4.1 and change one setting at a time; the default results are those of Table 2. Of the settings we ablate, 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. We discuss the first four below (Table 4) and the others in the paper’s Appendix J.",
    "<b>Selected sub-networks per step.</b> Besides training all five sub-networks at every step (All, the default), we test three rules for 𝒜<sub>τ</sub> in the joint objective: Sandwich (Low and Max plus one other), Max + random (Max plus one other) and Random (one sub-network, as in MatFormer), drawing the random ones uniformly at random. The loss of each sub-network follows how often it is trained: Sandwich and Max + random match All at Max but are up to 0.024 and 0.086 nats worse at those they select less often, and Random is worse at every sub-network, by 0.095 nats at Max. All gives the lowest mean loss, 0.012 to 0.059 nats below the others, and the highest mean 5-shot accuracy.",
    "<b>In-place distillation.</b> Elastic recipes often distill the narrower sub-networks from the widest one. We add to the losses of Low to Extra a KL divergence toward the output distribution of Max (forward KL with weights 0.1 to 1; reverse and symmetric KL with weight 0.5). This 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, but raises that 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.",
    "<b>Stage-residual gradient scaling.</b> Since the readout is linear, the logits of width s<sub>k</sub> are a sum of k terms, one from the last-layer coordinates of each of its groups, so the supernet can be viewed as a sequence of stages: the core gives a base prediction and each shell adds a correction, as in residual learning. To make each shell learn this correction, we keep the forward pass but, in the backward pass, scale by γ the gradient that each sub-network’s loss sends to the trunk through the narrower sub-network it contains; the gradient through its own shell is unchanged. Thus γ = 1 is the usual joint training, while with γ = 0 each loss reaches the trunk only through its own shell, which therefore learns only to correct the prediction of the narrower sub-network. 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 the mean loss by 0.047 nats, as the last-layer features of the narrower groups no longer adapt to the wider losses; we keep γ = 1. The stage view still points to a property that MatFormer and fully elastic models lack: since no group reads a later one and every group is normalized on its own, a new shell can be appended to a trained supernet and trained as a new stage with all trained weights frozen, so every trained sub-network keeps its exact outputs."
   ]
  },
  "ensembles": {
   "where": "§4.4 · App. J",
   "text": [
    "One trunk pass yields the output distributions p<sub>s</sub> of all five sub-networks, which may err on different tokens, so a per-token ensemble Σ<sub>s</sub> π<sub>s</sub>p<sub>s</sub> could beat Max as a predictor at little extra cost, or replace it as a stronger teacher (details in the paper’s Appendix J).",
    "<b>Ensemble as predictors.</b> We evaluate four ensembles of a trained supernet: uniform weights, confidence weights (π<sub>s</sub> ∝ max<sub>v</sub> p<sub>s</sub>(v)), the most confident readout per token, and the best fixed ensemble, whose weights minimize the loss. We also train a router, an MLP that maps the last-layer hidden state of Max to per-token weights π<sub>s</sub>, together with its own supernet on the π-weighted sum of the five losses. None of the ensemble methods beat Max alone (Figure 5 and the paper’s Table 33): the best fixed ensemble only matches it, the other three are 0.03 to 0.065 nats worse in train loss, and the router puts all weight on Max within the first 50 steps (Figure 5b), so Max improves by 0.013 nats while the other sub-networks, now barely trained, lose up to 0.365 nats (Figure 5c).",
    "<b>Oracle.</b> An oracle that picks per token the readout with the highest target probability is 0.19 nats below Max (the paper’s Table 34), as a narrower sub-network is better on 71% of tokens (Figure 5e). Two supernets trained separately with the same settings differ far more: averaging their Max lowers the loss by 0.056 nats at twice the cost, and their oracle gains 0.23 nats against 0.19 for the five sub-networks on held-out text. Our ensembles show that finding the better sub-network per token without the target is hard, but we believe a better ensemble could close part of this gap, a promising way to substantially improve ShellElastic at almost no extra cost.",
    "<b>Ensemble as teachers.</b> Every sub-network, including Max, distills (forward KL, weight 0.5) toward the loss-weighted ensemble (π<sub>s</sub> ∝ ω<sub>s</sub>), which predicts 0.012 nats worse than Max, or the soft oracle (π<sub>s</sub> ∝ p<sub>s</sub>(y)<sup>β</sup>), a smoothed oracle that predicts 0.14 nats better (the paper’s Table 33). Yet neither teacher helps Max, not even the soft oracle: like in-place distillation with the same KL (Figure 5d), they improve Low by up to 0.018 nats and worsen Max by up to 0.010 nats."
   ]
  }
 },
 "captions": {
  "figure1": "Core and shells of ShellElastic preserve supported prefixes. (a) Each colored block row produces output group 𝒢<sub><i>k</i></sub> from the input prefix ending at <i>s</i><sub><i>k</i></sub><i>g</i>. Group 𝒢<sub>1</sub> is the dense core; subsequent groups are shells. (b) <i>N</i><sub><i>k</i></sub> normalizes group 𝒢<sub><i>k</i></sub> using its own statistics.",
  "figure2": "Overview of μ-Elastic. The ratio <i>q</i> = <i>m</i> · <i>s</i> of the paper’s Equation 7 scales the HPs of every group and readout. (a) μ-Transfer: the (σ₀, η₀, λ₀) tuned on the proxy are reused (i) for an elastic supernet and (ii) for a wider one with more shells (the paper’s Corollary 1). (b) μ-Internal: block row <i>W</i><sub><i>k</i></sub> has fan-in <i>s</i><sub><i>k</i></sub><i>g</i> in every sub-network that contains it, so <i>q</i> = <i>ms</i><sub><i>k</i></sub> keeps it maximally updated in all of them (the paper’s Theorem 1). The bars show each group’s pre-activation size at initialization and update size: with the proxy’s (σ₀, η₀), as under SP, all groups, even the core, exceed the proxy by √<i>q</i> and <i>q</i>, so no shared value fixes them all (the paper’s Proposition 2); under μ-Elastic, all match it.",
  "figure3": "Sweeps of the learning rate η₀ and initialization scale σ₀, the other at its base value, for the proxy and each elastic model at both scales under μ-Elastic and SP. Δ val. loss is the mean validation loss of the five sub-networks minus the minimum of the μ-Elastic line of the same model and scale (for the proxy, its own line). The paper’s appendix shows the sweeps for each sub-network.",
  "figure4": "Architecture comparison and search against parameters at step 25,000. (a, b) Comparison with StairFormer: validation loss and average 5-shot accuracy (%) over 13 tasks. (c) Architecture search: validation loss of each supported width. Full results are in the paper’s appendix.",
  "figure5": "Ensembles of the sub-networks. (a) Train loss minus that of Max during training for the five predictors and the oracle, which uses the target. (b) Router weight on each tier. (c) Validation loss of each tier minus the default when training with the router. (d) Validation loss of each tier minus the default when distilling toward Max, the loss-weighted ensemble or the soft oracle. (e) Share of held-out tokens on which each tier assigns the target the highest probability. The default is the mean of the three ShellElastic runs of Table 2.",
  "table1": "μ-Elastic for ShellElastic. Scaling rules set by the μ-Transfer factor <i>m</i>, the μ-Internal factor <i>s</i> (the paper’s Equation 7), and the FFN expansion <i>e</i> = <i>D</i><sub>ffn</sub>/<i>D</i><sub>model</sub>; <i>s</i> 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.",
  "table2": "Final validation loss and mean 5-shot accuracy over 13 downstream tasks at three scales, labeled by sub-network size. Entries at the 119M–285M and 375M–906M scales average three runs; those at 675M–1.69B are single runs. †Five separately trained models, together taking 5× the steps of any other row. Bold marks the best of the four methods that train all five sizes in one run.",
  "table3": "Effect of μ-Elastic on each elastic model at the 119M–285M and 375M–906M scales. For each model, the first row is the best of the 12 (13 at 375M–906M) SP configurations of Figure 3 at that scale (lowest mean validation loss), and the w/ μ-Elastic row is its result in Table 2, with the base values of the paper’s Section 4.1 and no re-tuning. The paper’s appendix gives the values with four decimals.",
  "figure12": "<b>Per-sub-network sweeps for MatFormer.</b> The sweeps of Figure 3 for each sub-network, from the smallest (Low) to the largest (Max; widths in the paper’s Table 7), with the loss of each sub-network in place of the mean over the five.",
  "figure13": "<b>Per-sub-network sweeps for the fully elastic model</b>, as in Figure 12.",
  "figure14": "<b>Per-sub-network sweeps for ShellElastic</b>, as in Figure 12.",
  "figure15": "<b>Weight-decay sweeps.</b> Sweeps of λ₀ for the proxy and for ShellElastic at the 119M–285M scale under μ-Elastic (solid) and SP (dashed, open markers); Δ val. loss is as in Figure 3.",
  "figure16": "<b>Per-sub-network weight-decay sweeps.</b> The sweeps of Figure 15 for each sub-network of ShellElastic, from the smallest (Low) to the largest (Max); lines and Δ val. loss as there, with the loss of each sub-network in place of the mean over the five, and one y-range for all panels.",
  "table4": "Ablations at the 119M–285M 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. The paper’s appendix gives all ablations."
 },
 "figure_sizes": {
  "figures/arch.png": [
   2020,
   790
  ],
  "figures/architecture.png": [
   2020,
   1100
  ],
  "figures/ensemble.png": [
   2020,
   1492
  ],
  "figures/parameterization-overview.png": [
   2020,
   1872
  ],
  "figures/parameterization-rules.png": [
   2020,
   649
  ],
  "figures/transfer-sub-networks-fully-elastic.png": [
   2020,
   1636
  ],
  "figures/transfer-sub-networks-matformer.png": [
   2020,
   1745
  ],
  "figures/transfer-sub-networks-shellelastic.png": [
   2020,
   1636
  ],
  "figures/transfer-weight-decay-sub-networks.png": [
   2020,
   974
  ],
  "figures/transfer-weight-decay.png": [
   2020,
   1167
  ],
  "figures/transfer.png": [
   2020,
   1181
  ]
 },
 "bibtex": "@article{jin2026shellelastic,\n  title   = {{ShellElastic}: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer},\n  author  = {Jin, Can and Kang, Yan and Gong, Yifan and Ge, Chongjian and Xu, Mingze and Xiong, Yuanjun and Metaxas, Dimitris N. and Peng, Hongwu},\n  journal = {arXiv preprint arXiv:2610.15124},\n  year    = {2026},\n  url     = {https://arxiv.org/abs/2610.15124}\n}"
}
