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.