Neural Networks Provably Learn Spectral Representations for Group Composition

Abstract: 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.

• The paper proves that gradient descent induces Fourier spectral specialization, with each neuron converging to a single nontrivial irreducible representation.
• It shows that training leads to rank-one compression and cross-layer rotational alignment, underpinning the network's algebraic structure learning.
• Empirical results across Abelian and non-Abelian groups validate the theory, confirming the network’s ability to solve group composition tasks.

Provable Spectral Feature Emergence in Neural Networks for Group Composition

Introduction and Background

This paper addresses the precise internal structure learned by neural networks when trained to model group composition tasks over arbitrary finite groups. Specifically, it establishes that two-layer neural networks trained to predict $g_1 \star g_2$—where $\star$ denotes the group operation—organize their internal representations along the irreducible representations (irreps) of the group. The analysis proceeds by lifting the parameter dynamics to the group Fourier (representation-theoretic) domain, revealing that gradient-based training induces emergent spectral structure, rank compression, and layer alignment in the Fourier basis. These phenomena generalize empirical observations made previously in modular addition ($\mathbb{Z}_p$ arithmetic) to all finite groups, both Abelian and non-Abelian.

The approach builds on the established link between group structure and harmonic analysis, as exemplified by the group Fourier transform. Recent mechanistic interpretability literature has observed that neural networks solving modular arithmetic often encode solutions in Fourier-aligned weights, but formal theory verifying and generalizing these patterns remained incomplete. This work closes the gap by providing a full characterization, showing that gradient flow in the neural parameter space amounts to Riemannian gradient ascent on a representation-theoretic energy in spectral space.

Architecture, Training Regime, and Fourier Lifting

The task considered is the prediction of group products from input pairs $(g_1, g_2)$ over group $G$, where non-Abelian structure demands separate embeddings for left and right operands. The network is a two-layer MLP with quadratic activation; weights $\theta_m^1$, $\theta_m^2$, $\xi_m$ are vector-valued, with $m$ indexing the neuron. The cross-entropy risk, computed over the complete group composition table, is minimized in two phases: feature learning via projected gradient flow on the sphere for the direction parameters, followed by a scale-maximization phase for per-neuron coefficients $a_m$.

Crucially, for mathematical analysis, the parameter dynamics are “lifted” via group Fourier transform. This reparametrizes weights $\star$0 as functions over $\star$1 and decomposes them via irreducible unitary representations $\star$2. For $\star$3 non-Abelian, the basis is block-matrix-valued with sizes equal to the irrep dimension $\star$4, consistent with the orthogonality relations fundamental to representation theory.

Figure 1: Cayley graphs and spectral basis heatmaps illustrating the algebraic and representation-theoretic structure of $\star$5 (Abelian) and $\star$6 (non-Abelian) groups.

Main Theoretical Results: Spectral Specialization, Rank Compression, and Alignment

The central theorem demonstrates: with probability one (over random initialization), each neuron converges to a single nontrivial irrep, suppressing all others, and the surviving matrix-valued Fourier coefficients become strictly rank-one and rotationally aligned across layers.

These findings are encapsulated as:

• Single Irreducible Representation: For each neuron $\star$7, only one $\star$8 and its dual (in the conjugate class) have nonzero weights; all other Fourier blocks trend to zero.
• Rank-One Rotational Alignment: Within the active $\star$9, the Fourier coefficients at all network layers are rank-one, and the cross-layer matrices satisfy rotational proportionality, e.g.,

$\mathbb{Z}_p$0

for some $\mathbb{Z}_p$1; analogous cyclic relations hold for permutations of the three.

• Stability and Saddle Avoidance: The dynamics are characterized as Riemannian gradient ascent on a compact manifold (the product of spheres for each parameter in Fourier space), with strict saddle points and measure-zero traps almost surely avoided via generalizations of the stable-manifold theorem.

Figure 2: Riemannian gradient flow evolving on the constrained parameter manifold; illustration of tangent space splitting near strict saddles relevant for saddle-avoidance arguments.

• Abelian Case: Diversification and Uniform Phase Distribution: For Abelian $\mathbb{Z}_p$2, the winning irrep is uniformly distributed, and the absolute phase is Haar-random and independent of the representation, leading to uniform spectral coverage (“diversification”), enabling exact cancellation of noise terms in the final predictor.

Empirical Verification

Extensive experiments are provided, notably over the Frobenius group $\mathbb{Z}_p$3 for non-Abelian cases and over direct products like $\mathbb{Z}_p$4 for Abelian cases.

Figure 3: Block-sparse structure in DFT domain—each neuron concentrates on a single conjugate irrep pair, with all other blocks negligible, validating the single-representation prediction.

• Block-sparse spectral structure: For every neuron, DFT heatmaps reveal high magnitude only in the active irrep block and its conjugate, the remainder being suppressed.
• Layer alignment and rank collapse: As training proceeds, cross-layer cosine similarities approach 1 (perfect alignment), and the ratio of singular values in the active matrix blocks drops to zero, evidencing strict rank-one structure.

Figure 4: Accuracy curves (left) showing convergence to perfect test accuracy, and metrics (center, right) confirming rotational alignment and singular value collapse to rank-one structure across neurons.

• Two-stage dynamical regime: Stage I yields the spectral structure; scale growth in Stage II sharpens softmax prediction, and cross-entropy loss drops logarithmically per the rates predicted.

Figure 5: Loss trajectory (left) with sharp drop only after scale maximization begins; scaling factor growth under tied and untied (per-neuron) scenarios both demonstrate the expected logarithmic growth.

• Dynamics in Abelian setting: Phase-alignment and representation competition (the “lottery ticket” mechanism) proceed exponentially quickly, with the ultimate winner dictated by initial spectral magnitudes.

Figure 6: (Left) Phase alignment trajectories converge exponentially at neuron-dependent rates. (Right) Intra-neuron competition among irreps clearly demonstrates the lottery mechanism, with the winner being established early and competitors exponentially decaying.

• Visualization of individual neuron weights: DFT projections for neurons in the Abelian group case confirm single-frequency specialization with the expected Hermitian structure in the learned parameters.

  Figure 7: DFT heatmaps for neurons trained on $\mathbb{Z}_p$5, showing single active frequencies and their conjugates.

  

  Figure 8: Extension to $\mathbb{Z}_p$6 again reveals strict frequency sparsity and conjugate-pair structure in the learned weights across neurons.

Mechanistic and Practical Implications

• Theoretical analysis confirms that neural networks trained with gradient descent do not rely on arbitrary superposition of features but coordinate at the spectral level: neurons specialize, compress, and align predictably in the group’s harmonic basis.
• In Abelian cases, diversification and phase randomness enforce ensemble “majority vote” solutions, yielding predictors with noise-cancelling properties and robust correct-label selectivity—characterized by a “flawed indicator” function peaked at the true group product.
• In the non-Abelian case, the same spectral organization emerges, but the underlying irreps are no longer scalars but low-rank matrices, with rank compression acting as a “low-dimensional bottleneck” for high-dimensional representations.

The unification of these mechanisms under a single theoretical framework suggests a robust intrinsic bias in deep learning toward the discovery and compression of compositional algebraic structure in group-theoretic or equivariant regimes. This has direct implications for the analysis and design of neural models tasked with algebraic, combinatorial, or symmetry-exploiting problems.

Outlook for Future Research

The results suggest several future directions:

• Extending to higher-order or infinite groups: Characterizing learning and spectral dynamics in more complex group settings, continuous or Lie groups.
• Investigating generalization and “grokking”: The analysis focuses on the population loss; extending the theory to understand train-test split behaviors and the onset of generalization remains open.
• Connections to equivariant network design: The results explain why architectures enforcing equivariance and spectral sparsity are efficient; the theory may help refine architectural choices for group-aware machine learning.

Conclusion

This paper provides a rigorous, comprehensive theory and empirical verification for the representation-theoretic organization of neural networks solving group composition. It establishes a two-stage feature learning and scale-maximization process with spectral specialization, rank-one compression, and cross-layer alignment as necessary outcomes, all operating in the Fourier domain for arbitrary finite groups. The results unify numerous empirical phenomena within a provable mechanistic account and offer explicit predictions for a broad class of deep learning models with structured, compositional, or symmetric data (2606.02993).