Weight decay is widely used as a regularization technique in neural network training, yet its role in neuron condensation (parameter direction alignment) remains unclear. Starting from a parameter initialization in the neural tangent kernel regime, we characterize training dynamics under weight decay through three stages: rapid fitting, amplitude compression, and neuron condensation. Using a two-layer ReLU network, we analyze a residual correlation field that governs both neuron amplitudes and directions. During rapid fitting, the residual approaches a quasi-static equilibrium maintained by weight decay while the tangent kernel remains nearly unchanged. In the early stage of amplitude compression, kernel decay amplifies the residual correlation field, whose isolated attracting extrema provide candidates for condensation directions. As neuron amplitudes stabilize, we bound the drift of attracting extrema and demonstrate contraction of neuron directions around them, leading to neuron condensation. This staged analysis provides a dynamical understanding of how weight decay promotes a condensed representation, beyond reducing parameter norms.
Figures & tables
Figure 1: Parameter norm contraction, condensation and generalization in a Transformer. Panels (a) and (b) show the first layer WQ and WK weights for λ=0 and λ=0.8 , respectively. Curves show their Frobenius norms, and insets show row-wise cosine similarities at epochs 50 and 300. Weight decay contracts the norms and induces row vectors alignment. Panel (c) shows that λ=0.8 achieves substantially higher generalization accuracy on the held-out anchor pairs (3,4) and (4,3) than λ=0 . Experimental details are provided in Appendix G
Figure 2: Training overview. (a) Data MSE, weight decay penalty, total objective, and test MSE; vertical markers indicate the loss minimum (epoch 447) and an epoch near the amplitude plateau (epoch 16,000). (b) Learned functions at epochs 447, 16,000, and 80,000. The four black dots are training samples. (c) Directions and amplitudes of all 200 neurons at different epochs.
Figure 3: Diagnostics of the three stages. (a,b) Data MSE and residual norms ∥e∥2 , ∥eqs∥2 ; the inset in (b) shows the full run, including ∥e−eqs∥2 (dotted). (c) Individual amplitude differences Ak(t)−Ak(∞) (left axis), with Ak(∞) approximated by the value at epoch 80,000, and total amplitude variation VA(t)=∑k∈N(t)∣A˙k(t)∣ over all non-dead neurons (green curve, right axis). (d) Normalized depths −qt(θ∗(t))/(λ/2) of the smooth negative minima associated with the two neuron groups. The inset shows qt(θ)/(λ/2) at epoch 5,000. The normalized depth plateau in Stage III can be explained by the balanced amplitude dynamics in Appendix A.2 . (e,f) Angular diameters and angle snapshots of two neuron groups.
Figure 4: Residual correlation field evolution and surviving neurons. Curves show normalized qt(θ)/(λ/2) , computed from the actual residual e (black solid) and its quasi-stationary approximation eqs (pink dashed), which nearly overlap. Blue and orange dots represent positive and negative neurons respectively. Black triangles mark smooth local minima. During Stage II (up to epoch 16,000), the field deepens substantially, while neuron directions remain dispersed around its extrema. During Stage III, neurons condense at these extrema.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: NTK spectrum under large initialization. (a) Full training trajectory: all four eigenvalues remain above λ=10−3 , decrease during Stage II, and approach plateaus during Stage III. (b) Zoom into the first 1,000 epochs, showing the comparatively small changes in the eigenvalues during initial fitting. The horizontal dotted line indicates λ=10−3 .
Figure 6: Fitting speed and final condensation at different initialization scales. (a) Training MSE during the first 1,000 updates. (b–d) Outputs at epoch 80,000, with targets (dotted) and training samples (black dots). (e–g) Final directions and unsigned amplitudes Ak=∣ak∣rk . Columns compare large initialization with λ=0 and 10−3 , and small initialization with λ=0 . Colored circles show non-dead neurons; gray crosses show neurons inactive on every training input.
Figure 7: Fitting precedes condensation for the targets in equation ( 54 ) (rows A–C). Left: target, training samples, early output, and final output. Middle/right: neurons after rapid fitting/at epoch 80,000, with identical marker sizes and separately scaled amplitude axes. Circles and triangles denote nonnegative and negative output weights.
Setting
Value
Layers / heads per layer
2 / 1
Embedding / feedforward dimension
400 / 1,200
Query, key, and value dimension
64
Input length / vocabulary size
9 / 120
Initialization rate
0.48
Optimizer
AdamW
Appendix
Table 1: Settings shared by the two Transformer runs in Figure 1 ; only weight decay differs.
Figure 8: Weight contraction and condensation across Transformer layers. The top and bottom rows show layers 1 and 2, respectively. Panel (a) corresponds to training without weight decay ( λ=0 ), and Panel (b) corresponds to training with weight decay ( λ=0.8 ). Within each panel, the columns show WQ and WK . Curves show Frobenius norms during training, while the insets show row-wise cosine similarities at epochs 50 and 300. Weight decay induces norm contraction and structured directional alignment in both layers.
In modern deep learning, weight decay is often credited with "stabilizing" training dynamics, diverging from its classical role as a static regularization penalty. We investigate a fundamental question: does weight decay stabilize training dynamics, and if so, through which mechanism? Indeed, training stability is understood through different but related notions in the literature. We consider how weight decay affects the parameter-space dynamics and loss sharpness by analyzing its effects at the \emph{Edge of Stability} (EoS). We show that weight decay robustly slows *progressive sharpening}. Furthermore, we uncover a striking architecture-dependent phase transition. In CNNs, weight decay dampens the oscillations at the EoS, while in MLPs, increasing weight decay causes a phase transition in which the sharpness stabilizes at a threshold significantly below the theoretical η2 boundary. We develop a mathematical framework that accurately models these phenomena and identify the global alignment of the parameter vector and the sharpness gradient as the mechanistic driver of the phase transition. Importantly, we show that these phenomena translate into stability in terms of search in function-space (NTK). Last, this shows that curvature thresholds obtained from convex/quadratic heuristics may not be reliable stability diagnostics under regularization.
Massive activations, residual-stream coordinates with magnitudes far larger than typical activations, are associated with attention sinks in transformers, but how their scale is regulated during training remains incompletely understood. Combining training-trajectory analyses and controlled interventions, we trace their emergence, growth, and consolidation. Sink-carrying channels vary across random seeds but stabilize early within each run. Over longer training, surrounding channels erode and the sink concentrates onto a few redundant carriers. Across ablations, gradient attenuation follows the sink token's collective root-mean-square magnitude rather than any single channel, making collective scale central to understanding their effects. Our central result is that weight decay causally controls the turnover of global activation scale. In controlled continuations, removing decay near the peak allows this scale to keep rising, whereas retaining it produces decline even at constant learning rate. We develop a balance model for the rise and peak of massive-activation magnitude, in which AdamW-preconditioned growth opposes weight decay. Sweeping the decay coefficient λ shifts peak timing approximately log-linearly and yields peak magnitudes scaling approximately as λ−1/2, consistent with this balance. Optimizer measurements further show that preconditioning sustains the large-channel cohort against decay even when raw maintaining forces are too small to do so. Together, these findings connect the observed life cycle to scale-regulating training dynamics and establish weight decay as a training-time lever on activation magnitude.
S. Aaron McClendon, Jorge Gallego-Feliciano, Antonios Saravanos
Aimpoint Digital Labs, Atlanta, GA, USA · New York University, New York, NY, USA
Why does weight decay work? We prove that, in any fixed-precision regime, the smallest weight norm of a looped neural network outputting a binary string equals the Kolmogorov complexity of that string, up to a logarithmic factor. This implies that weight decay induces a prior matching Solomonoff's universal prior, the optimal prior over computable functions, up to a polynomial factor. The result is norm-agnostic: in fixed precision, every weight norm collapses to the non-zero parameter count up to constants, so the same sandwich bound holds for any norm used as a regulariser. The proof has two short reductions: any program for a universal Turing machine can be encoded into neural weights at unit cost per program bit, and any fixed-precision network can be described by enumerating its non-zero parameters with logarithmic addressing overhead. Both bounds are tight up to constants, with the logarithmic factor realised by permutation encodings: a network whose parameters encode a permutation produces a string whose Kolmogorov complexity is the non-zero parameter count times its logarithm. The fixed-precision assumption is essential: with infinite precision, neural networks can encode non-computable functions and the weight norm loses its relevance.