PaperlayerView exact revision on arXiv

Source-linked AI summary

Hidden Gauge Controls Feature Specialization in ReLU Networks

Tongxi Wang

arXiv:2608.06766v2cs.LGcs.AI

TL;DR

The paper asks how initialization determines which neurons acquire useful features, beyond whether a network recovers predictive accuracy. It studies function-preserving weight rescalings and finds that they can change neuron roles despite identical initial predictions, with an early learner limiting the others’ adaptation.

  • Problem

    The paper asks what makes some neurons acquire useful features while others become redundant, since accuracy recovery alone does not reveal how a network’s representation is organized.

    Source: §1 · p. 1

  • Method

    The authors rescale ReLU neurons’ input and output weights reciprocally to preserve initial features and contributions, then analyze learning dynamics in a Gaussian teacher–student model.

    Source: §1 · p. 2

  • Results

    Changing relative weight scales can make one student learn the teacher while every other coefficient vanishes, and experiments show scale can change assignment between two competing features.

    Source: §1 · p. 3

  • Takeaways & Limitations

    In the studied Gaussian ReLU model, parameter scale can determine neuron specialization despite fixed initial features and functional contributions.

    Source: §5 · p. 13

    The duplicate-student analysis assumes an initial angle to the teacher satisfying 0 < ϕ ≤ π/3.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

The success of deep learning depends on learning useful representations, yet predicting how training organizes these representations across neurons remains difficult. In this work, we show that changing the scale of initial weights can determine which neurons learn a feature without altering any neuron's initial contribution. We construct ReLU networks with identical initial features and predictions that reach the same final predictions with different roles for their neurons. In one, all neurons share the learned feature. In the other, one neuron acquires it while every other neuron's contribution vanishes. The only change is the relative scale of each neuron's input and output weights. Our analysis explains how an initial learning advantage persists through convergence: as one neuron learns the target, it reduces the error driving the others and limits their subsequent adaptation. We prove this outcome in a nonlinear model under gradient flow and small-step gradient descent, and quantify how scale changes the speed and path of feature learning. Experiments verify the predicted dynamics and show that scale also changes feature assignment when two features compete. These results reveal how initialization can control the organization of a learned representation without changing what the network initially represents.

1 Introduction

The introduction asks how initialization shapes which neurons acquire useful features, distinguishing recovery of predictive accuracy from understanding a network’s internal representation. It shows that relative input/output weight scales can change neuron roles despite identical initial features and contributions, and explains how an early advantage persists through training.

1 Introduction

  • Pruning studies differ on whether retained weights matter: resetting successful sparse subnetworks to their original weights works, while random reinitialization can degrade performance; structured-pruning results find random initialization can match or outperform fine-tuning.The introduction distinguishes recovering accuracy from explaining which neurons carry a representation.Source: §1 · p. 1
  • Positive reciprocal rescaling of a neuron’s output and input weights preserves its feature and contribution, but their relative scales affect learning under gradient descent.This function-preserving rescaling freedom is the gauge symmetry named in the title.Source: §1 · p. 2
  • Changing only relative weight scales lets one neuron acquire a feature while the others’ coefficients vanish, despite identical initial features and contributions and the same final prediction function.With matching relative scales, neurons instead remain identical and share the learned feature equally.Source: §1 · p. 2, Figure 1 · p. 2
  • As the selected student fits the target, prediction error decays and limits the other students’ subsequent rotation, leaving their coefficients to vanish.The proof establishes this in a Gaussian one-teacher model under an angle condition; the result extends to small-step population gradient descent.Source: §1 · pp. 2–3
  • Single-neuron analysis finds a quadratic separation in specialization times and inequivalent learning paths; experiments verify specialization and coefficient decay and show scale-dependent assignment between two competing features.Stability estimates extend finite-time selection to nearby initializations and convergence to small-step population gradient descent.Source: §1 · p. 3

2 Related work

The related-work discussion places the paper among studies of scale-dependent feature learning, teacher–student dynamics, symmetry, and kernel or distributional analyses. It distinguishes its focus on individual neurons and fixed Euclidean training from approaches centered on other models, geometries, or limiting regimes.

2 Related work

  • Kunin et al. (2024) analyze how relative layer scales affect nonlinear feature learning; their signed radial dynamics yield the coefficient and direction equations used here when both learning rates equal one.Source: §2 · p. 3
  • Teacher–student work studies population gradients, specialization, local convergence, regularized recovery, and overparameterization; this paper uses a quantitative Gram-matrix condition to identify individual coefficients.The cited work includes Tian (2017, 2020), Zhou et al. (2021), Akiyama and Suzuki (2021), Safran et al. (2021), and Xu and Du (2023).Source: §2 · p. 3
  • Prior work links homogeneous-network gradient flow to conserved differences of squared layer norms; this paper keeps Euclidean training fixed and uses conserved imbalance to compare trajectories from the same predictor.Related methods instead alter conditioning or balance, or use geometry invariant to node rescaling.Source: §2 · p. 3
  • Kernel, infinite-width, mean-field, and wide or deep linear analyses relate initialization to learning, whereas this work follows individual neurons through convergence at finite width.Source: §2 · p. 3

3 Theory

The theory separates coefficient adjustment from feature-direction rotation, showing that function-preserving weight rescalings alter learning paths and speeds. In duplicate-neuron models, a sufficiently imbalanced scale assignment selects one neuron while redundant coefficients vanish, including under small-step gradient descent.

3.1 Exact coefficient–direction dynamics

  • The exact dynamics separate coefficient adjustment from direction rotation, and reversing the imbalance sign leaves coefficient velocity unchanged while altering direction-update speed.Thus neurons at the same functional state can evolve at different angular speeds.Source: §3.1 · pp. 4–5

3.2 Function-preserving scaling and specialization time

  • Opposite imbalances reach the same learned feature along different paths, and the specialization-time ratio grows quadratically with D.The single-student comparison starts from a common coefficient and alignment.Source: §3.2 · p. 5, Figure 2 · p. 6

3.3 Feature assignment among duplicate neurons

  • When all imbalances agree, symmetry keeps duplicate students identical and each learns the teacher with coefficient 1/m.The mixed scale assignment instead gives the designated neuron the teacher feature and leaves redundant contributions vanishing.Source: §3.3 · p. 6
  • With one +D and m−1 −D imbalances, a sufficiently large D makes the designated student converge to the teacher while every redundant coefficient tends to zero.The theorem assumes Gaussian inputs, duplicate initializations, and a fixed neuron count and initialization angle.Source: §3.3 · pp. 6–7

3.3.1 Why the initial advantage persists

  • The limiting analysis aggregates redundant students into a total coefficient whose adjustment rate depends on their number, while the loss landscape remains unchanged.At finite imbalance, trajectory controls and symmetry ensure convergence with each redundant coefficient tending to zero.Source: §3.3.1 · pp. 7–8

3.4 Robustness and discrete gradient descent

  • For nearby initial states and sufficient imbalance, the designated neuron reaches the teacher neighborhood first, while the negative-imbalance units remain outside it at that time.The proof controls total redundant amplitude and each unit’s deviation from the mean.Source: §3.4 · p. 8
  • For sufficiently small steps and sufficiently large imbalance, full-batch gradient descent converges to the selected representation, with each redundant coefficient tending to zero.The stated result uses learning rate η = h/D.Source: §3.4 · p. 8

4 Experiments

The experiments compare four dynamical descriptions and test specialization, scaling and stability, and two-teacher assignment. They show that changing scale can change neuron roles even when initial features are fixed.

4.1 Specialization and coefficient decay

  • Across four dynamical descriptions, the selected neuron aligns with the teacher while redundant coefficients decay; their directions stay nearly fixed, so coefficient-weighted direction error falls through coefficient decay.The selected angle approaches zero, whereas the redundant angle remains nearly fixed.Source: §4.1 · p. 9, Figure 3 · p. 10
  • With the same initial coefficient and alignment, positive imbalance produces rapid alignment, while negative imbalance gives coefficient adjustment followed by slower rotation and a different specialization time and path.The specialization times differ by the predicted quadratic factor.Source: §4.1 · p. 9

4.2 Scaling and stability

  • Selection time follows D^-1 and subsequent redundant-direction motion follows D^-2; empirical trajectory error decreases approximately as N^-1/2, while gradient-descent error decreases linearly with h = ηD.Selection persisted under the tested perturbations of initial coefficients, directions, and imbalances.Source: §4.2 · pp. 10–11, Figure 4 · p. 11

4.3 Competition between two teacher features

  • Reassigning positive imbalances changes which nonduplicate students learn each teacher; when both designated students start nearer the first teacher, other students can retain part of the second feature.Both tested assignments succeeded across seeds, while the same-side configuration could distribute the second feature among additional neurons.Source: Figure 5 · p. 12, §4.3 · p. 12

5 Conclusion

Parameter scales can control which neurons specialize without changing their initial features or functional contributions. The paper proves selection and coefficient decay under specified training dynamics and demonstrates scale-dependent assignments in two-teacher experiments.

5 Conclusion

  • In a Gaussian ReLU model, a designated student learns the teacher while redundant coefficients vanish; this holds at finite imbalance under gradient flow and small-step population gradient descent.Two-teacher experiments show scale can change feature assignment among nonduplicate students.Source: §5 · p. 13

Appendices

The appendices provide derivations, convergence and stability arguments, and experimental protocols. Appendix A.1 introduces scale-aware coefficient and direction coordinates for population gradient flow and relates them to the original neuron parameters.

Appendices

  • Appendices A–D derive the dynamics and single-neuron results, establish convergence and finite-imbalance selection, and provide experimental protocols and additional results.The appendices use O(·) for upper bounds and Θ(·) for two-sided bounds, with constants allowed to depend on indicated parameters.Source: Appendices · p. 15
  • The notation O(·) denotes upper bounds and Θ(·) denotes two-sided bounds, with subscripts recording allowed dependence of constants.Generic positive constants may change between estimates.Source: Appendices · p. 15

A.1 Population gradient flow in coefficient and direction variables

  • Appendix A.1 tracks each neuron’s predictor coefficient, direction, and scale in coordinates that smoothly map the regular neuron parameters to (q, s, δ).For the finite network, the particle measure is a sum of point masses at neuron parameters, making its integral exactly the network sum.Source: §A.1 · pp. 15–16
  • Pushing the particle measure through the coordinate map yields the coefficient-and-direction dynamics, and characteristic solutions correspond to original-parameter solutions until leaving the regular coordinate domain.The theorem states the transformed continuity equation and its coordinate velocities.Source: §A.1 · p. 16
  • The loss decay separates exactly into coefficient and directional contributions, and the entire δ-marginal is conserved.Source: §A.1 · p. 16
  • The coefficient rate depends on |δ|, while the directional rate also depends on its sign; symmetric weight decay reduces |δ|, so retaining the unregularized O(D) negative-imbalance scale requires λwdD to be small.Source: §A.1 · p. 17

A.2 Effect of node rescaling beyond two layers

Reciprocal rescaling preserves a ReLU network’s function but changes how its parameters move under Euclidean gradient flow. In the Gaussian teacher–student model, the resulting dynamics admit a closed description and converge to teacher alignment and zero loss.

A.2 Effect of node rescaling beyond two layers

  • Reciprocal node rescaling preserves the network function but changes parameter gradients and the relative speeds of incoming and outgoing directions.The rescaled gradient blocks acquire distinct factors, so the joint direction velocity is not proportional to the original when both velocities are nonzero and ρ ≠ 1.Source: §A.2 · p. 18
  • For fixed q, reciprocal rescaling gives input-direction speeds ρ² and ρ⁻² times the balanced speed, a ratio of ρ⁴; as D grows, ρ⁴ = Θ(D²).Source: §A.2 · p. 19
  • The node-rescaling gradient result applies across data distributions, losses, and depths, but the global results require additional trajectory control and Gaussian structure.Those global results are proved for bias-free two-layer networks.Source: §A.2 · p. 19

A.3 Gaussian ReLU kernel and closed ODE

  • Combining classical Gaussian ReLU population gradients with the conserved imbalance δ yields an exact closed ODE for population gradient flow on the regular region.Source: §A.3 · pp. 19–20

A.4 Invariant region and global convergence

  • The invariant region keeps q bounded while alignment c increases to 1; the coefficient converges to q⋆ and the loss to zero.Source: §A.4 · pp. 20–21

A.5 Function-preserving rescaling and quadratic timescale separation

Function-preserving rescalings with opposite imbalance produce trajectories that cannot be made identical by reparameterizing time. Their large-imbalance behavior separates into fast alignment for positive imbalance and coefficient relaxation followed by slow alignment for negative imbalance.

A.5 Function-preserving rescaling and quadratic timescale separation

  • For opposite imbalances ±D from the same initial state, the trajectories cannot coincide under any strictly increasing time reparameterization.Away from the coefficient nullcline their vector fields are non-collinear; on it their phase curves have different curvature.Source: §A.5 · p. 22

A.6 Fast and slow large-imbalance limits

  • Positive imbalance permits fast alignment, whereas negative imbalance first relaxes the coefficient and then aligns slowly.Source: §A.6 · p. 23

The limiting solution is therefore

B Global fast dynamics and local geometry

This section analyzes convergence of duplicate-unit ReLU training to a selected representation and the stability of its dynamics. It compares exact dynamics with a symmetric reference flow, establishing bounds involving finite imbalance, perturbations, gradient descent, and the reference trajectory’s time to reach a specified amplitude.

B Global fast dynamics and local geometry

  • Every selected-index choice gives the same initial function and loss, and on the duplicate-redundant manifold gradient flow reduces exactly to equations for the selected amplitude q and aggregate redundant amplitude P.The redundant units remain identical by permutation symmetry.Source: B Global fast dynamics and local… · p. 25

B.1 Scaled Cartesian dynamics

  • Scaled Cartesian coordinates retain feature information when the selected coefficient vanishes; in fast time, the D → ∞ limit freezes b and v while z and p evolve.The limiting frozen values are b = 1 and v = s0.Source: §B.1 · pp. 26–27

B.2 The limiting system

  • In the limiting system, q, P, and θ follow a weighted gradient flow, while the energy landscape is independent of M and M enters through the flow metric.The limiting initial state is specified by (B.14).Source: §B.2 · p. 27

B.4 An invariant region for the limiting system

  • The inequality Fϕ(x) > 0 for 0 < x ≤ π − ϕ supports exclusion of other stationary points, while the invariant region yields a positive lower bound on q.The lower bound is uniform in m on compact angle intervals.Source: The function · p. 27, §B.4 · pp. 28–29, §B.6 · p. 30

B.5 Local loss and gradient bounds

  • Near the selected solution, a positive transverse Hessian gives local quadratic loss bounds and a gradient-to-loss bound, uniformly in m on compact angle intervals.The Hessian coordinates are the selected coefficient, selected tangent direction, and redundant coefficient.Source: §B.5 · pp. 29–30

B.6 Uniform fast convergence

  • For every m ≥ 2 and 0 < ϕ ≤ π/3, the duplicate-initialized fast flow exists globally and converges to (1, 0, 0), with uniform bounds on compact angle intervals.The constants are independent of m.Source: §B.6 · p. 30

B.7 Continuation through a loss barrier

  • If loss controls the local coordinates, decays exponentially, and bounds the motion of remaining coordinates, a trajectory stays in the neighborhood and its local coordinates converge to zero.The continuation lemma gives L(t) ≤ L(t0)e^−λD(t−t0).Source: §B.7 · p. 31

C.1 A common stability estimate

  • The next analysis separates the sum of redundant amplitudes from their differences; imbalance perturbations enter the common stability system only through its initial point.It studies how finite imbalance, initial-state perturbations, or discrete updates change the limiting trajectory.Source: C Finite imbalance, perturbations, and… · p. 31, §C.1 · p. 32
  • The common stability estimate controls selected-vector and total-amplitude errors with constants independent of m, up to a stopping time.Its stated constants have β > 1/2 and depend only on Iϕ and C0.Source: §C.1 · p. 32
  • Convexity of K0 and a Hessian bound control the selected-vector interaction error, while weighting by 1/M cancels coupling and decay terms.The reference total amplitude is the only amplitude required to be nonnegative.Source: §C.1 · p. 33
  • Proposition C.2 bounds deviations of the exact flow’s active state and redundant output from the reference trajectory.The proposition also notes that its exponent is not optimized.Source: §C.1 · p. 33
  • The reference trajectory needs at least fast time 1/2 to reach q∞ ≥ 3/4, starting from q∞(0) ≤ 1/2.Source: §C.1 · p. 34

C.2 Specialization and coefficient decay

With a sufficiently large scale imbalance, duplicate-initialized neurons specialize: one aligns with the teacher while redundant coefficients vanish. The results extend to locally identifiable nonduplicate features and small-step gradient descent under stated conditions.

C.2 Specialization and coefficient decay

  • Under exact duplicate initialization, the selected neuron converges to the teacher while redundant coefficients vanish and the loss decays exponentially.Permuting the unique positive imbalance selects a different neuron without changing the initial predictor.Source: §C.2 · p. 34

C.3 Local identifiability and asymptotic convergence

  • With quantitative independence of local feature variations, trajectories entering the selected neighborhood under the signed-imbalance condition converge to the selected solution and redundant coefficients vanish.The condition ensures redundant features remain distinguishable after removing first-order changes of the selected unit.Source: §C.3 · pp. 35, 37

C.4 Discrete full-batch gradient descent

  • For discrete training, rescaling the original parameters makes each update an exact Euler step, enabling trajectory comparison before selection and a bound on later imbalance drift.Source: §C.4 · p. 37

C.4.1 Local descent and direction motion

  • For sufficiently large D and small fast-time step h, gradient descent enters the selected neighborhood, stays nearby, and converges with every redundant coefficient tending to zero.The theorem gives entry within ⌈Tc/h⌉ iterations.Source: §C.4.1 · p. 40

D Experimental and numerical validation

The experiments compare reduced, original-parameter, and empirical dynamics against the large-D prediction. They report agreement across descriptions and scaling behavior consistent with the predicted laws over tested multiplicities.

D Experimental and numerical validation

  • For fixed multiplicity, the exact symmetric flow is predicted to differ from its large-D limit by O(D^-2) over a fixed fast-time horizon.The trajectory comparison uses the selected coefficient and direction alongside total redundant coefficient and direction.Source: D Experimental and numerical validation · p. 40

D.1 Agreement across dynamical descriptions

  • The reduced and original-parameter systems agree to numerical precision across the population grid, and empirical training agrees with its symmetry-reduced form.The tested descriptions select the positive-imbalance unit while the redundant coefficient decays as loss falls.Source: §D.1 · p. 41

D.2 Scaling with imbalance magnitude and multiplicity

  • The fitted scaling exponents remain stable from m = 2 through m = 32, and both rescaled quantities remain bounded across tested multiplicities.Source: §D.2 · pp. 41–42, Table D.2 · p. 42

D.3 Finite-sample trajectory scaling

Finite-sample experiments examine how full-trajectory error scales with sample size at D = 32, using 20 coupled seeds.

D.3 Finite-sample trajectory scaling

  • The log trajectory error is approximately linear in log sample size, with slopes consistent with N^-1/2 in the tested fixed-dimensional configurations.Source: §D.3 · p. 42, Table D.3 · p. 42
  • Every empirical run satisfies the study’s selection criterion.Source: §D.3 · p. 42
  • The metric weighting keeps the trajectory error well defined as the redundant coefficient approaches zero.Source: §D.3 · p. 42

D.4 Robustness and discrete training

The validation covers robustness and discrete training, post-selection direction motion, two-teacher competition, and a single-student comparison. Across these studies, scale changes are tested while preserving the initial predictor or neuron features as specified.

D.4 Robustness and discrete training

  • The robustness and discrete-training studies vary initial perturbations and imbalance multipliers while preserving imbalance signs, and include discrete-to-continuous trajectory comparisons.Source: §D.4 · p. 42, Table D.4 · p. 43

D.5 Post-selection direction motion under gradient descent

  • For each fixed h, the rescaled post-selection redundant-direction path is nearly constant across tested D, consistent with predicted D^-2 scaling.Source: §D.5 · p. 43, Table D.5 · p. 43

D.6 Two-teacher competition with nonduplicate students

  • Configurations A and B achieve their specified teacher assignments across all seeds, while same-side assignments also produce solutions distributing the second feature among additional neurons.In those same-side outcomes, both designated neurons can remain near the first teacher.Source: §D.6 · p. 45, Table D.6 · p. 45

D.7 Equal-scale control and single-student comparison

  • The equal-scale control matches each student's initial functional contribution to the mixed-imbalance configuration; the single-student comparison tests the predicted quadratic separation in alignment times.Source: §D.7 · pp. 45–46

Related papers

Loading 2608.06766v2…