PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Qifan Yu, Xinyu Ma, Zhijian Zhuo, Minrui Wang, Deyi Liu, Shiyi Zhan, Yiyuan Ma, Liang Xiang, Xingyan Bin, Di He

arXiv:2602.02472v2cs.LGcs.CL

TL;DR

Mid-stage width expansion is understudied despite its importance for reducing progressive-learning costs. SPARKLING combines RMS-scale preservation with asymmetric symmetry breaking, outperforming scratch training across settings while reducing costs by up to 35% for 2× width expansion.

  • Problem

    Width expansion in progressive learning lacks a general, systematic mechanism and has mostly been studied only during the early stages of training.

  • Method

    SPARKLING combines RMS-scale consistency with asymmetric optimizer-state reset and learning-rate re-warmup for mid-stage width expansion.

  • Results

    Across dense and MoE models, width axes, and optimizer families, SPARKLING consistently outperforms scratch training and reduces training costs by up to 35% at 2× width.

  • Takeaways & Limitations

    SPARKLING provides an effective and efficient framework for mid-stage width expansion across architectures, expansion stages, width axes, and optimizer families.

  • Takeaways & Limitations

    A unified principle for simultaneous width and depth expansion remains unestablished, and tuning-free hyperparameter transfer after expansion remains future work.

Abstract

from arXiv · show

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.

1. Introduction

SPARKLING addresses the underexplored challenge of mid-stage width expansion by combining RMS-scale consistency for signal preservation with targeted symmetry-breaking interventions. Across dense and MoE models, width axes, and optimizer families, it improves downstream performance over full-scale-from-scratch training while reducing costs by up to 35% under 2× expansion.

  • Motivation: Mid-stage width expansion remains underexplored, as prior width-based methods largely focus on the first 10–30% of training tokens.A general and systematic width-expansion mechanism has yet to be established.
  • Signal Preservation: RMS-scale mismatch changes layerwise signal magnitudes through residual streams, destabilizing optimization even without an immediate loss spike.This motivates treating signal preservation as activation-scale consistency rather than only loss continuity or function preservation.
  • Symmetry Breaking: Copy-based expansion preserves forward continuity but duplicates parameter subspaces, causing identical gradients and functionally redundant new capacity.The resulting backward symmetry occurs under both element-wise and non-element-wise optimizers.
  • SPARKLING: SPARKLING combines RMS-scale consistency with asymmetric optimizer state reset and asymmetric learning rate re-warmup to stabilize expansion and break gradient symmetry.Its targeted interventions act on backward dynamics while preserving the copied forward function.
  • Empirical Validation: 35% lower training cost is achieved under 2× width scaling while SPARKLING consistently outperforms full-scale-from-scratch training on downstream evaluations.The framework generalizes across dense and MoE architectures, hidden and expert-intermediate width axes, and AdamW and Muon.

2. Related Work

Progressive Learning expands architectures during training, with prior work emphasizing depth growth and stabilization of post-growth optimization dynamics. The paper identifies a limitation in treating forward initialization and backward optimization separately, motivating a framework that combines RMS preservation with symmetry-aware optimization.

  • Progressive Learning: Progressive Learning accelerates training by expanding a small base model toward the target scale during training.Prior depth-expansion strategies grow models by stacking layers or inserting blocks.
  • Post-growth optimization: Post-growth methods stabilize optimization through accelerated decay schedules, weight-norm-based gradient rebalancing, or dynamics-preserving growth operators.These approaches target optimization behavior after model expansion.
  • Limitations and framework: Existing strategies typically address forward initialization or backward optimization dynamics in isolation, whereas this work establishes a framework balancing both perspectives.The framework interprets function-preserving initialization as RMS preservation and redesigns optimization to address symmetry issues caused by preservation-focused initialization.

3. RMS Scale Consistency of Activation

RMS-scale consistency preserves activation statistics during width expansion, keeping residual dynamics within the trained operating regime across normalization architectures. The framework covers fan-out, fan-in, and RMSNorm expansion, with experiments showing improved late-stage convergence.

  • Signal preservation: RMS-preserving expansion keeps activation scales unchanged, preventing hidden states from drifting outside the pre-expansion operating regime.This preserves well-formed representations and meaningful features for subsequent blocks.
  • Residual dynamics: Under pre-normalization, RMS preservation maintains the calibrated residual mixing between the main stream and transformed branch.An RMS mismatch can make the branch overwhelming or leave the block nearly identity.
  • Residual dynamics: RMS-preserving expansion remains necessary for post-normalization because residual weighting still depends on the RMS ratio between the input and branch outputs.The constraint is therefore architecture-agnostic across pre-norm and post-norm variants.
  • Expansion cases: Fan-out expansion preserves activation RMS when added output channels are copied or sampled consistently with the original weight distribution.Fan-in expansion instead requires scaling, with distinct treatment for random or one-side copied and both-sides copied coordinates.
  • Empirical validation: RMS-preserving expansion robustly improves late-stage convergence across expert-inner and hidden-dimension growth and diverse initialization strategies.Naive Init, No Scaled has a smaller immediate loss gap but worse late-stage convergence, while RMS-Preserved Scaled reaches lower final loss.
  • Empirical validation: Both-sides copied initialization significantly underperforms other RMS-preserved variants, and immediate loss-spike magnitude does not predict final convergence.One-side zero initialization should treat the zero-initialized side as random rather than strictly loss preserving at expansion.

4. Breaking the Symmetry Lock

Copy-based width expansion can preserve forward outputs yet lock duplicated parameters into identical backward updates, leaving added capacity redundant. SPARKLING breaks this lock through asymmetric optimizer-state reset and learning-rate re-warmup, improving recovery and final loss across expansion settings.

  • Symmetry Lock: Copy initialization duplicates gradients, so expanded components evolve identically and fail to diversify into distinct features.This copy-induced backward symmetry explains slower post-expansion recovery and higher eventual loss despite forward-output preservation.
  • Symmetry Lock: Symmetrically initialized optimizer states enforce identical updates under both AdamW and Muon, keeping the expanded model in its original lower-dimensional subspace.The resulting symmetry lock makes width scaling inefficient because duplicated neurons remain functionally redundant.
  • Symmetry Lock: Newton-Schulz orthogonalization in Muon cannot spontaneously break copy-induced symmetry because its polynomial Gram-matrix update preserves block-exchange symmetry.The update retains identical column blocks in the analyzed fan-in expansion case.
  • Asymmetric Optimizer States: Asymmetric Reset improves recovery speed and final loss, whereas symmetric Drop Opt. and Copy Opt. underperform substantially.Resetting optimizer states only for new channels breaks the symmetry lock; additionally scaling optimizer states provides no gain.
  • Asymmetric Learning-Rate Re-warmup: Across Inner 2×, Hidden 2×, and Hidden 2× & Inner 2× expansions, asymmetric learning-rate re-warmup lowers eventual loss under the same token budget for both zero-copy and copy-copy initialization.The gain is largest for copy-copy, where re-warmup closes the post-expansion gap to zero-copy and achieves the lowest final loss.

5. Discussions

SPARKLING combines RMS-preserving scaling with copy-based initialization and asymmetric optimizer interventions, delivering reliable downstream performance across width-expansion settings. It also reduces compute substantially while retaining or exceeding the target-width model’s performance, with complementary gains from both underlying principles.

  • Framework: SPARKLING combines RMS-preserving scaling, copy-based initialization, asymmetric optimizer state reset, and asymmetric learning-rate re-warmup.The framework is evaluated as an integrated method for mid-stage width expansion.
  • Overall performance: SPARKLING achieves the best downstream average among expansion variants and matches or exceeds the from-scratch expanded baseline on most tasks.This holds despite a small remaining gap in final pre-training loss relative to training the expanded model from scratch.
  • Ablation studies: The two underlying principles provide complementary, additive gains, with neither signal preservation nor symmetry breaking alone matching the full framework.The ablation isolates RMS-preserved scaling and asymmetric strategies on final pre-training loss and downstream performance.
  • Generality: SPARKLING generalizes across optimizer families and architectures, with RMS-preserved scaling and asymmetric re-warmup lowering final loss under Muon and transferring to dense models.The reported validation extends beyond the MoE-with-AdamW setting.
  • Compute efficiency: 20 %–35 % training FLOPs are saved under the same 200 B-token budget, with up to a 1.49× measured wall-clock speed-up under 2× width expansion.SPARKLING matches or even exceeds the from-scratch expanded model while reducing training costs.
  • Iso-compute scalability: At matched compute, SPARKLING achieves lower final loss across all three MoE width axes and dense models, with larger gains under iterative two-stage expert-inner expansion.The reported iso-compute advantage indicates more favorable compute–loss scaling and persists when expansion proceeds from 256 → 512→1024.

6. Conclusion and Future Work … B. RMS Scale Under Zero Initialization

SPARKLING stabilizes mid-stage width expansion by preserving activation RMS scale and breaking copy-induced gradient symmetry, while identifying unified width–depth expansion and µP compatibility as future directions. The derivations explain when RMS preservation reduces to fan-in variance invariance, and the zero-initialization analysis shows that zero-initialized blocks quickly behave like random ones after updating.

  • 6. Conclusion and Future Work: SPARKLING combines RMS-scale consistency with asymmetric optimizer-state reset and learning-rate re-warmup for stable mid-stage width expansion.The framework targets signal preservation and symmetry breaking during expansion.
  • 6. Conclusion and Future Work: Future work includes establishing a unified principle for simultaneous width and depth expansion and testing whether RMS preservation satisfies the µP condition.The µP condition would naturally support transferability of optimal hyperparameters after expansion.
  • A.1. Eq.: Under fan-in independence and centered homoscedastic terms, output variance decomposes additively across input dimensions.The assumptions include independence between W and x, E[wij] = E[xj] = 0, and equal weight variance across fan-in dimensions.
  • A.1. Eq.: When sin is fixed, preserving output RMS is equivalent to preserving Var(yi), which reduces to keeping dinσ2_x invariant under the stated assumptions.This provides the fan-in variance constraint used to derive RMS preservation.
  • A.2. Eq.: Copying both newly introduced fan-in weights and inputs violates independence, causing duplicated terms to contribute through covariance rather than a simple sum.The duplicated terms can contribute quadratically through covariance.
  • A.2. Eq.: For one-to-one copying with 0 < c ≤1, preserving the original variance with unchanged input scale requires rescaling the expanded-layer weights.The derivation defines copied and remaining index sets and obtains the corresponding post-rescaling variance rule.
  • A.3. Identical Gradients Under Copy Expansion for Fan-In Expansion: Under fan-in copy expansion, the two copied parameter blocks receive identical gradients, and a uniform scalar α does not remove this symmetry.This explains why copy-based initialization can hinder feature diversity.
  • B. RMS Scale Under Zero Initialization: Under RMS-preserving scaling, zero-initialized sides should be treated as random because the first gradient update quickly makes their effective statistics resemble random initialization.The analysis compares random-zero and zero-copy regimes using whole-MLP input and output activation RMS.

C. Ablation of RMSNorm Weight Expansion · D. Detailed Experimental Setup · D.1. Baseline Model Configuration

The RMSNorm-weight ablation finds random and copy initialization effectively equivalent after width expansion, motivating copy as the default. The setup uses pre-norm attention and MLP blocks, per-head q/k normalization, and a special correction for tied embedding-output projections.

  • C. Ablation of RMSNorm Weight Expansion: Random and copy RMSNorm-weight initializations quickly align after hidden-dimension copy-copy expansion and maintain nearly identical final losses.The ablation keeps SPARKLING’s linear-layer copy-copy initialization, RMS-preserving scaling, and asymmetric strategies unchanged.
  • C. Ablation of RMSNorm Weight Expansion: Copy initialization is used as the default RMSNorm initialization for all hidden-dimension expansion experiments.This choice follows the ablation’s finding that random and copy initialization remain consistent through training.
  • D.1. Baseline Model Configuration: The baseline model adopts a pre-norm design by placing RMSNorm before both attention and MLP sublayers.Detailed architectural and hyperparameter settings are provided for the pre-expansion baseline model.
  • D.1. Baseline Model Configuration: Per-head q/k normalization normalizes each projected query and key vector over its head-dimension d_head.Normalization is applied separately within each attention head.
  • D.1. Baseline Model Configuration: Tied word embeddings and output projections make hidden-size expansion act as fan-out on embeddings but fan-in on outputs.This asymmetry requires distinct RMS-preserving scaling factors for the two roles.
  • D. Detailed Experimental Setup: The tied-projection special case compensates the fan-in factor by multiplying its coefficient after the final output projection.This correction preserves the intended RMS-preserving scaling when one shared matrix serves both roles.

D.2. Training Hyperparameters · E. Hidden-Dimension Expansion: RMS-Preserving Scaling

The paper specifies a fixed AdamW-based training recipe and distinguishes learning-rate selection for scratch training from width expansion. For hidden-dimension expansion, it evaluates RMS-scale preservation when doubling width mid-training.

  • D.2. Training Hyperparameters: AdamW uses (β1, β2) = (0.9, 0.95), ϵ = 10−8, weight decay 0.1, and applies decay to norms and embeddings.The stated default applies unless otherwise specified, such as Muon experiments.
  • D.2. Training Hyperparameters: The default schedule uses cosine decay with linear warmup over 3% of total steps.Experiments run on 64 × NVIDIA A100 GPUs with 80 GB memory each and global batch size 768.
  • D.2. Training Hyperparameters: The per-device microbatch size is 3.This complements the reported global batch size of 768.
  • D.2. Training Hyperparameters: Scratch-trained models use Li et al. (2025)’s step-law-optimal peak learning rate with batch-size scaling, whereas expanded models retain the pre-expansion peak learning rate.The expanded model is trained with the same peak learning rate as before expansion.
  • D.2. Training Hyperparameters: Table 3 presents the baseline model configuration, while Table 4 summarizes the training hyperparameter configuration.These tables are identified as configuration references for the experiments.
  • E. Hidden-Dimension Expansion: RMS-Preserving Scaling: The hidden-dimension experiment doubles model hidden size from 1024 to 2048 at 100 B tokens and continues training to 200 B tokens under the same recipe.This provides the hidden-dimension counterpart of the RMS-scale analysis in Sec. 3.3.
  • E. Hidden-Dimension Expansion: RMS-Preserving Scaling: Fig. 6 reports that naive unscaled initialization can produce a smaller instantaneous loss discontinuity than RMS-preserving scaling at expansion.The passage states that enforcing RMS-scale consistency yields the same qualitative conclusion as expert-inner growth, but the supplied text ends before specifying the full continuation.

F. Analysis between Perturbation and Final Loss · G. A Sample Asymmetric Re-warmup Learning Rate Curve · H. Hyperparameter Search for Asymmetric Re-warmup

The analysis shows that immediate post-expansion perturbation does not determine final convergence, while RMS-preserving scaling improves late-stage convergence. Asymmetric re-warmup keeps original parameters on the baseline schedule, briefly boosts new parameters, and performs best with a modest ratio and short window.

  • F. Analysis between Perturbation and Final Loss: Immediate loss-spike magnitude does not directly predict final convergence: random-copy perturbs more than random-zero but reaches lower final loss.This indicates that strict function preservation is neither necessary nor sufficient; post-expansion training dynamics matter more.
  • F. Analysis between Perturbation and Final Loss: RMS-preserving rescaling consistently improves late-stage convergence over naive unscaled expansion under hidden-dimension 2× expansion.The comparison uses expansion from 1024 →2048 at 100 B tokens and mirrors the expert-inner growth pattern.
  • F. Analysis between Perturbation and Final Loss: At expansion step te, original parameters retain the baseline cosine schedule while newly introduced parameters re-warm from ηe = η(te) to ˆηmax = ρ ηe.The new-parameter schedule runs for τw steps before following the same cosine tail toward ηmin.
  • G. A Sample Asymmetric Re-warmup Learning Rate Curve: The representative curve uses continuity for original parameters and a slightly higher, short-lived learning-rate peak for newly introduced parameters before decay.This trajectory illustrates the asymmetric re-warmup schedule in Eq. (28).
  • H. Hyperparameter Search for Asymmetric Re-warmup: ρ = 1.3 and τw = 250 are adopted as the default re-warmup configuration for experiments involving re-warmup.This setting is also empirically suitable for hidden-dimension expansion.
  • H. Hyperparameter Search for Asymmetric Re-warmup: ρ ≈1.25–1.3 and τw ≈0–250 steps achieve the lowest final loss in the expert-inner 2× expansion hyperparameter sweep.The search varies the re-warmup ratio ρ and re-warmup steps τw.
  • H. Hyperparameter Search for Asymmetric Re-warmup: A broad stable region favors a modest, short-lived learning-rate boost rather than prolonged or overly strong re-warmup.The finding comes from sweeping ρ and τw under expert-inner 2× expansion.

I. Ablation for RMS-Preserving Scaling and Asymmetric Strategies · J. Comparison to Prior Function-Preserving Symmetry-Breaking Heuristics

The ablation shows that RMS-preserving scaling and asymmetric strategies provide complementary, additive benefits, with full SPARKLING achieving the strongest results across initialization pairs. Compared with prior function-preserving symmetry-breaking heuristics, SPARKLING remains consistently superior while exhibiting a transient, rapidly recovering loss increase at expansion.

  • I. Ablation for RMS-Preserving Scaling and Asymmetric Strategies: The ablation evaluates 2× expert-inner expansion across four settings, isolating naive expansion, RMS-preserved scaling, asymmetric strategies, and the combined SPARKLING framework.The asymmetric strategy consists of asymmetric optimizer-state reset and learning-rate re-warmup.
  • I. Ablation for RMS-Preserving Scaling and Asymmetric Strategies: SPARKLING combines RMS-preserving scaling with asymmetric optimizer-state reset and learning-rate re-warmup to address signal preservation and symmetry breaking.The ablation compares configurations without either component, with RMS scaling only, with asymmetric strategies only, and with the full framework.
  • I. Ablation for RMS-Preserving Scaling and Asymmetric Strategies: The two SPARKLING principles deliver complementary, additive gains: asymmetric strategies narrow but do not eliminate the gap caused by omitting RMS-preserving scaling.Adding RMS-preserving scaling on top of asymmetric strategies still reduces loss by a meaningful margin in the LLM pre-training regime.
  • I. Ablation for RMS-Preserving Scaling and Asymmetric Strategies: SPARKLING attains the lowest final loss and best downstream average across all tested initialization pairs.These results confirm the joint contribution of signal preservation and symmetry breaking.
  • J. Comparison to Prior Function-Preserving Symmetry-Breaking Heuristics: All alternative symmetry-breaking strategies converge to a higher final loss than SPARKLING under expert-inner copy-copy expansion.The alternatives include fixed or randomized uneven splitting, symmetric ± perturbation, global re-warmup, and naive function-preserving scaled initialization.
  • J. Comparison to Prior Function-Preserving Symmetry-Breaking Heuristics: Prior methods use Uneven Splitting or equal-and-opposite perturbations to break symmetry while preserving the expanded function.The comparison also includes global parameter re-warmup and scaled initialization without symmetry breaking.
  • J. Comparison to Prior Function-Preserving Symmetry-Breaking Heuristics: Around expansion, the alternatives show a transient loss up-shift followed by fast recovery, consistent with targeted exploration enabled by asymmetric re-warmup.Despite introducing asymmetry by construction, the prior heuristics remain consistently weaker than SPARKLING.

K. Comparison to Prior Dynamics-Based Strategies … O. Generalization to Multi-Stage Expansion

Across prior dynamics-based strategies, optimizers, architectures, and expansion schedules, SPARKLING consistently preserves stable training and strong performance. It achieves the lowest loss or strongest downstream results in reported comparisons, generalizes beyond MoE and AdamW, and composes across successive expansions with lower compute.

  • K. Comparison to Prior Dynamics-Based Strategies: SPARKLING achieves the lowest final loss against three representative dynamics-based baselines by jointly enforcing RMS-scale consistency and targeted backward symmetry breaking.The baselines address only individual aspects of post-expansion dynamics, whereas SPARKLING combines both principles.
  • L. Downstream Performance of Initialization and Dynamic-based Strategies: 2× expert-inner expansion: SPARKLING attains the lowest final pre-training loss and strongest average downstream performance across initialization and dynamics-based strategies.All methods use the same training recipe and token budget.
  • M. Effectiveness Under Muon: Muon experiments show that RMS-preserving scaling lowers final loss versus naive unscaled initialization, while asymmetric learning-rate re-warmup further improves late-stage convergence.These results indicate that both components remain effective beyond element-wise optimizers such as AdamW.
  • N. Generalization to Dense Models: 2× dense-model inner-dimension expansion: SPARKLING matches or outperforms the from-scratch baseline on most downstream tasks.The dense-model comparison uses intermediate size 4096 while keeping other components identical to the MoE setting.
  • N. Generalization to Dense Models: SPARKLING outperforms naive expansion and from-scratch training at the same target width despite a reduced compute budget, confirming generalization beyond MoE.The evaluation reports both final pre-training loss and downstream performance.
  • O. Generalization to Multi-Stage Expansion: 2-stage expansion from 256→512 at 50B tokens and 512→1024 at 100B tokens tracks direct 1-stage expansion closely, with only a slightly higher final loss.The iterative trajectory continues to 200B total tokens and keeps SPARKLING components identical at each expansion point.
  • O. Generalization to Multi-Stage Expansion: The multi-stage trajectory starts from a smaller model and consumes substantially less compute, revealing a trade-off between final pre-training loss and compute efficiency.The results indicate that SPARKLING composes across multiple successive expansions rather than being limited to one expansion.

P. Iso-Compute Performance and Scalability

Under matched FLOPs budgets, SPARKLING achieves lower loss than from-scratch training on dense and MoE models across multiple width-expansion settings. Iterative 2-stage expansion improves the loss margin over 1-stage expansion while also saving compute.

  • Experimental setup: The iso-compute comparison matches SPARKLING and from-scratch training on FLOPs across dense and MoE architectures and single- and multi-stage expansion.The iso-token setting instead matches training tokens while allowing different compute amounts.
  • Dense models: 0.0306 lower absolute loss than the from-scratch baseline is achieved by SPARKLING for dense models at matched compute.The comparison gives both methods the same FLOPs budget.
  • Iterative expansion: 0.0175 lower loss from 2-stage expansion exceeds the 0.0163 improvement from 1-stage expansion while providing 25% additional compute savings.Both iterative variants outperform the from-scratch baseline under the same compute budget.
  • MoE models: SPARKLING reaches lower final loss across Inner 2×, Hidden 2×, and joint Hidden 2× & Inner 2× MoE expansion at matched compute.This extends iso-compute gains across all three tested width axes.
Loading 2602.02472v2…