Source-linked AI summary

Learning Orthogonal Multi-Index Models Beyond Small Initialization: Incremental Learning, Competitive Dynamics and Symmetry

Mo Zhou, Weihang Xu, Simon S. Du, Maryam Fazel

arXiv:2609.10879v1cs.LGstat.ML

TL;DR

The paper addresses how shallow networks learn orthogonal multi-index targets when standard initialization leaves neuron interactions active, beyond analyses based on small initialization or related simplifications. It proves Hermite-ordered incremental learning and competitive parameter-mass reallocation using a symmetry-based finite-width approximation. The analyzed modified gradient flow reaches small population loss with polynomial width and samples, while vanilla gradient descent shows the same qualitative dynamics empirically.

  • Problem

    Prior analyses often simplify dynamics through small initialization, correlation loss, or layer-wise training, leaving standard-initialization feature learning with interacting neurons less understood.

  • Method

    The paper analyzes polynomial-width two-layer networks using modified gradient flow and a symmetry-based finite symmetrized-network approximation.

  • Results

    The population loss decreases in Hermite order, while parameter mass shifts into the target subspace and concentrates on neurons aligned with individual target directions.

  • Takeaways & Limitations

    Competitive mass reallocation is an additional feature-learning mechanism under standard initialization, beyond independent amplification of weakly aligned neurons.

  • Takeaways & Limitations

    The analysis uses modified gradient flow, and its polynomial width and sample bounds are not intended to be optimal.

Abstract

from arXiv · show

Recent work has identified incremental learning in shallow networks trained on single-index and multi-index models. However, existing analyses often rely on simplifying settings, such as small initialization, correlation loss, or layer-wise training. These choices reduce neuron interactions and leave some feature learning dynamics under standard initialization unexplored. We study training dynamics for polynomial-width two-layer networks learning orthogonal multi-index targets under standard initialization using polynomially many samples. We first prove that incremental learning still occurs: the loss decreases sequentially according to the Hermite expansion of the target, with lower-order components learned before higher-order components recover the individual target directions. In this standard initialization regime, training also shows a competitive reallocation of parameter mass: after the total mass fits the target mean and stabilizes, mass shifts into the target subspace and then concentrates on aligned neurons. Our theoretical analysis uses slightly modified gradient flow, while vanilla gradient descent empirically exhibits the same qualitative dynamics. Technically, we introduce a symmetry-based finite-width approximation via symmetrized networks, rather than comparing directly with an infinite-width limit. This yields better control of approximation errors and may be of independent interest.

1 Introduction

The paper studies incremental learning and feature-learning dynamics for orthogonal multi-index targets under standard initialization, where neuron interactions make prior analyses insufficient. It shows that learning proceeds through Hermite-ordered fitting and competitive reallocation of parameter mass, using a symmetry-based finite-width analysis.

  • 1 Introduction: Standard initialization leaves neuron interactions active, unlike simplifying regimes that make dynamics approximately target-driven and decoupled.The network output is non-negligible from the start, so neurons interact through their shared prediction.
  • 1 Introduction: The population loss decreases in Hermite order: the mean is fitted first, followed by target-subspace recovery and individual direction identification.Lower-order components precede higher-order components in the analyzed dynamics.
  • 1 Introduction: The theorem establishes small population loss with polynomial width and samples for modified gradient flow, while vanilla gradient descent exhibits the same qualitative stagewise behavior empirically.Figure 1 uses d = 100, r = 20, m = 250, and n = 50000.
  • 1 Introduction: After the mean is fitted, parameter mass is competitively reallocated from irrelevant directions to the target subspace and then to aligned neurons.This mechanism differs from independently amplifying weak initial correlations and has a Lotka–Volterra-type effective structure.
  • 1 Introduction: The analysis uses a finite symmetrized network that preserves target and population-dynamics symmetries instead of directly coupling the finite network to an infinite-width limit.This addresses coupling difficulties caused by finite-width particle selection near target directions.

2 Preliminaries and problem setup

The setup uses Gaussian inputs, orthogonal multi-index targets, and a mean-field-scaled two-layer network with an even activation under standard random initialization. Training uses empirical square loss and evaluates population loss through a Hermite and tensor decomposition.

  • 2 Preliminaries and problem setup: The model uses Gaussian inputs and an orthogonal multi-index target whose relevant directions are orthonormal.Rotational invariance permits taking the target directions to be coordinate vectors without loss of generality.
  • 2 Preliminaries and problem setup: The learner is a mean-field-scaled two-layer network with normalized two-homogeneous parameterization.The scaling uses a factor 1/m, while neuron parameters are represented through direction and norm.
  • 2 Preliminaries and problem setup: The activation is even, with ReLU and symmetrized smooth ReLU-type activations included under the stated Hermite-coefficient conditions.Symmetrization is defined by σ(x) = ϕ(x) + ϕ(−x).
  • 2 Preliminaries and problem setup: Training minimizes empirical square loss on i.i.d. Gaussian samples and evaluates performance using population loss.The dynamics use gradient flow with weight decay and standard random initialization.
  • 2 Preliminaries and problem setup: Standard initialization gives each neuron an approximately unit norm, unlike small initialization where the network output remains initially negligible.Symmetric initialization also supports ReLU and related non-even activations through paired neurons.
  • 2 Preliminaries and problem setup: The population loss decomposes by Hermite order, with network and target moment tensors compared within each degree.Gaussian orthogonality makes the Hermite expansion central to the loss and fitting analysis.

3 Main Results

Under standard initialization, polynomial-width two-layer networks learn orthogonal multi-index targets with polynomial samples while the population loss decreases incrementally in Hermite order. The theorem also characterizes target recovery and competitive parameter-mass dynamics, using a symmetrized-network coupling strategy.

  • Polynomial-width networks learn orthogonal multi-index targets with polynomial samples under standard initialization, within the theorem’s stated parameter and sample regime.The result uses width d^3 ≤ m ≤ poly(d) and the stated polynomial sample requirement.
  • The analyzed dynamics use a modified gradient flow with suitable weight decay and stepsize, achieving target recovery within the theorem’s stated time bounds.For activation σ(x)=|x|, the loss reaches L(W) ≤ ε within T_ε = O(r/ε).
  • The population loss decreases sequentially according to the target’s Hermite expansion, reducing each component after its corresponding transition.The theorem characterizes not only final target recovery but also the order in which target components are learned.
  • The result extends incremental-learning analysis beyond small initialization and fixed-kernel comparisons by treating interacting standard-initialization dynamics with polynomial complexity.The paper contrasts its guarantee with small-initialization analyses and notes that fixed-kernel methods require super-polynomial sample complexity for these targets.
  • A finite symmetrized network is analyzed first, then coupled to the actual empirical network to control the training trajectory.This proof strategy constructs population dynamics exhibiting incremental behavior before transferring the result to the empirical network.

4 Symmetry in Dynamics

The analysis exploits symmetries of the target and Gaussian population dynamics by constructing a finite-width symmetrized network. This symmetry removes odd moments, simplifies Hermite tensors, and enables direct comparison with the actual finite-width network.

  • The finite group flips coordinate signs and permutes irrelevant coordinates, leaving the target invariant.
  • Population gradient flow preserves particle-measure invariance under this group.
  • The resulting conditional sign symmetry eliminates odd monomials and makes the second-moment tensor diagonal.
  • This finite symmetry replaces an infinite-width comparison while retaining a direct coupling to the actual network.
  • Symmetrization copies each neuron along its group orbit, producing a finite network exactly invariant under the group.

5 Training dynamics of symmetrized network under population loss

The symmetrized population dynamics learn the target incrementally through successive Hermite orders. After mean fitting, competitive mass reallocation first recovers the target subspace and then concentrates mass on aligned neurons to fit individual directions.

  • 5.1 A recurring competitive ODE: The competitive system reallocates mass from competing components to a favored component while keeping total scale nearly constrained.Its ratio q/p decreases while p remains below one, providing the mechanism for subspace learning and aligned-neuron fitting.
  • 5.1 A recurring competitive ODE: Weight decay preserves a positive competitive advantage near equilibrium and accelerates exponential decay of the competing component.Without weight decay, the toy model instead predicts polynomial decay, typically q(t) = Θ(1/t).
  • 5.2 Training dynamics in stages: The mean-fitting stage rapidly increases total mass to the target scale and produces the first loss drop.
  • 5.2 Training dynamics in stages: During second-order learning, target-subspace mass displaces irrelevant mass, taking time O(log d) and producing the second loss drop.Neurons become almost entirely supported on the target subspace, but their within-subspace directions remain near initialization.
  • 5.2 Training dynamics in stages: Higher-order terms amplify initial alignment so that at least one neuron reaches each individual target direction.This discovery occurs in time O(r/log m), while aligned neurons initially carry too little norm to change low-order population terms substantially.
  • 5.2 Training dynamics in stages: Later norm fitting transfers mass from dense non-aligned neurons to aligned neurons, fitting the target up to weight-decay bias before the unregularized stage reaches loss at most ε.

6 Coupling the symmetrized network and finite-width network

The proof transfers dynamics from the finite symmetrized population network to the actual finite-width empirical network through a direct coupling. This yields a polynomial-width control that preserves the incremental and competitive behavior.

  • The analysis couples the actual empirical network directly to the finite symmetrized network rather than to an infinite-width limit.
  • Three coupled processes separate empirical dynamics, symmetrized representatives, and the finite group-average reference network.
  • The output discrepancy decomposes into particle-coupling error and finite-width group-sampling error.
  • A stage-wise stability and concentration argument controls the coupling through time T2, with Δmax(t) ≤ ε1/6 under the theorem conditions.
  • The resulting coupling lets the finite-width empirical network inherit the symmetrized population network’s incremental and competitive dynamics.
  • The polynomial width bound is enabled by stage-wise analysis because a crude e^T stability estimate would require super-polynomial width.

7 Conclusion

The paper establishes incremental learning and competitive parameter-mass reallocation for orthogonal multi-index targets under standard initialization, using symmetry-based finite-width analysis.

  • The paper studies overparameterized two-layer networks learning orthogonal multi-index targets under standard initialization.The analysis uses modified gradient flow and exploits symmetry in the target and dynamics.
  • Incremental learning persists beyond small initialization, while parameter mass is competitively reallocated during training.The conclusion identifies this reallocation as an additional feature-learning mechanism.
  • The analysis uses a finite symmetrized counterpart rather than comparing the finite network directly with an infinite-width limit.This approach is intended to control the actual finite-width trajectory.
  • The paper defines conditional symmetry as invariance of a measure under coordinate-wise sign flips.Under this symmetry, odd polynomials vanish and the coordinate dynamics admit useful simplifications.
  • The resulting gradient dynamics evolve each coordinate multiplicatively, helping control coordinates that begin small.The paper also notes a discrepancy in a prior gradient formula while stating that the prior results should remain unchanged after proof modifications.

A.2 Omitted results and proofs for Section 3

The omitted results establish order-by-order convergence under the paper’s finite-width setting and explain how symmetrization decomposes and controls the training loss.

  • The parameter requirements include width d^3 ≤ m ≤ poly(d) and sample size n ≥ d^(3+c), with bounds presented as sufficient rather than optimal.The paper emphasizes trajectory control rather than optimizing sample complexity.
  • The analysis focuses on finite-width square-loss dynamics with non-negligible network output, where low- and high-order Hermite components coexist during training.This distinguishes the setting from target-driven analyses with decoupled neuron dynamics.
  • Symmetrization decomposes the loss into a symmetric-network term and a discrepancy term measuring distance from the symmetrized model.The symmetric term can be interpreted as bias and the discrepancy term as variance in function space.
  • Controlling the discrepancy is nontrivial because it may increase as the scale of the network output grows during training.The paper develops bounds to control this variance term.

B Initialization

The initialization analysis establishes high-probability regularity of neurons and sets up the first two learning stages: fitting the mean and recovering the target subspace.

  • With sufficient overparameterization, the initialization analysis conditions on several high-probability regularity properties of neuron coordinates and norms.These include bounds on potential neurons, coordinate gaps, and the absence of bad neurons.
  • The initialization lemmas use Gaussian concentration, spherical-coordinate tail bounds, and union bounds over neurons and target directions.The resulting failure probabilities are controlled polynomially in the width.
  • Stage 1.1: learn 0-th order: Stage 1.1 learns the 0th-order component by growing neuron norms while leaving neuron directions essentially unchanged.This stage fits the target mean without materially changing directions.
  • Stage 1.2: learn 2nd order: Stage 1.2 learns the 2nd-order component by increasing norm within the target subspace and decreasing norm in the irrelevant subspace.The directions in both subspaces remain essentially unchanged during this stage.

C.1 Stage 1.1: learn 0-th order

Stage 1.1 fits the target’s mean while preserving neuron directions, supported by an induction argument that simplifies the dynamics to the 0th-order term.

  • Stage 1.1 learns the 0th-order component while neuron directions remain essentially unchanged and norms grow to fit the mean.The stage ends after a controlled time and maintains the required regularity conditions.
  • The induction hypothesis tracks neuron directions, norms, and regularity until the stage endpoint.The proof shows the endpoint occurs before any tracked condition fails.
  • During Stage 1.1, the dynamics are dominated by the 0th-order term after higher-order contributions are controlled as errors.This simplification enables the convergence analysis for the aggregate parameter scale.
  • The resulting coordinate dynamics preserve the relevant initialization structure and yield convergence of the aggregate scale during the stage.The proof derives this from the simplified dynamics and bounds on higher-order terms.

C.2 Stage 1.2: learn target subspace (2nd-order term)

Stage 1.2 learns the target subspace through the second-order component while preserving neuron-direction balance, completing within O(log d) time.

  • The target-subspace component grows from Θ(r/d) to 1−o(1), while irrelevant-subspace directions remain unchanged.
  • Stage 1.2 ends within O(log d) time and makes the second-order term small.
  • The stage dynamics are controlled by induction hypotheses for neuron directions, norms, coordinate lower bounds, and regularity.
  • The simplified dynamics track the population zeroth- and second-order terms and establish their convergence.
  • Before convergence, target-subspace coordinates remain small while irrelevant coordinates retain their initial-scale behavior.

C.2.1 Omitted proofs for Stage 1.2

The omitted-proof analysis establishes Stage 1.2 by controlling neuron-level dynamics, deriving population ODEs, and proving convergence in logarithmic time.

  • The proof combines induction hypotheses, simplified dynamics, and integral bounds to obtain the Stage 1.2 estimates.
  • The proof shows that Stage 1.2 ends before direction, norm, or regularity conditions fail, with T12≤O(log d).
  • The population dynamics reduce to a two-variable system whose convergence is analyzed through exponential decay of a scaled ratio.
  • At an intermediate time before convergence, the leading population coefficient remains small while the second coefficient changes by only a constant factor.

D Stage 2: learn target directions (high-order term)

Stage 2 recovers individual target directions through potential neurons, then grows their norms and removes weight decay to reduce the loss.

  • Within T2−T1≲r/log m, every target direction has at least one nearly aligned neuron and the loss falls to ε.
  • The analysis separates potential neurons from small-and-dense neurons and proves that at least one potential neuron per target direction aligns.
  • Stage 2.1: Stage 2.1 recovers directions while other neuron directions remain almost unchanged, but aligned-neuron norms have not yet grown significantly.
  • Stage 2.2: Stage 2.2 rapidly increases the norms of aligned neurons until their total mass reaches the target mass.
  • Stage 2.3: Stage 2.3 sets weight decay to zero and completes recovery through continued growth of aligned-neuron norms.

D.2 Stage 2.1

Stage 2.1 recovers every ground-truth direction through competitive neuron dynamics while keeping aligned-neuron mass small. Potential neurons align selectively, whereas other coordinates remain suppressed.

  • Stage definition: Stage 2.1 is defined to end when every ground-truth direction has been learned, subject to the induction-hypothesis conditions.The stage endpoint tracks directional recovery together with balance, regularity, and norm-control conditions.
  • Stage outcome: At Stage 2.1's end, target directions are learned but the total mass of aligned neurons remains small.This stage separates directional recovery from subsequent norm growth and loss reduction.
  • Population dynamics: Population second-order terms remain balanced across target directions, supporting nearly independent directional dynamics during this stage.The proof uses this balance to simplify neuron and population dynamics and control irrelevant coordinates.
  • Neuron classes: The stage's potential neurons are a small subset, and aligned neurons remain stable under the subsequent dynamics.The analysis characterizes potential neurons at the start of Stage 2 and uses their stability for later stages.
  • Neuron alignment: Every ground-truth direction gains a nearly perfectly correlated neuron within O(r/log m) time, while non-leading coordinates remain small.The analysis compares each potential neuron's largest coordinate with all other coordinates and shows directional separation during Stage 2.1.

D.3.2 Proof of Induction Hypothesis Lemma D.2 in Stage 2.2

Stage 2.2 stabilizes recovered directions and grows aligned-neuron norms while suppressing the remaining neurons. The proof maintains induction conditions through several regularity and stability cases.

  • Induction maintenance: The proof verifies that direction, norm, stability, regularity, and balance conditions persist throughout Stage 2.2.Each condition is checked against the corresponding stopping time, showing the stage completes before any required condition fails.
  • Directional stability: The dynamics preserve balance among target directions while suppressing coordinates outside the aligned sets.The argument bounds non-aligned coordinates and maintains directional stability after the short learning-rate adjustment.
  • Norm adjustment: After Stage 2.1, aligned-neuron norms approach a common equilibrium while residual neuron mass decreases to εnorm.The resulting balanced state simplifies the later dynamics and preserves the recovered directions.
  • Transition to Stage 2.3: Once weight decay is removed, aligned norms grow toward one while residual terms remain small, producing arbitrarily small loss in the next stage.The Stage 2.3 transition begins from the Stage 2.2 state with aligned norms near equilibrium and residual mass at most εnorm.

E (Optional) Stage 3: local convergence

The optional local-convergence stage starts after Stage 2 reaches inverse-polynomial loss and uses a local gradient bound to reach arbitrary error. Its analysis also establishes finite-width coupling through symmetrized dynamics.

  • Local convergence: The local phase reaches ε loss within O(r/ε) time once the loss is below the threshold ε0 = 1/poly(r).Stage 2 supplies the required initial neighborhood, after which the local gradient lower bound yields convergence.
  • Local convergence: The local-convergence argument uses the absolute-value activation and a sample size scaling as n ≳ dr^4 log d / ε^2.These conditions are inherited from the cited local convergence result used in this stage.
  • Scope: The paper notes that extending the local-convergence argument to more general activations would require additional local analysis.The absolute-value choice is used to leverage the existing local convergence result rather than to establish a general activation theorem.
  • Finite-width coupling: The analysis couples the actual finite-width network to a symmetrized network through an intermediate process and controls their average coupling error.The decomposition is ef ≈ f̌ and f̌ ≈ f, with width-dependent concentration and Lipschitz bounds controlling the differences.
  • Symmetry construction: Symmetrization works by representing each initialized particle through a canonical orbit representative and a uniformly random group transformation.Gaussian invariance makes the transformation independent and uniform, enabling concentration of the intermediate network around the symmetrized one.

F.2 Bound ∆max and ∆

This section proves bounds on average and maximum coupling errors across training stages using Jacobian estimates, ODE comparisons, and a stopping-time argument.

  • Main bound: The main lemma combines stage-wise bounds on the decomposition terms A_t, B_t, and C_t through a Gronwall-type argument.The proof first controls average error and then individual error because the dynamics of the latter depend on the former.
  • Main bound: The analysis assumes n > d^3+c and m ≥ d^2+c, with c tied to the small exponent in ε_dir = m^−Θ(1).Under these sample-size and width conditions, the stated bounds hold through time T_2.
  • Stage-wise dynamics: The proof tracks both average error ∆ and normalized maximum error ∆max across multiple stages, including aligned neurons receiving an increased stepsize γ = m.The maximum error is normalized by the corresponding reference-neuron norm, while Stage 2.2 focuses on aligned neurons.
  • Stopping-time argument: Starting from zero errors, a stopping-time contradiction shows that the stopping time reaches T_2, so the improved bounds hold throughout the analyzed interval.Before the stopping time, the estimates imply ∆(t) ≤ ε^1/5 and ∆max remains below its stopping threshold.
  • Jacobian bounds: The Jacobian analysis exploits symmetry, including vanishing odd-order terms, and separately bounds contributions from small-and-dense and potential-neuron groups.These bounds are used to control the A_t term in the stage-wise error analysis.
Loading 2609.10879v1…