📊 statistics

Neural Networks Provably Learn Spectral Representations for Group Composition

This paper proves that two-layer neural networks trained on finite group composition tasks provably learn spectral representations by converging to irreducible representations with exponential rates, driven by a Riemannian gradient ascent on a representation-theoretic energy functional that induces low-rank compression and feature diversification.

Original authors: Jianliang He, Leda Wang, Fengzhuo Zhang, Siyu Chen, Zhuoran Yang

Published 2026-07-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jianliang He, Leda Wang, Fengzhuo Zhang, Siyu Chen, Zhuoran Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are watching a team of tiny, digital detectives try to solve a massive, complex puzzle. This isn't a mystery about who stole the cookies; it's a puzzle about how computers learn to understand the hidden rules of the universe. In the world of artificial intelligence, we often wonder: when a neural network (a computer brain made of layers of math) gets really good at a task, what does it actually "learn" inside its head? Does it just memorize answers, or does it discover deep, elegant structures? This paper dives into that question by giving the computer a very specific, mathematical game to play: learning how to combine things according to the rules of a "group."

To understand the game, you need to know what a "group" is. Think of a group as a set of moves or objects that follow strict rules. For example, imagine a clock face. If you move the hand forward by 3 hours and then by 4 hours, you end up in the same spot as if you moved it forward by 7 hours. The rules of how these moves combine are consistent and predictable. In math, this is called a "group composition." The researchers wanted to see if a neural network, when trained to predict the result of combining any two moves in such a group, would naturally discover the secret "language" that describes these rules. That language is called "representation theory," which is basically a way of breaking down complex patterns into simple, fundamental building blocks, much like how a prism breaks white light into a rainbow of colors.

The paper, titled "Neural Networks Provably Learn Spectral Representations for Group Composition," takes a two-layer neural network and trains it on this group-combination game. The researchers didn't just watch the network learn; they used advanced math to prove exactly how it learns. They found that the network doesn't just guess; it organizes itself in a very specific, beautiful way.

Here is what they discovered. When the network starts, its internal parts (called neurons) are like a chaotic crowd, all trying to do everything at once. But as the training happens, something magical occurs. Each neuron stops trying to be everything and decides to specialize in just one specific "frequency" or pattern. In the world of math, these patterns are called "irreducible representations." It's as if every neuron in the crowd picks a single instrument to play, and they all agree on the exact same note.

But it gets even more interesting. The paper proves that these neurons don't just pick a note; they align perfectly with each other. The researchers showed that the network compresses its complex, multi-dimensional data into a "rank-one" structure. Imagine a tangled ball of yarn that suddenly untangles itself into a single, straight, perfect thread. This happens for every neuron, and they all line up in a specific rotational order, like dancers in a synchronized routine.

The study also looked at what happens when the group is "Abelian," which is a fancy word for groups where the order of operations doesn't matter (like adding numbers: 2 + 3 is the same as 3 + 2). In this case, the researchers proved that the network doesn't just pick one pattern; it picks all the possible patterns, but in a perfectly fair way. Every possible "note" gets played by a different neuron, and their phases (the timing of their notes) are spread out evenly, like a perfect circle of dancers. This creates a "majority vote" system where the noise cancels out, and the correct answer pops out clearly.

The authors proved that this happens with almost total certainty, provided the network starts with random settings. They showed that the network avoids getting stuck in bad spots and naturally flows toward this perfect, organized state. They also found that this learning happens in two distinct stages. First, the network figures out the right patterns and aligns them (the "feature learning" stage). Second, it turns up the volume on these patterns (the "scaling" stage) to make the final answer super clear and accurate.

In short, this paper proves that when you teach a neural network to understand the rules of combining things, it doesn't just memorize. It discovers the fundamental, spectral "music" of those rules, organizing itself into a highly efficient, low-rank, and perfectly aligned structure. It's a mathematical guarantee that these digital brains are capable of finding deep, elegant order in the chaos of data.

Technical Summary: Neural Networks Provably Learn Spectral Representations for Group Composition

Problem Statement
The paper investigates the emergence of structured internal representations during the training of neural networks, specifically focusing on the task of learning group composition. Given a finite group $(G, \star)$ , the objective is to train a two-layer neural network to predict the product $g_1 \star g_2$ for any pair of elements $(g_1, g_2) \in G \times G$ . While prior work has established that networks learning modular addition over cyclic groups ( $\mathbb{Z}_p$ ) discover Fourier features (single-frequency specialization), the behavior of networks learning arbitrary finite groups—particularly non-Abelian groups where representations are matrix-valued rather than scalar—remained theoretically uncharacterized. The central question is whether neural networks discover the irreducible representations (irreps) of a general group $G$ and how the principles of phase alignment and low-rank structure observed in Abelian settings generalize to this broader context.

Methodology
The authors analyze a two-layer fully connected neural network with quadratic activation ( $\sigma(x) = x^2$ ) trained on the complete composition table of a finite group $G$ using cross-entropy loss. The network architecture employs separate input embeddings for the two operands to accommodate non-commutative operations.

The theoretical analysis proceeds in two stages:

Stage I (Feature Learning): The scaling factors of the neurons are fixed at a small constant $a$ , and the directional parameters (input and output embeddings) evolve via projected gradient flow on the unit sphere. The authors lift the training dynamics from the Euclidean parameter space to the spectral domain using the Group Discrete Fourier Transform (DFT). They demonstrate that the projected gradient flow is equivalent to a Riemannian gradient ascent on a specific energy functional $\Omega_m$ defined over the manifold of Fourier coefficients.
Stage II (Scale Optimization): Once the directional parameters converge, the scaling factors are optimized to maximize the margin, driving the cross-entropy loss to zero.

The analysis relies heavily on harmonic analysis on finite groups, specifically the orthogonality of irreducible representations and Schur's orthogonality relations. The authors employ tools from dynamical systems, including the classification of critical points (equilibria) and saddle avoidance arguments extended to continuous Riemannian gradient flows, to prove convergence properties.

Key Contributions and Results

General Group Representation Learning (Theorem 4.3): For arbitrary finite groups (including non-Abelian), the authors prove that under random initialization, the gradient flow drives each neuron to converge almost surely to a state characterized by:
- Single Representation: Each neuron selects exactly one non-trivial irreducible representation $\check{\rho}_m$ (and its dual $\check{\rho}_m^\vee$ ), causing all Fourier coefficients corresponding to other representations to vanish.
- Rank-One Rotational Alignment: The Fourier coefficients of the surviving representation become rank-one matrices. Furthermore, these coefficients across the input and output layers exhibit a specific rotational alignment: $\hat{\xi}_m[\rho] \propto \hat{\theta}_{2m}[\rho] \hat{\theta}_{1m}[\rho]$ , and cyclic permutations thereof. This generalizes the scalar phase alignment observed in modular addition to matrix-valued representations.
- Mechanism: The proof establishes that the training dynamics maximize an energy functional. Non-target equilibria (negative energy, zero energy with trivial support, or higher-rank positive energy) are shown to be either measure-zero traps or strict saddles, which are avoided almost surely by the gradient flow.
Abelian Group Characterization (Theorem 5.1 & 5.3): For Abelian groups, the authors provide a complete population-level description:
- Diversification: Under uniform spherical initialization, the learned representations are uniformly distributed over all non-trivial irreps, and the phases are Haar-uniform on the unit circle, independent of the chosen representation.
- Flawed Indicator: The resulting ensemble predictor acts as a "flawed indicator," assigning the highest logit to the correct group product $g_1 \star g_2$ , with secondary "ghost" peaks at $g_1^2$ and $g_2^2$ . This structure arises from the exact cancellation of noise terms via the orthogonality of representations and the averaging over Haar-uniform phases.
- Convergence Rates: The paper establishes explicit exponential convergence rates for both phase alignment and representation competition. The selection of the winning representation follows a "lottery ticket" mechanism, determined by the initial spectral magnitudes.
Stage II Dynamics (Theorem 4.5): The authors prove that if the mean-field predictor satisfies a perfect accuracy condition (which holds for Abelian groups and is empirically verified for non-Abelian cases), the scaling factors grow logarithmically, sharpening the softmax distribution and driving the loss to zero at a rate of $O(1/T)$ .

Significance and Claims
The paper claims to provide the first rigorous theoretical characterization of how neural networks learn group composition for general finite groups. By unifying the phenomenon of "single-frequency specialization" in cyclic groups with the broader framework of representation theory, the authors demonstrate that gradient-based training naturally selects irreducible representations and compresses the learned matrix-valued coefficients to rank one.

The significance lies in:

Bridging Mechanistic Interpretability and Representation Theory: The work formalizes the link between emergent neural features and the mathematical structure of the underlying group, showing that the "Fourier features" in modular addition are simply the scalar case of learning irreducible representations.
Novel Low-Rank Compression: The discovery of rank-one rotational alignment in matrix-valued representations for non-Abelian groups is a novel structural finding that extends the low-rank bias observed in other deep learning contexts to the spectral domain of group operations.
Dynamical Explanation: The paper moves beyond empirical observation to provide a dynamical systems-based proof (via Riemannian gradient flow and saddle avoidance) explaining why and how these structured solutions emerge from random initialization.

The authors note that while their results cover general finite groups, open problems remain regarding the theoretical characterization of limiting distributions for groups with high-dimensional irreps or self-conjugate representations, as well as the analysis of generalization gaps (grokking) in finite-sample settings.

More like this