Source-linked AI summary
Weight Decay and Neuron Condensation: A Three-Stage Analysis of Two-Layer ReLU Networks
Cheng Xu, Pengxiao Lin, Zhangchen Zhou, Zhi-Qin John Xu
TL;DR
The paper studies whether weight decay can induce neuron condensation under large initialization, where training otherwise fits quickly with little feature movement. It analyzes two-layer ReLU training through a residual correlation field and three stages, finding that weight decay amplifies directional attraction and supports condensation as amplitudes stabilize.
Problem
Whether weight decay induces neuron condensation under large initialization remains unclear, despite rapid fitting in the lazy regime where features move little.
Source: §1 · pp. 1–2Method
The paper analyzes amplitude and angular dynamics in two-layer ReLU networks using a residual correlation field across rapid fitting, amplitude compression, and condensation.
Source: §1 · p. 2, §3.3 · p. 5Results
Weight decay promotes neuron condensation under large initialization as early amplitude compression amplifies the field and stabilized amplitudes accompany contraction around slowly moving attracting directions.
Source: §1 · p. 2, §6 · p. 12Takeaways & Limitations
The staged dynamics explain how weight decay can form a condensed representation beyond reducing parameter norms.
Source: §6 · p. 12The analysis focuses on two-layer ReLU networks, and it does not establish that the attracting directions coincide with those in the target representation.
Source: §5 · p. 12
Abstract
from arXiv · showhide
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.
1 Introduction
The paper asks whether weight decay can produce neuron condensation even with large initialization, where lazy training fits quickly but features move little. It analyzes this process in two-layer ReLU networks and connects its dynamics to empirical Transformer observations.
Motivation
- Overparameterized networks can fit the same data in many ways, and training may find solutions whose functions need far fewer neurons than the network contains.The introduction frames this as a concrete form of simplicity relevant to generalization.Source: §1 · p. 1
- Neurons sharing an input-weight direction compute proportional features and act as one, so a wide network can behave like a narrower one.Condensation is the concentration of many neurons’ input weight vectors, including biases, around a few directions.Source: §1 · p. 2
- Small initialization and dropout are prior routes to condensation, but small initialization delays fitting as its scale decreases.The paper contrasts its large-initialization setting with these earlier approaches.Source: §1 · p. 2
- Large initialization supports rapid fitting in the lazy regime, but features move little relative to their initial scale and condensation does not occur.The paper studies whether weight decay can induce condensation in two-layer ReLU regression despite this starting regime.Source: §1 · p. 2
Main contributions
- Weight decay induces condensation under large initialization through rapid fitting, amplitude compression, and neuron condensation, with field amplification and angular contraction around attracting directions.The analysis distinguishes this process from condensation under small initialization and lazy training under large initialization.Source: §1 · p. 2
Empirical observations
- In Transformer query and key projections, weight decay is associated with row alignment and substantially higher accuracy on held-out anchor pairs (3, 4) and (4, 3) than without decay.The paper presents these observations as empirical evidence beyond two-layer networks.Source: §1 · p. 2, Figure 1 · p. 3
2 Related Work
Related work covers condensation at small initialization, transitions beyond lazy training, and other routes to neuron condensation. The paper distinguishes its focus on within-group alignment of hidden-neuron directions from other geometric properties.
Condensation and early directional alignment
- Small-initialization studies analyze feature quantization, initialization-dependent regimes, early alignment, and directional convergence near certain saddles.For shallow ReLU networks, prior work also studies alignment under orthogonal inputs and early alignment for more general data.Source: §2 · p. 3
Weight decay and the transition beyond lazy training
- Prior work studies regularized kernel dynamics, transitions from kernel behavior to parameter-space maximum-margin bias, delayed generalization, and convex reformulations of two-layer ReLU training.The cited transition analysis assumes smooth homogeneous models and vanishing weight decay.Source: §2 · p. 3
Other mechanisms leading to condensation
- Dropout, initialization changes, and weight decay are among other reported routes to condensation, while neural collapse and low-rank structure describe different geometric properties.The paper focuses on within-group angular diameter of hidden-neuron directions, not concentration of sample representations or matrix rank alone.Source: §2 · p. 4
3 Preliminaries
The preliminaries set up a two-layer ReLU gradient-flow model and describe neurons through amplitudes and directions. They define the residual correlation field, condensation, and the observed three-stage training progression.
3.2 Network and training objective
- The experiments use a two-layer ReLU network with m hidden neurons, independently initialized parameters at large initialization, and gradient flow with weight decay on all parameters.The main target is a sum of two ReLU functions, as specified in the experiment setup.Source: §3.2 · p. 4
3.3 Amplitude and angular dynamics
- Each neuron is represented by an amplitude and direction, with ϕθ as its directional ReLU feature and ψθ as the feature’s angular derivative away from activation boundaries.The network output is expressed through signed amplitudes multiplying directional features.Source: §3.3 · pp. 4–5
- The residual correlation field qe pairs the training residual with the directional feature, and its angular derivative pe pairs the residual with the angular derivative feature.These quantities are evaluated along training as qt and pt.Source: §3.3 · p. 5
The training residual satisfies
- Weight decay acts directly on parameter scales, while angular motion is governed by the field derivative; initially strictly inactive neurons lose amplitude exponentially while their directions remain fixed.The scale imbalance also decays exponentially, motivating a balanced description of late training.Source: The training residual satisfies · p. 5
3.4 Neuron condensation
- Neuron condensation means that the angular diameter of a fixed group of active neurons tends to zero, so distinct neurons approach a common direction.Different groups may condense around different directions.Source: §3.4 · p. 5
3.5 Three-stage overview
- Training progresses from rapid fitting with broadly distributed directions, through amplitude compression and strengthening attraction, to approximately exponential angular contraction after amplitudes nearly stabilize.The overview describes features evolving from scattered to narrowing and finally tight angular groups.Source: §3.5 · p. 6
4 Main results
The analysis follows neuron dynamics through the residual correlation field: rapid fitting leaves a residual, compression amplifies the field, and stabilized amplitudes support directional contraction.
4.1 Stage I: Fast residual equilibration
- After rapid fitting, the residual approaches its quasi-static reference while the kernel remains close to its initial value, leaving a nonzero residual that drives subsequent neuron motion.For λ > 0 and y ≠ 0, weight decay sustains the quasi-static residual.Source: §4.1 · pp. 7–8
4.2 Stage II: Discrete extrema and amplification of qe
- For a fixed residual, smooth extrema in non-dead activation cells are equilibrium candidates: positive neurons are attracted to negative minima, and negative neurons to positive maxima.The extrema form a discrete set that can serve as condensation directions.Source: §4.2 · pp. 8–9
- During early amplitude compression, residual tracking persists as weight decay reduces the kernel and amplifies the correlation field, while leading-order extrema locations remain unchanged.In the actual dynamics, smooth extrema also move during Stage II.Source: §4.2 · p. 9
4.3 Stage III: Slow drift of extrema and neuron condensation
- Amplitude variation controls angular energy and bounds extrema drift; after the initial angular-energy contribution decays, the drift speed is bounded by a term proportional to amplitude variation.The reported small amplitude variation during Stage III supports a small drift bound.Source: §4.3 · pp. 10–11
- Under balanced amplitudes, positive-neuron groups inside a fixed band around a sufficiently deep, slowly moving negative minimum remain in-band and contract exponentially, yielding condensation.If the minimum's drift tends to zero, each neuron's direction converges to it.Source: §4.3 · pp. 11–12
5 Discussion
The analysis is for two-layer ReLU networks and characterizes condensation around residual-field extrema, without establishing that those directions match the target representation.
5 Discussion
- The dynamical theory is limited to two-layer ReLU networks; it does not establish that residual-field extrema match target directions, and corresponding Transformer dynamics remain open.Additional target-function experiments show the same progression from rapid fitting to late condensation, while also exhibiting more condensation directions than the target representation.Source: §5 · p. 12
6 Conclusion
The analysis describes how weight decay can produce neuron condensation after rapid fitting, through amplitude compression and the attraction of neuron directions to extrema of a residual correlation field. It also develops neuron dynamics and balance relations and establishes initialization properties for the kernel analysis.
6 Conclusion
- The analysis characterizes rapid fitting, amplitude compression, and neuron condensation, with rapid fitting preceding substantial directional alignment during continued weight-decay training.The residual correlation field has isolated attracting extrema, and stabilized amplitudes support contraction around slowly moving attracting directions.Source: §6 · p. 12
A.1 Neuron dynamics and balance
- The neuron-dynamics analysis separates radial and angular motion and establishes balance relations, while strictly inactive neurons decay exponentially without changing direction.For a strictly inactive neuron, a_k(t) and v_k(t) scale as e^−λt and its squared amplitude A_k(t) as e^−2λt.Source: §A.1 · p. 16
A.2 General and balanced amplitude dynamics
- Under exact balance, the amplitude dynamics reduce to a system in which positive amplitude remains positive on finite intervals with bounded correlations.For surviving positive neurons near attracting minima with nearly stationary amplitudes, the reported relation is −2q_t(θ*)/λ ≈ 1.Source: §A.2 · p. 17
B.1 Proof of Theorem 2: initialization
- With fixed sample count and γ ≥ 1/4, the Gaussian initialization yields ∥e(0)∥₂ = O_p(1), and the normalized initial kernel converges to a positive-definite limit.The kernel convergence is in operator norm in probability.Source: §B.1 · p. 18
B.2 Proof of Theorem 2: kernel stability
The proof controls the residual and kernel over an initial fitting horizon, then shows that the residual tracks a quasi-static equilibrium as the kernel changes little.
B.2 Proof of Theorem 2: kernel stability
- On a high-probability localized event, the residual norm is uniformly bounded over the fitting horizon, and the initial fitting time lies within that horizon with probability tending to one.The residual norm's upper rate of change is bounded by 2λ∥y∥₂.Source: §B.2 · pp. 19–20
B.3 Proof of Theorem 2: residual equilibration
- The residual tracks the quasi-static equilibrium with error o_p(ε_m) through the fitting interval, while that equilibrium changes by only o_p(ε_m).The result follows from kernel stability and the positive limiting minimum eigenvalue.Source: §B.3 · p. 20
C.1 Proof of Proposition 3
The proofs bound the residual correlation field's extrema and control how the kernel, amplitudes, and neuron directions evolve around attracting minima. Under stated balance and stability conditions, the angular diameter of a group contracts even if its attracting minimum does not converge to a fixed angle.
C.1 Proof of Proposition 3
- The field has between one and 4n strict local extrema, each either inside an activation cell or at an activation boundary.There are at most 2n activation boundaries and 2n open cells; boundary extrema have fixed positions while they persist and attract positive neurons under the stated one-sided condition.Source: §C.1 · p. 21
C.2 Proof of Theorem 4
- During the controlled interval, the kernel remains positive definite, and the actual correlation field is compared with its quasi-static reference as the kernel decays.The initial tracking mismatch relaxes on the kernel time scale, while reference motion on the decay time scale produces a small lag.Source: §C.2 · pp. 22–23
D.1 Amplitude budget and angular energy
- Under Theorem 5's hypotheses, the set of surviving neurons is constant after T, and exact balance and energy dissipation yield the amplitude-budget argument.The supplied proof also distinguishes output-velocity vectors from the scalar sum of absolute amplitude derivatives.Source: §D.1 · p. 24, §D.2 · p. 24
D.2 Proof of Theorem 5
- The proof differentiates the output sum to relate its velocity to neuron changes and bounds the motion of a smooth minimum through its depth.It also derives velocity estimates using the scalar sum of absolute amplitude derivatives.Source: §D.2 · pp. 24–25
D.3 Consequences of small amplitude variation
- When the total absolute amplitude variation tends to zero, the angular-energy bound tends to zero, and residual and minimum velocities then tend to zero.A uniform bound on amplitude variation alone gives a stated upper bound on the limiting angular energy.Source: §D.3 · p. 25
D.4 Angular tracking
- An exactly balanced positive-output neuron remains in a band around a smooth negative-minimum branch when the branch's drift is sufficiently small relative to its depth.Within the band, the tracking estimate contracts the angular offset at rate c_d = 2R_0 cos d.Source: §D.4 · p. 26
D.5 Pairwise angular contraction
- For exactly balanced positive-output neurons sharing a band around a sufficiently deep smooth minimum, their angular diameter contracts to zero even if the minimum does not converge.The result assumes the shared band lies inside one open activation cell.Source: §D.5 · p. 26
D.6 Proof of Theorem 6
Under the theorem's conditions, tracking and pairwise contraction imply that the directions of the neuron group converge toward one another around the moving attracting minimum.
D.6 Proof of Theorem 6
- The group's angular diameter tends to zero under the positive contraction rate supplied by the tracking and pairwise-contraction bounds.The tracking estimate also implies that each group's direction approaches the moving minimum when its drift tends to zero.Source: §D.6 · p. 27
E.1 Target and training data
The experiments use a four-point ReLU target and compare training behavior under large initialization and weight decay. The NTK eigenvalues change little during initial fitting, decline during amplitude compression, and later stabilize above the decay coefficient.
E.1 Target and training data
- The four-point target is f⋆(x) = 0.2 ReLU(x −1/3) + 0.2 ReLU(−x −1/3) on [−1, 1], with test MSE evaluated at 10,000 equally spaced points.The target and evaluation domain define the experiment’s function-fitting task.Source: §E.1 · p. 28
E.2 Network and optimization
- Training uses 200 hidden neurons, initialization scale γ = 0.4, 80,000 full-batch SGD updates, learning rate 0.1, and weight decay λ = 10−3.Parameters are initialized independently from N(0, m−2γ).Source: §E.2 · p. 28
- All four NTK eigenvalues begin much larger than λ for sufficiently wide networks with γ < 1/2, then decline after rapid fitting and stabilize above λ during Stage III.The observed trajectory is consistent with early kernel decay Kt ≈ e−2λ(t−TI)KTI under the theorem’s assumptions.Source: §E.2 · p. 28, Figure 5 · p. 28
F.1 Initialization scale and weight decay
Large initialization with weight decay yields rapid fitting followed by neuron condensation, while large initialization without decay fits rapidly but does not condense. Across three additional targets, fitting precedes the formation of narrow directional groups.
F.1 Initialization scale and weight decay
- Large initialization with λ = 10−3 and small initialization without decay both produce two dominant angular groups, but large initialization fits within hundreds of updates before late condensation.Small initialization without decay fits more slowly.Source: §F.1 · p. 29
- With large initialization and λ = 10−3, all non-dead neurons condense, including a low-amplitude group near 90°, while dead neurons remain dispersed as their amplitudes vanish.Some low-amplitude non-dead neurons remain dispersed with small initialization and no decay.Source: §F.1 · p. 29
- Large initialization without decay fits rapidly but does not produce a condensed representation.This contrasts with the large-initialization run using weight decay.Source: §F.1 · p. 30
F.2 Targets with three ReLU breakpoints
- For all three targets with three ReLU breakpoints, outputs track targets after rapid fitting, while neurons with appreciable amplitudes later form narrow angular groups and final outputs show regularization bias.The targets include a convex case and two cases with peak and trough in opposite orders.Source: §F.2 · p. 31
G Transformer experiment on anchor function composition
The Transformer experiment compares matched runs that differ only in weight decay on a compositional anchor-function task. The reported weight-decay run shows norm contraction and structured directional alignment in query and key weights across both layers.
G Transformer experiment on anchor function composition
- The two Transformer runs use identical saved data and initial parameters, with weight decay coefficients 0 and 0.8 as the sole comparison difference.The experiment complements the two-layer network analysis with Transformer observations.Source: G Transformer experiment on anchor… · p. 32
G.1 Anchor functions and data construction
- Each input contains a consecutive anchor pair (a, b) and a preceding key x; the four elementary anchor functions are translations, and the target is their composition ga(gb(x)) = x + da + db.Because the translations commute, the pair order does not affect the target; (3, 4) and (4, 3) both yield x −10.Source: §G.1 · p. 32
- Training excludes anchor pairs (3, 4) and (4, 3), while testing includes all 16 pairs, so those two test groups assess generalization to unseen compositions.For the other 14 seen pairs, the test split assesses held-out key–position combinations.Source: §G.1 · p. 32
G.2 Network and initialization
- The model is a two-layer decoder-only Transformer with one causal self-attention head per layer.Its input includes learned positional embeddings.Source: §G.2 · p. 32
G.3 Training objective and weight decay comparison
- Training minimizes cross entropy at the last output position, and AdamW applies decoupled weight decay to all trainable parameters.The learning rate warms up over 40 epochs and then follows a cosine schedule with a 360-epoch period.Source: §G.3 · p. 33
G.4 Cosine similarities
- With λ = 0.8, weight decay induces norm contraction and structured row-wise directional alignment in WQ and WK across both Transformer layers.Figure 8 shows Frobenius norms over training and row-wise cosine similarities at epochs 50 and 300; the comparison run uses λ = 0.Source: §G.4 · pp. 33–34, Figure 8 · p. 34