Neural Networks Provably Learn Spectral Representations for Group Composition

Jianliang He; Leda Wang; Fengzhuo Zhang; Siyu Chen; Zhuoran Yang

Source-linked AI summary

Neural Networks Provably Learn Spectral Representations for Group Composition

Jianliang He, Leda Wang, Fengzhuo Zhang, Siyu Chen, Zhuoran Yang

arXiv:2606.02993v2Jun 2, 2026cs.LGmath.OCmath.RTmath.STstat.ML

TL;DR

The paper asks how structured, low-dimensional representations emerge when neural networks learn finite-group composition. It analyzes projected gradient flow in the group Fourier domain and proves that neurons specialize in irreducible representations with rank-one cross-layer alignment, while Abelian ensembles diversify uniformly and achieve exact noise cancellation. The population-level theory does not yet cover train-test generalization or grokking.

Problem
The paper addresses whether networks trained on finite-group composition discover the groups’ irreducible representations and whether observed alignment and low-rank structure generalize beyond cyclic groups.
Method
The paper lifts projected gradient flow for a two-layer group-composition network into the Fourier domain and analyzes it as Riemannian gradient ascent on a representation-theoretic energy.
Results
For general finite groups, each neuron almost surely selects one irreducible representation and develops rank-one, rotationally aligned Fourier coefficients; for Abelian groups, representations and phases are uniformly distributed and the ensemble achieves perfect accuracy through exact noise cancellation.
Takeaways & Limitations
The theory provides a representation-theoretic account of feature learning and identifies rank-one compression as a spectral structure for matrix-valued group representations.
Takeaways & Limitations
The analysis focuses on the population case, leaving rigorous train-test-split analysis and the delayed generalization phenomenon known as grokking open.

Abstract

from arXiv · show

Understanding how structured internal structure emerges during neural network training is central to the study of deep learning. We investigate this phenomenon through the group composition task, where a two-layer neural network is trained to predict $g_1 \star g_2$ for elements of a finite group $G$. By lifting the projected gradient flow to the Fourier domain, we demonstrate that the training dynamics are governed by a Riemannian gradient ascent on a representation-theoretic energy functional. We prove that, under random initialization, this flow drives each neuron to converge almost surely toward a single irreducible representation, while the cross-layer Fourier coefficients achieve a rotational rank-one alignment. This framework provides a representation-theoretic account of feature learning and characterizes a novel low-rank compression phenomenon for matrix-valued group representations. Moreover, for Abelian groups, we provide a complete population-level description: random initialization promotes uniform diversification across nontrivial representations and induces Haar-uniform phases, jointly approximating the indicator via a majority-vote mechanism. We further prove that both phase alignment and representation competition emerge with exponential convergence rates.

1 Introduction

The paper studies how neural networks learn internal structure on finite-group composition tasks, asking whether irreducible representations and low-rank organization emerge from training. Using Fourier-domain gradient-flow analysis, it proves spectral specialization, rank-one cross-layer alignment, and a fuller Abelian characterization.

Motivation: The paper asks whether networks trained on arbitrary finite-group composition discover irreducible representations and whether modular-addition alignment and low-rank principles generalize.
Approach: The analysis uses harmonic analysis and lifts projected gradient flow for a two-layer network into the Fourier domain.The dynamics are characterized through a Riemannian gradient flow driven by a representation-theoretic energy functional.
General finite groups: With probability one under continuous random initialization, each neuron converges toward a single non-trivial irreducible representation.The result applies to every neuron under the population-risk gradient flow.
General finite groups: The surviving matrix-valued Fourier coefficients become rank one and align rotationally across layers, yielding a low-rank compression principle for non-Abelian representations.This extends spectral organization beyond cyclic groups and characterizes the alignment through products of the other layers’ coefficients.
Abelian groups: For Abelian groups, learned representations diversify uniformly over non-trivial irreducibles, while phases are independently Haar-uniform and the ensemble predictor combines them into a flawed indicator.The paper also establishes exponential convergence for phase alignment and representation competition.

2 Learning Group Composition with Neural Network

The paper formalizes finite-group composition as a complete-table prediction problem for a two-layer network with separate operand embeddings. A two-stage projected-gradient procedure first learns directional spectral features and then optimizes neuron scales.

Task: The task is to predict g1 ⋆g2 for every pair of elements in a finite group, with associativity, identity, and inversion defining the group operation.Abelian groups additionally satisfy commutativity, with modular addition as a canonical example.
Architecture: The network uses separate input embeddings for the two operands, output embeddings, positive neuron scaling factors, and a quadratic activation.Separate inputs preserve the distinction between left and right operands in non-commutative groups.
Training objective: Training uses cross-entropy on the complete composition table containing all pairs (g1,g2) and their labels g1 ⋆g2.The loss rewards the correct-class logit while penalizing the log-partition term across classes.
Training procedure: Stage I constrains directional parameters to the unit sphere and trains them by projected gradient flow with fixed scales; Stage II freezes directions and optimizes the scales.The factorization separates feature learning from later margin maximization.
Initialization: The model is initialized with independent uniform samples on unit spheres, while a small fixed Stage-I scale creates an approximately uniform-softmax small-logit regime.

3 Warmup: Learning Generalized Modular Addition

The generalized modular-addition warmup extends cyclic-group Fourier analysis to direct sums of cyclic groups. Experiments and observations show that training concentrates each neuron on one frequency pair, aligns phases, and diversifies frequencies and phases across the ensemble.

Generalized modular addition: Finite Abelian groups decompose into direct sums of cyclic groups, reducing generalized modular addition to component-wise modular addition.The framework includes ordinary modular addition and bitwise XOR as special cases.
Fourier representation: The group Fourier transform decomposes parameter functions into harmonic components indexed by frequency tuples, with real parameters producing conjugate Fourier pairs.The analysis focuses on the non-self-conjugate case where all cyclic factors are odd.
Learned patterns: Each neuron retains a single non-trivial frequency and its conjugate while other Fourier coefficients vanish.This frequency concentration is the generalized modular-addition form of single-frequency specialization.
Learned patterns: The output phase equals the sum of the two input phases, and scatter plots concentrate near the corresponding diagonal relation.The relationship is the phase-level analogue of cross-layer alignment.
Diversification: Surviving frequencies are uniformly distributed over non-trivial frequencies, while input phases are independently uniform across neurons and operand positions.The ensemble therefore covers conjugate frequency pairs nearly equally and exhibits uniform phase coverage.
Learned predictor: Together, the observations make each neuron depend on one representation while the ensemble combines all non-trivial representations into a predictor peaking at the correct composition.The resulting predictor also has secondary ghost peaks at the squaring elements.

4 Main Results for General Group Learning

For arbitrary finite groups, gradient flow learns structured spectral representations in two stages: neurons select single irreducible representations with rank-one rotational alignment, then scaling grows to drive loss toward zero.

Stage I: Representation Learning: Each neuron converges almost surely to a single non-trivial irreducible representation, with other Fourier components vanishing.This extends single-frequency specialization from Abelian groups to arbitrary finite groups.
Stage I: Representation Learning: Active Fourier coefficients become rank one and satisfy cross-layer proportionality relations expressing rotational alignment.In non-Abelian groups, the distinct matrix relations must be characterized separately because multiplication is non-commutative.
Stage I: Representation Selection: Random initialization favors higher-dimensional irreducible representations despite their higher potential energy, because initialization geometry and competition amplify their scale advantage.The learned proportion is reported to correlate positively with representation dimension d_ρ.
Stage II: Scale Growth: Theorem 4.5 establishes a second stage in which tied scaling factors grow logarithmically and cross-entropy loss converges to zero at rate O(1/T), under the perfect-accuracy condition.The theorem requires M ≳ log(|G|^3/δ) for probability at least 1−δ and gives a logarithmic lower bound on a(t).
Main Takeaway: The complete mechanism combines single-representation learning, rank-one alignment, logarithmic scale growth, and softmax sharpening across finite groups.For Abelian groups, the perfect-accuracy condition is proved; experiments confirm it for non-Abelian cases.
Experimental Verification: Experiments show block-sparse Fourier coefficients for neurons, perfect accuracy, alignment metric distal approaching 1, and low-rank metric distr1 approaching 0.The heatmaps concentrate each neuron in one irreducible-representation block, while training trajectories support the alignment and rank-one predictions.

5 Mechanism and Training Dynamics of Abelian Group

For Abelian groups, training separates phase alignment from representation competition, while diversified neurons collectively implement a flawed-indicator predictor through majority vote.

Diversification: Uniform initialization distributes learned nontrivial representations uniformly and phases independently according to the Haar measure on the unit circle.The limiting measure is the push-forward of the product distribution Unif(Irr(G)̸=1) ⊗ Haar(D).
Neuron-level mechanism: Each neuron retains one nontrivial irreducible representation, with rank-one Fourier coefficients and aligned phases across layers.The resulting representation-theoretic structure is summarized as single-frequency specialization with aligned phases.
Ensemble mechanism: Averaging diversified neurons cancels noise terms while signal terms combine coherently, producing a closed-form majority-vote predictor.The correct label receives the largest logit, although the predictor is a flawed rather than perfect indicator.
Scale dynamics: The scaling parameter a grows logarithmically, sharpening the softmax output toward a one-hot prediction after the representation and phase structure forms.Figure 8 shows an initial loss plateau followed by a rapid drop toward 0 during Stage II.
Convergence: Under aligned phases, the initially largest-magnitude irrep wins, while all other irreducible representations decay.With uniform initialization, exchangeability makes each nontrivial irrep equally likely to win, supplying the microscopic basis for ensemble diversification.
Convergence: Phase alignment and representation competition are distinct training subprocesses, and both converge exponentially fast.The winning irrep is selected by the largest initial Fourier magnitude and amplified while competitors decay.

6 Conclusion

The paper develops a harmonic-analysis theory of group-composition learning, proving spectral organization for finite groups and a complete characterization for a restricted Abelian setting. Its scope remains limited by unresolved general-group distributions, self-conjugate representations, and train-test generalization.

Conclusion: The paper proves that gradient flow learns single irreducible representations with rank-one cross-layer alignment for finite groups, including non-Abelian groups.This establishes a representation-theoretic account of feature learning and low-rank structure in matrix-valued coefficients.
Open problems: Limiting distributions remain open for general groups with high-dimensional irreducible representations and Abelian groups containing self-conjugate representations.These are identified as unresolved theoretical cases.
Open problems: The paper studies the population case, leaving rigorous train-test analysis and the delayed generalization phenomenon known as grokking open.The conclusion specifically identifies train-test split analysis as an unresolved direction.
Scope and assumptions: The analysis uses a small-logit Stage I approximation in which the exact cross-entropy risk is replaced by a first-order approximate risk.The approximation relies on small scaling a, near-zero logits, and an approximately uniform softmax distribution.

A.2 Proof of Proposition 4.2

The proof lifts projected gradient flow into Fourier space and derives equivalent dynamics for matrix-valued Fourier coefficients. These dynamics are identified with Riemannian gradient ascent on an energy functional.

The group DFT converts projected gradient flow into a dynamical system for each neuron’s matrix-valued Fourier coefficient.The derivation uses Fourier expansions, Schur orthogonality, and the chain rule.
The spectral dynamics contain cross-layer coupling through products of the other two layers’ Fourier coefficients and projection terms enforcing sphere constraints.The same structure applies cyclically to the three layers.

B.2 Proof of Theorem 4.3: Representation Learning in Stage I

The proof analyzes each neuron as a Riemannian gradient ascent flow on a compact spectral manifold. It establishes the geometric formulation needed to characterize its limiting equilibria.

B.2.1 Step 1: Riemannian Lifting to the Spectral Manifold: Each neuron’s Fourier coefficients evolve on a product of unit-sphere constraints, forming the spectral manifold M.The original unit-norm constraints translate into Fourier-space constraints through Plancherel’s theorem.
B.2.1 Step 1: Riemannian Lifting to the Spectral Manifold: The projected Fourier dynamics coincide with Riemannian gradient ascent of the energy functional Ω on M.The representation basis makes the ambient Hessian block diagonal across irreducible representations.

B.2.2 Step 2: Critical Point Classification

The critical-point classification separates exceptional or saddle equilibria from the positive-energy rank-one states that remain as possible generic limits.

Zero-energy equilibria supported only on the trivial representation can attract trajectories only from a measure-zero initialization set.The exceptional set is a proper embedded submanifold of M.
Zero-energy equilibria with non-trivial representation support are strict saddles because their block structure yields a positive Hessian direction.Their Fourier ranks satisfy pairwise constraints induced by triple-annihilation relations.
Positive-energy equilibria with total rank at least two are strict saddles, leaving only positive-energy rank-one single-representation equilibria.The Hessian has a positive eigenvalue along suitable tangent directions.
Riemannian gradient flow avoids all strict saddles almost surely under absolutely continuous initialization.The stable manifold of the saddle set has zero Riemannian volume.

B.2.4 Assembly: Proof of Theorem 4.3

The assembly combines convergence on the analytic spectral manifold with the critical-point classification and saddle avoidance. Consequently, generic trajectories converge to the desired rank-one single-representation equilibria.

Compactness and real analyticity ensure that each gradient-flow trajectory converges to a single critical point rather than drifting along a continuum.The Łojasiewicz-based argument upgrades bounded energy and vanishing gradient to finite trajectory length.
The spectral dynamics decouple across neurons, so the proof analyzes each neuron independently on the compact real-analytic manifold M.The energy is a real polynomial and therefore real analytic on M.
Cases 1–2 require initialization in a measure-zero exceptional set, while Cases 3–4 are strict saddles avoided almost surely.Thus random initialization excludes all four undesired critical-point classes with probability one.
Almost surely, the flow converges to positive-energy rank-one equilibria supported on a single irreducible representation, yielding rotational rank alignment across layers.The equilibrium conditions establish the alignment relations and positive proportionality.

B.2.7 Proof of Lemma B.8: Case 3

The proof analyzes non-target equilibria through orthogonal block factorizations and constructs positive-Hessian directions, showing they are strict saddles. Consequently, only positive-energy, rank-one, single-representation equilibria remain relevant.

Steps 2–3: Hessian Analysis: For zero-energy equilibria with nontrivial representation support, triple-annihilation relations yield a positive Hessian direction.The construction proceeds by verifying tangency, simplifying the Hessian, and explicitly constructing an ascent eigenvector.
Step 1: Orthogonal Block Structure: The proof decomposes equilibrium Fourier coefficients into mutually orthogonal column and row spaces with shared rank.These spaces support compatible partial-isometry and block SVD factorizations.
Step 3: Eigenvector Construction: The proof constructs tangent eigenvectors with positive eigenvalues in both available-degree-of-freedom cases.Case 1 uses a free orthogonal block, while Case 2 pairs singular directions when no free block exists.
Conclusion: Therefore, every analyzed non-target equilibrium is a strict saddle, leaving positive-energy rank-one single-representation equilibria.These surviving equilibria yield the single-representation property and rank alignment with positive proportion.
Step 1: Orthogonal Block Structure: Each equilibrium coefficient is represented through shared partial-isometry bases and invertible core matrices.The core matrices are then shown to be proportional to unitaries.

B.2.9 Proof of Theorem B.11: Saddle Avoidance

The proof adapts center-stable manifold theory to Riemannian gradient flows and shows that strict-saddle stable sets have zero volume. Thus absolutely continuously randomized initialization avoids saddle convergence almost surely.

Flow Regularity: On a compact Riemannian manifold, the gradient flow exists for all time and each flow map is a diffeomorphism.Completeness follows from smoothness and compactness, while uniqueness supplies the inverse flow.
Center-Stable Manifold: Near a critical point, the non-positive Hessian eigenspace integrates into a local center-stable manifold.Trajectories remaining in a sufficiently small neighborhood must lie on this manifold.
Saddle Geometry: Positive Hessian eigenvalues define unstable directions that repel trajectories from a strict saddle.The local center-stable manifold therefore has lower dimension than the ambient manifold.
Global Saddle Avoidance: The global stable set of all strict saddles is contained in a countable union of lower-dimensional flow preimages and has zero volume.The countable-cover argument combines local center-stable manifolds with diffeomorphic flow maps.
Conclusion: With absolutely continuous random initialization, the gradient flow avoids convergence to every saddle point with probability 1.This is the saddle-avoidance conclusion of Theorem B.11.

B.3 Proof of Theorem 4.5: Growth Rate of Scaling Factor in Stage II

The proof reduces Stage II training to a scalar scaling-factor flow controlled by the predictor’s logit margins. A positive population margin transfers to finite width, yielding logarithmic scale growth and a cross-entropy convergence bound.

Step 2: Logarithmic Scale Growth: A positive margin lower-bounds the scale derivative and produces logarithmic growth of a(t).The proof obtains this by bounding incorrect-label softmax weights and integrating the resulting differential inequality.
Step 2: Logit Margin: The population predictor assigns the highest logit to the correct group-composition label, creating a positive population margin.Finite-width concentration then transfers this margin to all input pairs with high probability.
Step 1: Scalar Gradient Flow: Under tied output scaling, the network factorizes as a scalar a times a fixed finite-width predictor.Differentiating cross-entropy makes the scale derivative a softmax-weighted sum of correct-versus-incorrect logit gaps.
Step 3: Loss Convergence: To achieve R(T) ≤ ϵ, it suffices that T ≳ |G|/ϵ · (1 + (|G| −1)^−2).This follows after requiring a(T) ≳ log(|G|^2 · (|G| −1)/ϵ).
Abelian Magnitude–Phase Reduction: The Abelian analysis separates each complex Fourier coefficient into magnitude and phase and reduces phase dependence to relative alignment.The reduced variables are (αθ,m[ρ], αξ,m[ρ], φm[ρ]), whose dynamics form a closed system.

C.2 Proof of Theorem 5.1: Verification of (µ-PA) for Abelian Group

For finite Abelian groups, the limiting neuron distribution is uniform across nontrivial irreducible representations with Haar-uniform phases. Its induced predictor achieves perfect accuracy by giving the correct composition label the largest coefficient.

Limiting Distribution: The limiting parameter distribution is the inverse-Fourier pushforward of Unif(Irr(G)\{1}) ⊗ Haar(D).Permutation invariance across irreducible representations and rotational phase invariance yield the product measure.
Flawed Indicator: The correct label receives coefficient 2, two ghost labels receive coefficient 1 each, and all remaining labels have baseline −4/|G|.Despite the ghost-label weights, the correct label remains the largest coefficient.
Perfect Accuracy: This limiting distribution satisfies the perfect-accuracy condition for every input pair.The proof verifies the condition by computing the resulting population predictor.

C.2.1 Proof of Lemma C.3: Limiting Distribution

The proof characterizes each neuron's limiting state by a winning non-trivial irreducible representation and a limiting phase. Symmetry of initialization and dynamics yields uniform, independent representation–phase pairs, whose ensemble predictor favors the correct composition label.

Limiting structure: Each neuron converges toward a non-trivial irreducible representation, with all other representation magnitudes vanishing and the winner’s phase aligning to one.The limiting neuron is represented by the pair (ˇρm, um).
Limiting distribution: The limiting winner and phase are independent, with um ∼Haar(D), ˇρm ∼Unif(Irr(G)̸=1), and i.i.d. pairs across neurons.The conditional phase law remains Haar-uniform after conditioning on the reduced state.
Initialization: Uniform spherical initialization makes Fourier phases Haar-uniform and independent of representation magnitudes and relative phases.The DFT preserves the rotational symmetry of the initialization.
Representation selection: Permutation invariance of initialization and equivariance of the reduced ODE force every non-trivial irreducible representation to be equally likely to win.The symmetry acts on conjugate representation pairs and is preserved by the flow.
Population predictor: Averaging over representations and phases cancels noise terms and makes the correct label g1 ⋆g2 attain the largest expected logit.Representation orthogonality supplies the indicator contribution in the ensemble predictor.

C.3 Proof of Theorem 5.3: Convergence Rate for Abelian Groups

The Abelian-group dynamics reduce to coupled magnitude and phase equations under scale matching. Phase alignment and representation competition are analyzed separately, and both mechanisms exhibit exponential convergence toward a dominant aligned representation.

Reduced dynamics: Scale-matching initialization preserves equal input and output Fourier magnitudes for every representation throughout training.This permits the common magnitude variable αm[ρ](t) := |c θm[ρ](t)| = |c ξm[ρ](t)|.
Coupled mechanisms: Magnitude dynamics combine self-reinforcement for large, well-aligned representations with competition through the energy term, while phase alignment drives φm[ρ] toward 1.The phase-alignment rate is proportional to αm[ρ].
Convergence rates: The convergence analysis treats phase alignment and representation competition separately, with explicit time requirements for both mechanisms and suppression of the trivial representation.The energy is monotone non-decreasing along the Riemannian gradient-ascent flow.
Phase alignment: Under single-representation initialization, the single-frequency support is preserved and phase alignment converges exponentially fast.The dynamics reduce to a scalar phase ODE after the inactive representations remain zero.
Representation competition: Under perfectly phase-aligned initialization, phase alignment remains invariant and the initially dominant representation stays dominant as magnitude ratios diverge.The competition analysis uses log-ratio dynamics among representation magnitudes.

F Additional Results for Generalized Modular Addition

The generalized modular-addition results extend the spectral picture to product groups and self-conjugate representations. Most neurons retain conjugate-pair sparsity, while self-conjugate modes replace complex phases with real signs.

Odd-order product groups: On G = Z3 ⊕Z5, each neuron concentrates on one shared non-trivial frequency tuple and its conjugate partner across all three parameter families.Hermitian symmetry gives equal real parts and opposite imaginary parts for the active pair.
Self-conjugate representations: For G = Z2⊕Z3⊕Z5, 29 non-trivial representations consist of one non-trivial self-conjugate frequency and 14 conjugate pairs.The self-conjugate frequency is k = (1, 0, 0).
Learned structure: Most neurons still learn a single non-trivial frequency and its conjugate partner, while a neuron selecting the self-conjugate mode has one real active coefficient.The self-conjugate coefficient has no distinct conjugate partner and vanishing imaginary part.
Phase interpretation: Self-conjugate modes replace the continuous Haar-uniform phase with a discrete Rademacher sign in {±1}, without introducing a new learning mechanism.This is the boundary case of the same spectral framework.
Relation to prior analyses: The framework generalizes cyclic modular-addition analyses from one-dimensional characters to arbitrary finite Abelian products and even-order groups.For even-order groups, self-conjugate sectors are handled by real-valued signs rather than complex phases.