Learnable Multipliers: Freeing the Scale of Language Model Matrix Layers

TL;DR

This work identifies that WD-driven noise induces a suboptimal fixed scale for matrix layers during large-language-model pretraining. It introduces Learnable Multipliers (LRMs), first as a scalar on weight matrices and then as per-row and per-column multipliers, to learn data-dependent scales and free scale across dimensions. Through projector and MLP experiments, as well as depth- and width-scale analyses, LRMs yield richer feature scales, broaden internal scale diversity, and improve downstream benchmarks by roughly 1–1.2% over strong baselines, with robustness across Adam and Muon optimizers. They also address training dynamics by discussing symmetries, width scaling, and gradient clipping, demonstrating that LRMs provide a universal, architecture-agnostic path to better pretraining without extra inference cost. The work prompts future exploration of generalized scaling laws under LRMs and their applicability to a wider range of architectures and optimizers.

Abstract

Applying weight decay (WD) to matrix layers is standard practice in large-language-model pretraining. Prior work suggests that stochastic gradient noise induces a Brownian-like expansion of the weight matrices W, whose growth is counteracted by WD, leading to a WD-noise equilibrium with a certain weight norm ||W||. In this work, we view the equilibrium norm as a harmful artifact of the training procedure, and address it by introducing learnable multipliers to learn the optimal scale. First, we attach a learnable scalar multiplier to W and confirm that the WD-noise equilibrium norm is suboptimal: the learned scale adapts to data and improves performance. We then argue that individual row and column norms are similarly constrained, and free their scale by introducing learnable per-row and per-column multipliers. Our method can be viewed as a learnable, more expressive generalization of muP multipliers. It outperforms a well-tuned muP baseline, reduces the computational overhead of multiplier tuning, and surfaces practical questions such as forward-pass symmetries and the width-scaling of the learned multipliers. Finally, we validate learnable multipliers with both Adam and Muon optimizers, where it shows improvement in downstream evaluations matching the improvement of the switching from Adam to Muon.

Figures (11)

• Figure 1: Projector and MLP scale \ref{['eq:equilibrium_norm_scaling']} sweep experiments described in section \ref{['sec:whatlearned']}. Norm scales $S(\eta_P,\lambda_P)$ and $S(\eta_{MLP},\lambda_{MLP})$ use relative values of learning rate and weight decay. (Top left): the final loss of three projector norm configurations. (Top middle and right): trajectories of projector $\|W\|$ and multiplier $\|\mathbf{c}\|$ norms during training are compared between the experiment and Adam Brownian Motion simulation. (Bottom left and middle): logits and matrix/multipliers norms for the considered configuration. (Bottom right): the final loss of three MLP experiment configurations.
• Figure 2: (Left): a ratio of residual block output norms of the configuration with scalar multipliers to the configuration without. (Middle): values of scalar multipliers at the end of each block. (Right): values of internal scalar multipliers at selected locations of each block.
• Figure 3: Distributions of row norms $\|W_{i\,\bullet}\|$ of attention/SSM input and MLP gate projections, which correspond to internal features of these blocks. We collect the norms across all the model layers while normalizing norm values of each layer by their mean to align the scale of different layers and focus on within-layer distribution.
• Figure 4: Effect of light WD on (a) Q/K multiplier symmetry and (b) residual output symmetry. Light WD (orange) suppresses the drift and unbounded norm growth observed in the baseline without WD (blue).
• Figure 5: The width scaling of the norms of linear layers matrices, various activations throughout the model, and scalar multipliers attached to the selected model activations. We use geometric average to aggregate the values across the layers. Other experimental details can be found in section \ref{['sec:details']}, and we provide time evolution of selected output norms on figure \ref{['fig:width_mup_out_norms']}.