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

TL;DR

Progressive Learning can dramatically reduce pre-training cost by gradually expanding model width, but mid-stage width growth faces instability from activation statistics shifts and gradient symmetry. SPARKLING introduces a dual-pronged framework: RMS-scale consistency to preserve activation statistics during expansion, and asymmetric optimization interventions (optimizer-state resets and learning-rate re-warmups) to break gradient symmetry when copy-based expansion is used. Across MoE-based transformers and multiple width axes, SPARKLING consistently outperforms training from scratch while reducing compute by up to 35% at 2× width and delivering strong downstream performance. This work enables practical mid-stage width growth, with potential extensions to joint width-depth expansion and μP-style hyperparameter transfer.

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 resetting and learning rate re-warmup. Extensive experiments on 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.

Figures (9)

• Figure 1: RMS-preserving rescaling consistently improves late-stage convergence under MoE expert inner-dimension expansion. We expand the expert inner dimension from $512 \to 1024$ at $100$B tokens and plot reference-loss (relative to the pre-expansion reference) over the remaining training tokens. (a)--(e) sweep five (up_proj -- down_proj init) pairs. In every case, Naive Init, No Scaled yields a smaller immediate loss gap, while RMS-Preserved Scaled overtakes later and converges to a lower final loss. (f) compares the RMS-preserved late-stage results and highlights a notable pattern: both-sides copied significantly underperforms other RMS-preserved strategies.
• Figure 2: Optimizer-state handling under copy-based expansion. Symmetric treatments (Drop/Copy) exhibit a symmetry lock, yielding slower recovery and higher loss. Our asymmetric reset avoids this bottleneck, while state scaling provides no additional gain.
• Figure 3: Asymmetric re-warmup consistently improves convergence under mid-stage width expansion. Across Inner $2\times$, Hidden $2\times$, and joint expansion, re-warmup lowers the final loss for both RMS-preserved copy-copy and zero-copy. Copy-copy benefits most, achieving the best final loss, effectively mitigating copy-induced symmetry lock.
• Figure 4: Under zero initialization, RMS-preserved scaling enables the post-expansion activation RMS scale to quickly recover toward the original baseline, indicating that zero initialization should be treated as random under RMS-preserving expansion.
• Figure 5: Hidden-dimension expansion mirrors expert-inner growth. We repeat the RMS-preserving scaling comparison under hidden-dimension $2\times$ expansion ($1024\!\rightarrow\!2048$ at $100$ B tokens). Across initialization pairs, RMS-preserved rescaling consistently improves late-stage convergence relative to naive unscaled expansion, exhibiting the same pattern observed for expert-inner growth in Sec. \ref{['sec:rms_scale_experiments']}.