Papers
Topics
Authors
Recent
Search
2000 character limit reached

Preservation Is Not Enough for Width Growth: Regime-Sensitive Selection of Dense LM Warm Starts

Published 5 Apr 2026 in cs.AI | (2604.04281v1)

Abstract: Width expansion offers a practical route to reuse smaller causal-language-model checkpoints, but selecting a widened warm start is not solved by zero-step preservation alone. We study dense width growth as a candidate-selection problem over full training states, including copied weights, optimizer moments, and scheduler state. In a small-scale TinyStories proxy, we compare exact-copy, perturbative, asymmetric-reset, and structured non-clone warm starts under matched continuation budgets. We evaluate zero-step preservation, short-lag probe metrics, and downstream continuation utility in deterministic and stochastic regimes. The picture is mixed and partially replicated through a reduced-pool seed-1 check. Exact-copy symmetric warm starts rank first in every completed 16-step probe and in the completed stochastic 128-step continuations at seed-0 steps 1000 and 2000 plus reduced seed-1 step 2000. By contrast, the structured non-clone challenger wins deterministic 128-step continuation. Early escape from the inherited cloned subspace is therefore not a universal selector: it helps in long deterministic continuation, but it misleads at short lag and under stochastic continuation. The result is narrow but useful: for dense width growth at this scale, preservation is not a universal ranking criterion, and the best replacement signal depends on both regime and lag budget.

Authors (1)

Summary

  • The paper's main contribution reframes width growth as a candidate selection problem, emphasizing regime-sensitive warm start methods over traditional zero-step metrics.
  • It demonstrates that symmetric exact-copy methods excel in short-lag or stochastic regimes, while structured non-clone approaches perform better in deterministic long-horizon settings.
  • The findings highlight that probe KL is a robust low-cost selector, urging the integration of probe diagnostics for effective warm-start selection in dense LMs.

Regime-Sensitive Selection of Warm Starts in Dense Width Growth

Problem Motivation and Formulation

Dense width growth offers a pragmatic approach to scaling causal LLMs by reusing smaller model checkpoints and expanding their hidden dimensions post hoc. The central operational question addressed is how to optimally select among a set of width-expanded candidate training states—each consisting of weights, optimizer moments, and scheduler state—given a fixed parent checkpoint and a small continuation budget. Traditional zero-step preservation criteria (e.g., parent-child KL, immediate validation loss) are cheap but may fail to predict downstream utility, especially due to symmetry artifacts or inertia in the newly added capacity.

The study reframes width growth as a candidate-selection problem over full training states, instead of simply the structural widening map. The experiments are conducted in a controlled decoder-only TinyStories proxy where the width is exactly doubled while keeping depth and context fixed. Candidates include exact-copy symmetric, perturbative, asymmetric-reset, and structured non-clone (reference-subspace) warm starts. The regimes analyzed span both deterministic continuation (fixed order and seed, no dropout) and stochastic continuation (shuffled, moderate dropout injected only at continuation stage).

Experimental Protocol

All evaluations are performed on TinyStories with an LM architecture: 6 layers, width doubled from 256 to 512, and consistent optimizer hyperparameters (AdamW, fixed/cosine LR scheduling). A variety of candidate recipes are realized by controlled modifications to how the model and optimizer state are inherited and initialized. The main metrics are:

  • Downstream utility: Area under the continuation validation-loss curve (AUC, lower is better)
  • Zero-step metrics: Preservation KL, RMS drift
  • Low-cost probes: Mini-continuation probe KL, RMS, escape score (proportion of optimizer step leaving inherited subspace)
  • Selector regret: Gap in utility between best candidate and candidate chosen by a selector

Regimes are evaluated at short (16-step) and longer (128-step) horizons, cross seed and training step, and for both deterministic and stochastic continuations.

Main Empirical Results

Short-Lag and Stochastic Continuation Favors Exact-Copy

Across all short-lag (16-step) and stochastic (128-step) continuations, the symmetric exact-copy candidate dominates the downstream utility ranking. Reference-subspace and other non-clone or perturbed warm starts do not outperform pure cloning in these settings. Figure 1

Figure 1: Lag-budget and regime reconciliation; deterministic long-horizon reversals favor refsubspace_statecopy (negative), while short-lag and all stochastic runs favor exactcopy_symmetric (positive).

Deterministic Long-Horizon Favors Structured Non-Clone Candidates

A pronounced reversal emerges in deterministic 128-step continuations. Here, the reference-subspace (structured non-clone) candidate outperforms exact-copy, with negative AUC deltas reaching -3.96 in seed-0 and consistent wins across seeds/steps. This indicates the significance of symmetry-breaking and early escape from the inherited subspace, but only in sufficiently deterministic and long budget settings. However, the magnitude of the deterministic reversal is modest (~ -0.03 loss/step at most).

Selector Quality is Regime- and Lag-Sensitive

Probe KL emerges as the most reliable low-cost selector overall (mean regret 0.0178, zero-regret in 10/12 settings), but fails in some stochastic settings where probe RMS prevails. Probe escape only reaches zero regret in deterministic long-horizon, echoing its specificity as a signal for clone trapping. Zero-step metrics (preservation KL, loss) are consistently insufficient, with frequent ties and significant top-1 regret. Figure 2

Figure 2: Selector top-1 regret by metric; probe KL is strongest overall, probe escape is only exact for deterministic 128-step continuations, and neither zero-step loss nor KL work well across regimes.

Theoretical and Practical Implications

The findings decisively refute any universal selection heuristic based solely on zero-step preservation or on symmetry-breaking (escape) metrics. Instead, selection for downstream utility is sensitive to the continuation regime (deterministic vs stochastic) and the post-growth budget. Practically, for most realistic workflows characterized by stochastic continuation and limited budget, symmetric exact-copy with inherited optimizer state remains the best low-risk choice. Only when deterministic, long-horizon screening is feasible does the gains from more elaborate, structured non-clone initializations accrue. Theoretical implications touch on the need to balance preservation and symmetry breaking, a theme present in SPARKLING (Yu et al., 2 Feb 2026), but with additional nuance: escape is not universally optimal, and regime/budget context is critical.

For model developers and pipeline designers, the study suggests that candidate selection requires mixing probe-based diagnostics with awareness of update regime, and that expanding evaluation beyond zero-step metrics can prevent systematically poor choices. These findings are complementary to recent progressive-growth and morphism literature, reinforcing that practical transfer and warm-start strategies should preserve full training state and deploy short probe runs for selector evaluation.

Limitations and Future Directions

The study is constrained to single-dataset, dense decoder-only settings, small architectures, and limited seeds. The stochastic regime is operationalized rather than reflecting fully stochastic training, and the non-clone candidate space is intentionally narrow (reference-subspace only). Extensions to larger scale, other architectures (e.g., encoder-decoder, MoE), and deeper ablations on stochasticity components are warranted.

Potential future directions include:

  • Systematic disentangling of the effects of data shuffling, dropout, and sequence order noise
  • Extension to progressive token/context expansion and depth-wise growth
  • Automated or learned selector design using meta-probing or predictive modeling of candidate utility
  • Benchmarking these selection mechanisms at frontier scale and with established long-range evaluation metrics

Conclusion

Zero-step preservation metrics are insufficient as universal selectors for width growth in dense causal LMs. Selection is sharply regime- and lag-budget-dependent; symmetric exact-copy is safest for short or stochastic continuation, while structured non-clone (reference-subspace) candidates only win in deterministic long-horizon settings. Probe-based selectors, especially probe KL, are most robust overall. These results operationalize width growth as a candidate selection problem, with direct implications for efficient family construction and practical model scaling pipelines.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 1 tweet with 0 likes about this paper.