PaperlayerView exact revision on arXiv

Source-linked AI summary

Negligible in Size, Significant in Effect: On Scale Vectors in Large Language Models

Mingze Wang, Shuchen Zhu, Yuxin Fang, Binghui Li, Kai Shen, Shu Zhong

arXiv:2605.26895v1cs.LGcs.AIstat.ML

TL;DR

The paper addresses the unclear role of scale vectors in LLM normalization layers by systematically analyzing their expressivity, optimization, weight decay, and architectural design. It shows that these small components substantially affect pre-training and develops scale-vector designs that improve training performance.

  • Problem

    Scale vectors comprise a negligible fraction of LLM parameters, yet their effect on training and the appropriate use of weight decay remain unclear.

  • Method

    The paper combines theoretical analysis of scale-vector mechanisms with experiments and proposes heterogeneous, placement-based, reparameterized, and unified scale-vector designs.

  • Results

    Scale vectors substantially affect LLM pre-training, while theory and experiments show that weight decay benefits Input-Norm vectors but harms Output-Norm vectors, and improved designs consistently improve performance.

  • Takeaways & Limitations

    Scale-vector choices should be treated as meaningful optimization and architectural decisions rather than negligible implementation details.

  • Takeaways & Limitations

    Part of the weight-decay analysis uses a scalar-output model under a continuous-time SDE approximation, and the authors identify broader architectures and model-scale evolution as future work.

Abstract

from arXiv · show

Normalization layers in modern large language models (LLMs) consist of a deterministic normalization operation and a learnable scale vector. While the normalization operation has been extensively studied, the scale vector remains poorly understood despite its ubiquitous use. In this work, we present a systematic study of scale vectors in LLMs from the perspectives of expressivity, optimization, and architectural structure. First, we show empirically that although scale vectors constitute only a negligible fraction of model parameters, removing them substantially degrades LLM pre-training. Our theory further shows that, in Pre-Norm architectures, scale vectors do not increase expressivity; instead, they improve optimization through a self-amplifying preconditioning effect on subsequent linear mappings. Second, we investigate the role of weight decay for scale vectors. By distinguishing Input-Norm and Output-Norm layers, we theoretically show that weight decay is beneficial for the former but harmful for the latter, due to their distinct roles in optimization and expressivity. Third, motivated by this understanding, we propose three lightweight and complementary improvements to scale vectors: branch-specific heterogeneity, improved placement around linear mappings, and magnitude-direction reparameterization. Both theory and experiments show that each improvement yields consistent gains. Finally, we combine these improvements into a unified scale-vector strategy and evaluate it through extensive LLM pre-training experiments on dense and mixture-of-experts models ranging from 0.12B to 2B parameters, across multiple optimizers and learning rate schedules, under industrial-scale token budgets. The unified strategy consistently achieves lower terminal loss than well-tuned baselines and exhibits more favorable scaling behavior, while adding negligible parameter and computational overhead.

1 Introduction

The paper shows that scale vectors, despite their negligible parameter count, substantially affect LLM pre-training through optimization rather than added expressivity. It distinguishes when weight decay helps or harms and proposes lightweight designs that consistently improve training.

  • Scale vectors substantially affect LLM pre-training despite constituting a negligible fraction of model parameters.
  • In Pre-Norm architectures, scale vectors do not increase expressivity; their benefit arises from self-amplifying preconditioning of subsequent linear mappings.This effect accelerates training dynamics.
  • Weight decay benefits Input-Norm scale vectors but harms Output-Norm scale vectors because their optimization and expressivity roles differ.Experiments validate these theoretical insights.
  • Improving scale vectors: Branch-specific scale vectors tailor preconditioners to distinct query, key, and value branches with minimal computational and parameter overhead.
  • Improving scale vectors: New placements provide both row-wise and column-wise preconditioning, while magnitude-direction reparameterizations induce more anisotropic preconditioners and further accelerate training.
  • Extensive experiments: Across dense and mixture-of-experts models from 0.12B to 2B parameters, the unified strategy consistently achieves lower terminal loss and more favorable scaling behavior than well-tuned baselines.The experiments span AdamW and Muon, different learning-rate schedules, and industrial-scale token budgets.

2 Understanding Scale Vectors

Scale vectors are parameter-efficient but important for LLM training: in Pre-Norm architectures they are expressively redundant yet improve optimization, while weight decay should differ by normalization-layer role.

  • Scale vectors contain only 80,640 of 1,028,065,024 parameters in the studied Llama-1B model, or 7.84 × 10^-5 of model size.
  • Expressivity: In Pre-Norm architectures, scale vectors can be absorbed into following linear maps, so they add no expressivity.For any γ and linear map W2, choosing W1 = W2 diag(γ) gives equivalent outputs.
  • Optimization: Scale vectors nevertheless accelerate training through a self-amplifying preconditioner on subsequent linear mappings, with Theorem 2.2 giving Lf(t) < Lg(t) for every t > 0.The effective preconditioner is state-dependent: Pf,j(t) = γj(t)^2Ic + wf,j(t)wf,j(t)⊤.
  • Weight decay: Weight decay benefits Input-Norm scale vectors by balancing parameters and controlling Hessian sharpness, but harms Output-Norm scale vectors by restricting expressivity.With decay, Input-Norm scale-vector norms and λmax(∇2L), Tr(∇2L), and ∥∇2L∥F remain bounded; without it, they diverge along a sequence of times.
  • Weight decay: The paper therefore recommends individual weight decay: apply it to Input-Norm scale vectors but not to Output-Norm scale vectors.Input-Norm layers are immediately followed by linear transformations, whereas Output-Norm layers are not.

3 Improving Scale Vectors

The paper proposes three complementary ways to improve scale vectors: branch-specific heterogeneity, alternative placement around linear maps, and magnitude-direction reparameterization. These designs act as lightweight preconditioners that can accelerate optimization without increasing expressivity in the relevant settings.

  • Heterogeneity: Branch-specific scale vectors adapt preconditioning to distinct query, key, value, gate, and up projection dynamics.The heterogeneity design replaces shared vectors with separate vectors for attention and FFN branches.
  • Heterogeneity: Heterogeneous scale vectors do not increase expressivity and add only O(d) parameters.Their role is optimization-related, with negligible model-size overhead.
  • Placement: After-placement modulates outputs after linear mixing, whereas dual placement controls both input-side and output-side coordinates.DNP preserves two-sided modulation while inserting normalization for potential training stability.
  • Placement: Dual placement provides both row-wise and column-wise self-amplifying preconditioning, unlike standard single-sided modulation.Theorem 3.1 establishes instantaneous acceleration at the same effective state, while DNP additionally normalizes the intermediate representation.
  • Reparameterization: Magnitude-direction reparameterization separates global scale magnitude from direction and induces an anisotropic optimization preconditioner.The paper studies original-space and exponential-space variants; exponential parameterization preserves expressivity because signs can be absorbed into linear weights.
  • Unified view: A unified view treats Pre-Norm scale-vector designs as reparameterizations that alter effective-matrix optimization rather than linear expressivity.Experiments support the complementary optimization gains of these lightweight preconditioners.

4 Experiments

The experiments evaluate individual and unified scale-vector strategies across dense and MoE models, model sizes, optimizers, and learning-rate schedules. The unified strategy consistently improves training loss while adding negligible parameter and computational overhead.

  • 4.1 Validation of Each Strategy: Each of the four strategy classes—heterogeneity, placement, reparameterization, and individual weight decay—provides gains on Llama-0.12B.The step-by-step study finds DNP+OR strongest among the placement and reparameterization combinations, while OR+IWD achieves the lowest terminal loss in the weight-decay panel.
  • 4.2 Main Results of Unified Strategy: The unified strategy combines HG, DNP, OR, and IWD and consistently improves pre-training loss across dense and MoE models from 0.12B to 2B parameters.Experiments use tuned baselines and industrial-scale token budgets.
  • Dense Models: Dense-model experiments show lower terminal loss than tuned Llama baselines across 0.12B–1B models, with the gap widening during training.Fitted scaling laws suggest the gains may persist at larger scales.
  • MoE Models: MoE experiments achieve more than 0.02 lower terminal loss than well-tuned baselines across the evaluated model sizes.The advantage persists and widens throughout training, with a slightly steeper fitted loss-descent slope.
  • Alternative Training Configurations: 0.0156 terminal-loss reduction is reported for Dense-0.12B with the Muon optimizer, alongside corresponding gains for other tested models.The cited passage also reports 0.0167 for Dense-0.5B and 0.0187 for MoE-0.5B.
  • Alternative Training Configurations: The unified strategy’s advantage increases through the stable phase and does not diminish during decay under the wsd schedule.This pattern is observed on the tested Dense-0.12B and MoE-0.5B models.
  • Overhead: 0.033 terminal-loss reduction is obtained against a parameter-matched widened Llama-0.12B baseline, while the widened baseline gives only a marginal gain.The strategy adds only O(d) parameters versus O(d^2) Transformer matrix parameters; runtime and memory increases are minor.

5 Conclusion

The paper concludes that scale vectors, despite their negligible parameter count, substantially affect LLM pre-training and motivate improved designs. It also identifies broader architectural extensions as future work.

  • Scale vectors have a substantial impact on LLM pre-training despite accounting for a negligible fraction of model parameters.
  • The paper attributes scale-vector effects to optimization rather than expressivity and clarifies when weight decay should be applied.
  • Improved scale-vector designs based on heterogeneity, placement, and reparameterization consistently improve pre-training performance.
  • Future work includes broader architectures and scale-vector designs that evolve with model scale.

B Experimental Details

The experimental details specify that all experiments use H100 80G GPUs.

  • All experiments are conducted on H100 80G GPUs.
  • The experiments use H100 GPUs with 80G memory.
  • The reported hardware setting is consistent across the experiments: H100 80G GPUs.

B.1 Experimental details for Section 4

Section 4 evaluates scale-vector designs on dense Llama and LlamaMoE models under extended-token pre-training budgets and multiple optimization settings. The setup includes tuned baselines, AdamW, Muon, and Wsd schedules.

  • Models: The experiments use dense Llama models from 0.12B to 1B parameters and LlamaMoE models from 0.25B to 2B parameters.
  • Token Budget: Training uses approximately 100 times model parameters in tokens for dense models and 100 times activated parameters for MoE models.
  • Optimization: AdamW is the default baseline with β1 = 0.9, β2 = 0.95, weight decay λ = 0.1, gradient clipping at 1.0, and a cosine learning-rate schedule.
  • Optimization: Muon is applied to Transformer matrix blocks while AdamW handles scale vectors and other non-matrix parameters.
  • Optimization: The Wsd schedule uses warmup, a stable peak-learning-rate phase through 80% of training, and linear decay to zero.
  • Ablations: Several ablations retune peak learning rates or extend training steps for models without scale vectors.

B.2.3 Experimental details for Section 3.1

The experimental configuration for Section 3.1 matches the Llama-0.12B setup described in Appendix B.1.1.

  • The model in Figure 4 uses the same training configuration as Llama-0.12B in Appendix B.1.1.
  • Figure 4 follows the Llama-0.12B experimental setup.
  • Appendix B.1.1 provides the referenced configuration for the Figure 4 model.

C.1 Proofs in Section 2.1

The proof characterizes scale vectors as lightweight optimization aids rather than sources of additional linear expressivity. They induce channel-wise preconditioners and can accelerate optimization while preserving initialization behavior.

  • The proof establishes strict improvement in the compared error dynamics for nonzero teacher columns.
  • Scale vectors introduce only d0 additional parameters versus d0 · dL matrix parameters, yet can still improve optimization.
  • Scale vectors transform gradient flow into a self-amplifying preconditioned flow through state-dependent channel-wise preconditioners.
  • The induced preconditioner acts separately on each input channel, unlike the global spectral preconditioner induced by deep linear networks.
  • Initializing scale vectors at one preserves the initial function, unlike standard small random initialization of additional deep-linear factors.

C.2.1 Proof of Theorem 2.3

This proof analyzes scale-vector stochastic dynamics and connects weight decay to boundedness of scale norms and Hessian sharpness. It also explains how Hessian statistics govern deterministic and noise-induced loss changes in SGD.

  • The SDE analysis models SGD with noise intensity q = ησ2 and tracks effective parameters, scale factors, and Hessian statistics.
  • With weight decay µ > 0, the scale-vector norm and several Hessian sharpness statistics remain uniformly bounded over time.
  • Without weight decay µ = 0, scale norms and corresponding Hessian sharpness statistics diverge along a sequence of times.
  • The deterministic GD contribution is controlled by the Hessian’s maximum eigenvalue, while the SGD noise contribution depends on the assumed noise covariance.
  • Across the analyzed noise assumptions, noise-induced loss changes are governed by Hessian sharpness metrics including λmax(∇2L), Tr(∇2L), and ∥∇2L∥F.

D.1 Proofs in Section 3.2

The proof compares input-only and input-output scale-vector parameterizations under gradient flow. Input-output scales provide instantaneous, local, and—under matching-support conditions—global optimization advantages.

  • At the same effective state, input-output scale vectors achieve instantaneous acceleration relative to input-only scale vectors.
  • Input-output scale vectors have strictly faster early loss decrease than input-only scales from zero effective-weight initialization.
  • Under matching teacher support, the dynamics decompose into independent scalar problems for each nonzero teacher entry.
  • The scalar comparison shows a larger effective speed factor for the input-output parameterization at every positive effective weight.
  • Summing the scalar comparisons yields strictly lower loss for the input-output model throughout training under the matching-support condition.

D.2 Proofs in Section 3.3

The proof studies magnitude-direction reparameterization of scale vectors and shows that its anisotropic preconditioning can accelerate optimization. The advantage is instantaneous, local, and global under balanced teacher-column conditions.

  • The OR reparameterization is at least as fast as direct scale-vector dynamics at the same effective state.
  • For d > 1, OR is strictly faster during the early phase from the stated initialization.
  • For d > 1 and balanced nonzero teacher columns, OR remains strictly faster throughout training.
  • Theorem D.2’s claims receive numerical validation in Figure 11, whose simulations support the predicted acceleration induced by OR.
Loading 2605.26895v1…