SPARKLING: Balancing Signal Preservation and Symmetry Breaking for Width-Progressive Learning

AI summary

Overview

Research area: Efficient large language model pre-training — specifically Progressive Learning (PL), where a model is trained at small scale and then grown to a larger target scale during training rather than trained from scratch.

Technical level: Advanced. The paper combines activation-statistics analysis (RMS scale propagation through residual streams), gradient-symmetry derivations, optimizer-internals arguments (AdamW momentum and Muon's Newton–Schulz orthogonalization), and large-scale MoE pre-training experiments.

Scope: The paper proposes and empirically validates a framework for performing width expansion in the mid-stage of pre-training (after roughly half the token budget has already been spent), balancing two principles it names Signal Preservation and Symmetry Breaking.

What This Paper Is About

Progressive Learning can cut pre-training cost by starting small and growing the model, but prior work has focused almost entirely on growing depth, and the limited work on growing width only expands very early in training (less than 10–30% of training tokens), which saves almost nothing. The authors target width expansion at the mid-stage, where the savings are large but training becomes unstable: naive initialization disrupts activation statistics and causes loss spikes, while copy-based initialization (which perfectly preserves the forward function) creates gradient symmetry that leaves the new capacity functionally redundant. SPARKLING is their recipe for expanding width mid-training while keeping hidden-state statistics stable and forcing the duplicated new channels to actually diversify.

Key Contributions

1. Reframes width expansion around two principles. The authors identify that successful mid-stage width expansion hinges on (i) Signal Preservation, framed not as loss continuity but as maintaining the RMS scale of intermediate activations, and (ii) Symmetry Breaking, needed to resolve the gradient coupling that copy-based initialization induces.

2. Derives RMS-preserving scaling rules for expansion. They work out scaling factors for fan-out expansion, fan-in expansion (separately for random/one-side-copied, both-sides-copied, and one-side-zero cases), and RMSNorm weight expansion, so that any of the standard initialization regimes (copy, random, zero) can be applied without an activation-scale shock.

3. Diagnoses and breaks "symmetry lock." They formally show that duplicated channels receive identical gradients under both AdamW and Muon, and that Muon's Newton–Schulz orthogonalization cannot break this symmetry because it is a polynomial map of the Gram matrix that preserves the block-duplication structure. They then introduce two asymmetric interventions acting only on backward dynamics: optimizer state reset for new parameters only, and an asymmetric learning rate re-warmup.

4. Validates generality across architectures, width axes, and optimizers. Experiments cover dense and Mixture-of-Experts (MoE) models, width axes including hidden dimension and MoE expert intermediate dimension, and optimizer families including AdamW and Muon.

Main Findings

• Naive (unscaled) initialization spikes then converges worse. In the OLMoE expert-inner expansion sweep (512 → 1024 at 100B tokens), "Naive Init, No Scaled" gives a smaller immediate loss gap, but "RMS-Preserved Scaled" overtakes it later and converges to a lower final loss, across all five up_proj–down_proj initialization pairs.

• The immediate loss spike does not predict final convergence. Figure 1(f) shows the magnitude of the post-expansion loss spike is not predictive of final convergence.

• Both-sides copied underperforms other RMS-preserved variants. Even with RMS-preserving scaling, the "both-sides copied" configuration significantly underperforms other RMS-preserved strategies — which the authors attribute to backward symmetry rather than forward continuity.

• Symmetric optimizer-state handling locks in the symmetry. Comparing treatments under copy-copy initialization with RMS-preserving scaling: both "Drop Opt." (globally reset all states) and "Copy Opt." (duplicate states) underperform, showing slower recovery and higher converged loss. "Asymmetric Reset" (reset states only for new channels) improves both recovery speed and final loss.

• Scaling the optimizer states adds nothing. "Asymmetric Reset + Scaled Opt.," which rescales optimizer states to match the parameter scaling, yields no additional gain, suggesting strict state–parameter scaling alignment is not critical; any initial mis-scaling is quickly corrected by later gradient updates.

• Asymmetric re-warmup helps consistently. Across Inner 2×, Hidden 2×, and joint expansion, and across both copy-copy and zero-copy initialization, enabling re-warmup lowers eventual loss under the same token budget. The gain is largest for copy-copy, which then reaches the lowest final loss among variants — the authors attribute this to re-warmup injecting controlled optimization asymmetry that lets duplicated subspaces diversify.

• Downstream results match or beat the from-scratch expanded baseline. SPARKLING achieves the best average accuracy in all three 2× settings despite a slightly higher final pre-training loss than the from-scratch baseline:

  • Inner 2×: Baseline (expand) avg 59.64, Naive FP scaled 58.96, SPARKLING 60.30 (loss 2.3153 vs. baseline 2.3096).
  • Hidden 2×: Baseline (expand) avg 60.46, Naive FP scaled 59.25, SPARKLING 61.13 (loss 2.2933 vs. baseline 2.2795).
  • Hidden 2× & Inner 2×: Baseline (expand) avg 60.89, Naive FP scaled 60.89, SPARKLING 62.55 (loss 2.2415 vs. baseline 2.2225).

• Compute savings up to 35%. Under a fixed 200B-token budget with expansion at 100B tokens: Inner 2× saves 20% FLOPs (7.21 vs. 9.01 ×10²⁰ for from scratch; 66 vs. 84 wall-clock hours; 1.27× speed-up), Hidden 2× saves 25% (8.10 vs. 10.80 ×10²⁰; 75 vs. 96 hours; 1.29×), and joint expansion saves 35% (11.70 vs. 18.00 ×10²⁰; 140 vs. 209 hours; 1.49×). The baseline small-model run used 5.40 ×10²⁰ FLOPs and 48 wall-clock hours at 450M/2.56B activated/total parameters.

• RMSNorm expansion is insensitive to copy vs. random. An ablation (Appendix C) shows copy and random give nearly identical final loss because the RMSNorm parameter count is negligible relative to the linear layers, so the authors use copy initialization for RMSNorm weights in all hidden-dimension experiments.

Methodology in Plain English

The authors train an OLMoE model with 0.5B active and 2.5B total parameters for 200B tokens with AdamW, expand the width at the 100B-token mark, and then continue training the expanded model for the remaining 100B tokens under the same recipe, keeping the original OLMoE pre-norm configuration for a controlled comparison.

Preservation side. They argue that what really matters at expansion is that the ratio of output RMS to input RMS in each layer stays the same. If that ratio shifts, the balance between the residual stream and the branch output changes, and downstream blocks see out-of-distribution inputs. For fan-out expansion (widening the output dimension), copying or same-distribution random sampling of the new rows keeps this ratio naturally. For fan-in expansion (widening the input dimension), the weights must be rescaled by √(d_in / d′_in) in the random or one-side-copied case; in the both-sides-copied case the correlation between duplicated coordinates requires a different factor — 1/√(1+3c) when 0 < c ≤ 1 and 1/(1+c) when c > 1, equivalently √(d_in/(3d′_in − 2d_in)) for d_in < d′_in ≤ 2d_in and d_in/d′_in for d′_in > 2d_in. For RMSNorm, preserving the output RMS reduces to preserving the RMS of the gain parameter, so copying or sampling the new coordinates works without rescaling. The authors also note that a zero-initialized side should be treated as random rather than strictly loss-preserving at the expansion moment.

Symmetry side. Copy expansion makes the two duplicated weight blocks receive identical gradients, so with identical optimizer states they stay identical forever — the model remains in the lower-dimensional subspace it started in. To break this without discarding the original model's training signal, they keep the optimizer state for the original weights and zero the state for the new weights (asymmetric reset), and they give the new parameters their own short learning-rate re-warmup: starting at the current learning rate η_e, warming over τ_w steps to a peak of ρ·η_e, then following the shared cosine decay. They use ρ = 1.3 and τ_w = 250 steps, selected based on Appendix H.

Evaluation. They compare the small pre-expansion model, the from-scratch target-width model under the same token budget, a naive function-preserving scaled variant, and SPARKLING, reporting final pre-training loss plus accuracy on ARC-C, ARC-E, Arithmetic, BoolQ, CommonsenseQA, HellaSwag, MMLU, OpenBookQA, PIQA, SciQ, SocialIQA, and Winogrande. They also tabulate FLOPs, wall-clock hours, activated/total parameters, tokens, and expansion point for the cost comparison.

Why This Matters

Impact on research. The paper supplies what it argues was missing: a systematic mechanism for width expansion, rather than the scattering of initializations, uneven splitting, symmetric perturbations, and growth operators in prior work. It also proposes a unifying explanation — that widely used function-preserving initializations work because they are, at bottom, RMS-preserving — and shows that Muon's orthogonalization offers no escape from copy-induced symmetry. That reframing gives other researchers a clean target: preserve RMS on the way in, inject asymmetry on the way out.

Real-world applications:

• Cheaper LLM pre-training runs. A 20–35% FLOP and wall-clock reduction at a fixed 200B-token budget makes large-scale pre-training campaigns meaningfully cheaper.
• Mid-flight upgrades of existing models. Because expansion happens at 100B of 200B tokens, teams can start a training run at smaller width and decide to widen later, rather than committing to the target width on day one.
• MoE architecture search. The expert-inner-dimension axis (512 → 1024) is expanded independently of the hidden dimension, allowing the expert intermediate size to be tuned during training.
• Optimizer-agnostic adoption. Generality across AdamW and Muon means the recipe is not tied to one training stack.

Industry relevance. The work was done at ByteDance Seed with academic affiliation at Peking University, and the reported savings are measured in wall-clock hours on real pre-training runs (48, 84, 96, 209 hours for the compared configurations), which is the currency that matters for large-scale training budgets.

Future Directions

• What happens beyond 2× expansion, and beyond the joint Hidden 2× & Inner 2× setting? All reported experiments use 2× growth; the paper does not report larger ratios.
• Can the two asymmetric interventions be unified or further reduced? The authors found that optimizer state scaling adds nothing and RMSNorm copy vs. random is irrelevant, which suggests the minimal sufficient intervention set may be even smaller.
• How do the re-warmup hyperparameters generalize? ρ = 1.3 and τ_w = 250 were chosen for these runs; their behavior at other scales, token budgets, or expansion points is not established here.
• Why does the from-scratch baseline still achieve lower final pre-training loss while SPARKLING wins on average downstream accuracy in all three settings? The paper notes this gap but the provided content does not resolve it.

Note: the supplied paper content is truncated in Section 5.2, so any results reported after Table 1 in the discussion are not covered in this summary.

Target Audience

Researchers and engineers working on efficient pre-training and scaling laws, particularly those implementing progressive or growable architectures at industrial scale. It is most useful to practitioners who already understand Transformer residual streams, RMSNorm, and optimizer internals, since the paper's core arguments depend on gradient-symmetry and activation-statistics reasoning. Readers looking for an introductory treatment of model growth will find the derivations demanding, while those evaluating whether to replace a from-scratch pre-training baseline will find the FLOP and wall-clock tables directly actionable.

Authors’ abstract

Progressive Learning (PL) reduces pre-training computational overhead by gradually increasing model scale. While prior work has extensively explored depth expansion, width expansion remains significantly understudied, with the few existing methods limited to the early stages of training. However, expanding width during the mid-stage is essential for maximizing computational savings, yet it remains a formidable challenge due to severe training instabilities. Empirically, we show that naive initialization at this stage disrupts activation statistics, triggering loss spikes, while copy-based initialization introduces gradient symmetry that hinders feature diversity. To address these issues, we propose SPARKLING (balancing {S}ignal {P}reservation {A}nd symmet{R}y brea{K}ing for width-progressive {L}earn{ING}), a novel framework for mid-stage width expansion. Our method achieves signal preservation via RMS-scale consistency, stabilizing activation statistics during expansion. Symmetry breaking is ensured through asymmetric optimizer state reset and asymmetric learning rate re-warmup. Extensive experiments on dense and Mixture-of-Experts (MoE) models demonstrate that, across multiple width axes and optimizer families, SPARKLING consistently outperforms training from scratch and reduces training cost by up to 35% under $2\times$ width expansion.