Symmetry-Aware Feature Learning: A Polynomial Separation for Multi-Index Models

1 Introduction

• The study asks how much symmetry reduces SGD's sample needs to discover a relevant feature and how weight sharing compares with data augmentation.It analyzes high-dimensional, nonconvex SGD dynamics from random initialization rather than comparing hypothesis classes or estimators after optimization.Source: §1 · p. 3
• The teacher directions are symmetry-related copies of one direction, and the growing-rank setting makes them dependent competitors rather than independent parameters.The rank scales as r ≍ d^δ in the stated model.Source: §1 · p. 3
• nT = eΘ(d^(p-1)), nA = eΘ(d^(p-1)), and nU = eΘ(rd^(p-1)) for weak directional recovery with polynomial links of information exponent p ≥3.The bounds match up to logarithmic factors; at width s ≳ r log r, the untied and augmented procedures cover the teacher orbit within their respective sample-complexity scales.Source: §1 · p. 4
• For the pure quadratic Hermite link, tied and augmented learners need eΘ(d) samples for weak subspace recovery, versus eΘ(rd) for the untied learner.Individual orbit directions are generally unidentifiable in this regime.Source: §1 · p. 4
• For p ≥3, initialization fluctuations select a leading orbit direction, whose overlap grows to a macroscopic level while competing overlaps remain microscopic.For the quadratic link, the proof instead tracks multiplicative growth of squared projection onto the teacher subspace.Source: §1 · p. 4

2.1 Teacher model

• The teacher depends on one direction through its cyclic shifts, so the statistical model is invariant when that direction is shifted.The shift operator is a cyclic permutation with Π^r = I_d.Source: §2.1 · p. 5
• For polynomial links with information exponent p ≥3, the recovery target is an identifiable teacher-orbit direction; for the pure quadratic link, it is the orbit-spanned subspace.The quadratic case generally does not identify individual orbit directions.Source: §2.1 · p. 6

2.2 Learning procedures

• The tied student shares one trainable weight across transformed copies, while the untied student has separate weights not constrained to a common orbit.The augmented procedure uses the untied model class but replaces each sample loss with its group average.Source: §2.2 · p. 6
• Augmentation leaves the population objective unchanged and uses each fresh sample to generate its r transformed inputs, increasing per-iteration computation.The number of SGD iterations still equals the number of fresh samples.Source: §2.2 · pp. 6–7

3.1 Weak directional recovery (p ≥3)

• For p ≥3, tied and augmented procedures need eΘ(d^(p−1)) samples, while untied training needs eΘ(rd^(p−1)) for weak directional recovery.These are matching upper and lower bounds up to logarithmic factors under the stated dimension-dependent learning rates.Source: §3.1 · p. 8
• With width s ≥ (1 + ε)r log r, untied and augmented procedures can cover every teacher-orbit direction without increasing the sample-complexity scale.The tied procedure represents the full orbit through weight sharing; the stated width condition applies to the untied and augmented procedures.Source: §3.1 · pp. 8–9
• For Hermite-4 at δ = 1/2, Figure 1 compares single-weight recovery times with shared initialization; tied and augmented trajectories coincide.The plotted points are median samples to reach Mσ(w) ≥ 0.5 over 10 paired seeds, with pointwise 95% bootstrap intervals.Source: §3.1 · p. 8, Figure 1 · p. 10

3.2 Weak subspace recovery (p = 2)

• For the pure quadratic link, tied and augmented procedures need eΘ(d) samples, while the untied procedure needs eΘ(rd) for weak subspace recovery.For general centered polynomial links with information exponent 2, the matching upper bound is not established; it is proved only for the pure quadratic link.Source: §3.2 · p. 10

Learning low-dimensional representations

• Multi-index and neural-network studies examine extracting low-dimensional structure through moment-based methods or evolving network weights.Examples include filtered PCA and gradient-based refinement, representation learning under second-order nondegeneracy (Damian et al., 2022), and weight decay (Mousavi-Hosseini et al., 2023).Source: §4 · p. 11

Staircase structure and leap complexity

• Staircase and leap-complexity work studies how hierarchical structure supports learning additional coordinates or features.Abbe et al. (2021, 2022, 2023) establish Boolean and Gaussian learning results, including movement between successive saddles under modified SGD.Source: §4 · p. 11

Gradient dynamics and subspace discovery

• Online-learning studies connect information exponent, learning-rate scaling, fresh-batch updates, and Hermite structure to feature or subspace discovery.These include single-index dynamics (Ben Arous et al., 2021), multi-index gradient methods (Dandi et al., 2024a), and successive population-flow phases.Source: §4 · p. 11

Growing rank and subspace recovery

• Prior growing-rank work studies untied teachers, orthogonality-constrained SGD, and distinct stages of subspace versus directional recovery.Oko et al. (2024) report an upper bound with the same leading r and d dependence as this paper’s untied learner; Ren & Lee (2025) distinguish subspace discovery from identifying individual directions.Source: §4 · p. 12

Weight sharing and data augmentation

• Equivariant architectures and group-averaged training are established ways to exploit symmetry, with prior benefits studied in statistical and structured-image settings.The paper compares these approaches with untied training in online feature learning as teacher rank grows with dimension.Source: §4 · p. 12

5.1 Proof overview

• For p≥3, initialization selects a leading orbit correlation, whose gap widens during microscopic growth before it reaches weak recovery while competitors remain small.The proof controls interactions, stochastic fluctuations, and normalization errors so they do not erase the initial gap.Source: §5.1 · p. 13
• At the prescribed learning rates, an augmented weight has the tied trajectory’s effective drift and fluctuation scales, while an untied weight’s drift is slower by factor r.The upper-bound argument follows any fixed weight for the untied and augmented procedures, without choosing the best initialization.Source: §5.1 · pp. 13–14, §5.2 · p. 15

5.4 Non-escape and lower bounds

• The non-escape estimates keep all weights at microscopic correlation for the stated intervals, excluding directional and quadratic subspace recovery there.The resulting lower bounds apply to the specified polynomial links and procedure-specific learning rates.Source: §5.1 · p. 14, §5.4 · p. 16

5.2 Upper bounds for directional recovery

• For p≥3, the upper-bound results establish weak directional recovery for tied, untied, and augmented spherical SGD under their procedure-specific step sizes.For untied and augmented procedures, the stated upper bound holds for any fixed weight index.Source: §5.2 · pp. 14–15

5.3 Upper bounds for quadratic subspace recovery

• For the quadratic link, the upper bounds are eO(d) samples for tied and augmented procedures and eO(rd) for untied.The recovery target is the teacher-orbit subspace rather than an individual direction.Source: §5.3 · pp. 15–16

5.5 Proofs of the main results

• For p≥3, tied and augmented learners achieve weak directional recovery at eΘ(d^(p−1)) samples versus eΘ(rd^(p−1)) for untied; quadratic subspace recovery scales as eΘ(d) versus eΘ(rd).The proof combines upper bounds with non-escape lower bounds for the specified procedures and learning rates.Source: §5.4 · p. 16, §5.3 · p. 16, §5.5 · p. 17

6 Correlation dynamics and auxiliary estimates

• The section supplies the dynamical identities and stochastic estimates used for tied-student weak recovery, then explains how later sections adapt them to augmented and untied training.The tied student's correlations provide the cleanest form of the dynamics.Source: §6 · p. 17

6.1 Correlation recursion and population drift

• The exact correlation recursion combines population drift, higher-order and interaction terms, centered stochastic fluctuation, and spherical-normalization error.The population drift is derived using Hermite orthogonality and the cyclic teacher geometry.Source: §6.1 · pp. 18–20

6.2 Coordinate noise estimates

• Coordinate-noise estimates bound conditional moments uniformly over iterates on the teacher-geometry event.The bounds apply across coordinates and teacher configurations in that event.Source: §6.2 · p. 20

6.3 Uniform control of the stochastic gradient

• The stochastic-gradient norm is controlled with high probability uniformly over polynomially many SGD iterations.The stated bound covers every deterministic horizon T ≤ d^C for fixed C > 0.Source: §6.3 · p. 21

7 Proof of weak recovery for the tied student

• The leading polynomial dependence of tied-student recovery time is d^(p−1); higher Hermite components affect only constants and logarithmic factors.Source: §7 · p. 21

7.1 Proof strategy

• The proof has two phases: amplify one initialized teacher correlation to an intermediate scale, then grow it to the fixed recovery level ρ.The restart at the intermediate threshold uses the strong Markov property of the SGD trajectory.Source: §7.1 · pp. 21–22

7.2 Microscopic selection

• With high probability, a unique initialization leader is separated from competitors, and its advantage survives fluctuations as it reaches βd while competing correlations stay at scale αd.The selected leader remains unique through this microscopic-selection phase.Source: §7.1 · pp. 21–22

7.3 Growth to a macroscopic level

• From the intermediate scale, the selected overlap grows to fixed ρ while its unique-leader status persists and competing overlaps remain microscopic.The competitor dynamics include a contractive damping factor generated by the selected overlap.Source: §7.1 · p. 22, §7.3 · pp. 27, 30
• t_ρ := inf{n: L_n ≥ ρ} is bounded by C_{σ,ρ} ηd((1−ε)β_d)^{p−2}, and the selected coordinate reaches ρ before T.Choosing C_T sufficiently large ensures t_ρ ≤ T after the other stopping mechanisms have been excluded.Source: §7.3 · p. 30
• Mσ(w_{tρ}) ≥ ρ because α_d = o(β_d) keeps i* the unique favorable leader.This establishes Proposition 7.2.Source: §7.3 · p. 30

8.1 Componentwise reductions

• Each augmented weight follows exactly the tied SGD trajectory with reference step size ηd, driven by the same fresh samples.This componentwise identity yields the tied trajectory law for an augmented weight with independent uniform initialization.Source: §8.1 · p. 31

8.2 One-weight untied dynamics

• The single-weight untied population landscape matches the tied scalar overlap landscape, but its magnitude is smaller by a factor r^-1/2.The resulting one-weight dynamics provide the basis for the untied recovery analysis.Source: §8.2 · p. 32

8.3 Proof of the untied weak recovery bound

• The untied student’s selected overlap reaches ρ while competitors remain microscopic, with the proof scaling the recovery horizon by r relative to tied dynamics.The selected drift is reduced by r and stochastic increments by √r; the untied theorem applies the one-weight trajectory law to each fixed weight.Source: §8.3 · pp. 33, 35

8.4 Proof of weak orbit coverage

• With s ≥ (1 + ε)r log r weights, the proof combines initial-leader coverage and simultaneous recovery to establish full orbit coverage for the untied student.The argument bounds failure to initialize a suitable weight for each teacher direction, then applies simultaneous recovery of initial leaders.Source: §8.4 · pp. 35–37
• For augmented weights, symmetry makes the conditional probability of recovering any fixed teacher direction equal across directions, enabling a coverage bound across weights.The proof uses conditional independence of augmented trajectories across weights after conditioning on the teacher and samples.Source: §8.4 · p. 37
• The preceding estimates prove coverage by TU and TA for their respective procedures, allowing different directions to be recovered at different times.Source: §8.4 · p. 38
• Each iteration of (5) uses one fresh observation.Source: §8.4 · p. 38
• Theorem 5.4 supplies matching lower bounds, completing the corollary’s argument.Source: §8.4 · p. 38

9.1 Proof strategy

• Unlike the p ≥3 argument, the quadratic-link proof tracks total mass in the teacher-orbit subspace because the leading coordinate drift is linear.Once total overlap reaches a macroscopic level, it yields weak subspace recovery.Source: §9.1 · pp. 38–39

9.2 Total overlap and initialization

• Uniform spherical initialization gives total overlap R0 ≍ r/d with probability 1 − o(1), and the same initialization scale applies to each fixed untied or augmented weight.The projected mass onto the teacher subspace is comparable to R0.Source: §9.1 · p. 39, §9.2 · p. 40

9.3 Recursion and population growth

• Below any fixed threshold a < 1, the population contribution to total-overlap growth is comparable to eta d R.The quadratic-link identity expresses this contribution through the squared norm of the spherical gradient.Source: §9.3 · pp. 40–41

9.4 Noise and gradient moments

• Uniform noise and gradient moment bounds control relative fluctuations and yield high-probability bounds over polynomial horizons for both tied and untied processes.The estimates are uniform over current weights and teachers on the teacher-geometry event.Source: §9.4 · pp. 42–43

9.5 Growth and proof of the tied upper bound

• The tied process grows its total overlap to a fixed threshold, with high-probability control of the stopped relative noise.The proof tracks log overlap, whose increments have positive drift of order eta d; accumulated relative-noise, remainder, and normalization errors are o(1) over the stated growth horizon.Source: §9.1 · p. 39, §9.5 · p. 43
• With high probability, the tied learner’s initial overlap is at least cinit r/d and its growth reaches a level sufficient for weak subspace recovery.The proof translates total-overlap growth into subspace recovery and establishes the tied upper bound.Source: §9.5 · p. 47

9.6 Augmented and untied upper bounds

• Absorbing the ceiling into CT(ρ) proves part (T); recovery of this weight implies student recovery, so part (A) holds with CA = CT; absorbing the ceiling into CU(ρ) proves part (U).Source: §9.5 · p. 47, §9.6 · pp. 48–49
• For the untied effective update, drift and normalization coefficients are divided by r, while the noise coefficient is divided by √r.The proof uses this recursion to establish the untied recovery guarantee.Source: §9.6 · pp. 48–49
• When a tied weight achieves subspace recovery, all its cyclic shifts have the same projection mass in the teacher subspace.The teacher subspace and its orthogonal complement are invariant under the cyclic shift operator Π.Source: §9.6 · p. 49

10.1 Proof strategy

• All absolute correlations remain at initialization scale for a sufficiently small fraction of the microscopic growth time, excluding weak directional and subspace recovery across all three procedures.The non-escape argument applies even when the procedures share observations.Source: §10.1 · p. 50
• The proof bounds population drift by the largest absolute correlation raised to p −1, then controls accumulated coordinate noise until the first escape.The drift comparison retains exact spherical normalization.Source: §10.1 · p. 50

10.2 Initialization and escape times

• A simultaneous tail bound over initial weights and teacher directions, combined with the teacher-geometry event, establishes the required initialization-scale correlations with probability 1 −o(1).The argument uses a union bound over the initial correlations.Source: §10.2 · pp. 50–51

10.3 An upper comparison for the correlations

• Pathwise upper comparisons bound correlations for tied, augmented, and untied weights through the first exit step.The comparisons apply uniformly over weights and teacher directions.Source: §10.3 · pp. 51–53

10.4 Noise before escape

• Before escape, stopped coordinate fluctuations satisfy simultaneous partial-sum bounds over procedures, weights, and teacher directions for deterministic horizons up to d^p.The noise estimates require no independence between the coordinate processes.Source: §10.4 · pp. 53–54

10.5 Proof of the non-escape theorem

• Combining the drift and noise bounds gives non-escape horizons of Td for tied and augmented procedures and rTd for the untied procedure, with probability 1 −o(1).The constants depend only on σ, δ, and κ.Source: §10.5 · p. 55

11 Conclusion

• Symmetry-aware and symmetry-agnostic feature learning have an r-fold sample-complexity separation for weak directional recovery when p ≥3 and weak subspace recovery for the pure quadratic Hermite link.Full-group data augmentation matches the statistical benefit of weight sharing in this setting.Source: §11 · p. 56
• For d^1/2 ≲r ≪d, the authors conjecture the separation persists, although a different proof mechanism is needed.The initial largest correlation no longer remains dominant until weak recovery in this regime.Source: §11 · p. 56
• Open directions include general links with information exponent two, strong recovery, vanishing prediction error, other losses, and partial-group augmentation.The authors note that squared loss introduces interactions between filters and subset augmentation raises a sample-efficiency versus computation tradeoff.Source: §11 · p. 56

AI use statement

• Generative AI assisted with mathematical statements, proof strategies, literature and references, manuscript exposition, analysis code, and plots; the authors state they take responsibility for the final work.The authors describe iterative discussion, critical review, and substantive revision of the mathematics and exposition.Source: AI use statement · p. 56

Teacher geometry

• The teacher orbit is nearly orthogonal, making a random direction’s correlations with its r shifted teacher directions approximately r independent Gaussian coordinates.The largest relevant correlation therefore reflects the usual extreme-value factor.Source: A Teacher geometry and initialization · p. 61

Initialization scale

• With probability 1-o(1), max_k |m_k(u)| ≤ Cγ_d and max_k q_k(u) ≥ cγ_d for a random initialization u.These bounds apply to tied initialization and separately to any fixed component of untied or augmented initialization.Source: A Teacher geometry and initialization · p. 62

Initialization gap

• With probability 1-o(1), the largest relevant correlation uniquely exceeds its closest competitor by a relative gap of order (log d)^-2.For even p, the comparison is between absolute correlations; the gap supplies the initial asymmetry later amplified by nonlinear SGD.Source: A Teacher geometry and initialization · pp. 61, 63

• The stochastic-gradient estimates use Section 6 notation and the teacher-geometry event of Lemma A.1.Source: B Proofs of stochastic gradient… · p. 67

B.1 Feature-gradient covariance

• The feature-gradient covariance proof computes the feature gradient, uses cyclic invariance and Gaussian integration by parts, and obtains the announced statement.Source: §B.1 · p. 67

B.2 Coordinate noise moments

• For coordinate noise, conditional on the past, each centered teacher-direction projection is a Gaussian polynomial of degree at most 2P, so hypercontractivity supplies higher-moment bounds.Source: §B.2 · p. 69

B.3 Uniform control of the stochastic gradient

• Uniform stochastic-gradient control combines a conditional second-moment bound with higher-moment estimates and a union bound over the update times.The proof uses conditional Markov’s inequality to control the gradient uniformly over t up to T ≤ d^C.Source: §B.3 · pp. 70–71