Neural Networks Provably Learn Spectral Representations for Group Composition

Neural Networks Discover Group Irreducible Representations via Riemannian Gradient Flow

Researchers are trying to understand how useful internal structure emerges during neural network training. When data contains latent regularities, what internal representations do neural networks actually extract? How do those representations support computation and generalization? This is a central mystery in mechanistic interpretability (the field dedicated to reverse-engineering the "circuits" inside trained models).

Previously, much of this work focused on simple cyclic groups, such as modular addition ($x + y \pmod p$). Studies showed that networks trained on these tasks tend to extract trigonometric features aligned with the discrete Fourier basis. However, a massive gap remained. We did not know if these patterns were just artifacts of simple math. We needed to know if they generalized to more complex, non-commutative structures. The big question was whether a network learning to compose elements of an arbitrary finite group $G$ would discover the fundamental building blocks of that group, known as irreducible representations (irreps; the simplest, indivisible components of a group's symmetry).

This paper provides a definitive answer. It proves that for any finite group, the training process forces neurons to specialize in specific irreducible representations through a mathematically predictable optimization path.

Beyond the limits of modular addition

Current mechanistic interpretability often hits a wall when moving from Abelian groups (where $a \star b = b \star a$) to non-Abelian groups (where the order of operations matters). In the Abelian case, the "features" a network learns are scalar-valued frequencies. These are simple sinusoids that are easy to visualize and track.

But for non-Abelian groups, the underlying symmetries are matrix-valued. The representations are no longer single numbers. They are multidimensional objects where multiple coefficients are coupled together. Existing frameworks struggle to describe this transition. When we move to more complex algebras, we lose the ability to say "this neuron is learning frequency $k$." We lack a principled way to describe how a neuron organizes its internal matrix structure. Without this, we cannot tell if a network is truly learning the group's structure or just interpolating the training data.

The Riemannian ascent on spectral manifolds

The authors tackle this by modeling the training of a two-layer neural network through the lens of harmonic analysis. They "lift" the entire training process into the Fourier domain. The core of their approach relies on a two-stage training procedure. This design isolates feature learning from scale optimization.

1. Stage I (Feature Learning): The directional parameters (the embeddings $\theta_1, \theta_2,$ and $\xi$) are constrained to the unit sphere. During this phase, the scaling factors $a_m$ are held fixed. The authors prove that this projected gradient flow is equivalent to a Riemannian gradient ascent on a representation-theoretic energy functional $\Omega_m$. Essentially, the network is climbing a landscape where the "peaks" correspond to states that perfectly align with the group's irreducible representations.
2. Stage II (Margin Maximization): Once the directions are learned, the weights are frozen. Only the scaling factors $a_m$ are optimized. This stage acts as a "sharpener." It pushes the logits (the raw scores before softmax) toward the correct labels to minimize cross-entropy loss.

The architectural choice of using quadratic activations ($\sigma(x) = x^2$) is critical here. It allows for a clean spectral decomposition. This enables the authors to map the Euclidean gradient descent directly onto a manifold of Fourier coefficients. As seen in, this allows the transformation of abstract group operations into structured matrix-valued representations.

Converging to rank-one rotational alignment

The paper's primary strength lies in its rigorous convergence proofs. The authors report that for any arbitrary finite group, the gradient flow drives each neuron to converge almost surely toward a single irreducible representation and its conjugate.

More importantly, they identify a novel phenomenon called "Rank-one Rotational Alignment." In non-Abelian settings, the Fourier coefficients are matrices. The authors demonstrate that these matrices do not remain high-rank. Instead, they compress. They find that the surviving matrix-valued coefficients become rank-one. Furthermore, these coefficients across the different layers ($\theta_1, \theta_2, \xi$) become proportional to one another in a specific rotational order. This is empirically verified in [Figure 6c]. There, the singular value ratio $\sigma^{(2)}/\sigma^{(1)}$ decays toward zero. This decay signals that the matrix has collapsed into a low-rank structure.

For the simpler Abelian case, the results are equally striking. The authors show that the network achieves perfect accuracy through a "majority-vote" mechanism. Because the neurons are diversified across all possible representations and their phases are distributed uniformly (as shown in ), the individual "noise" from each neuron cancels out in the ensemble.

This leaves only the coherent signal of the correct group product.

Blind spots in the spectral theory

While the theoretical framework is robust, there are significant gaps that prevent immediate application to large-scale production models.

First, the proofs rely heavily on the "small-logit regime" assumption. The authors approximate the cross-entropy loss using a Taylor expansion. This assumes the initial output scale $a$ is sufficiently small. While they provide a bound in Proposition 4.2 to show the error is controlled, this might not hold in real-world training runs. In those cases, logits can fluctuate wildly.

Second, the theory focuses exclusively on the "population case." This means training on the complete composition table of the group. In practice, engineers deal with sampling, noise, and the train-test split. The paper acknowledges that the "grokking" phenomenon—where a model suddenly generalizes after a long period of overfitting—is not addressed here. We do not yet know how these spectral patterns emerge when the network only sees a fraction of the group elements.

Finally, the paper lacks a characterization of groups with very high-dimensional irreducible representations. The complexity of the "rotationally aligned" matrices increases with the dimension of the irrep. The current proofs do not fully capture the limiting distributions for these extreme cases.

The verdict: A blueprint for geometric deep learning

Is this worth your time? If you are working on geometric deep learning or designing architectures meant to be equivariant to specific symmetries, the answer is yes. This isn't just a theoretical curiosity. It is a proof that the internal structure of a network handling algebraic tasks is mathematically predictable.

The paper provides a rigorous justification for why certain architectural choices matter. For example, it shows why quadratic activations or specific embedding structures are effective for feature learning. It moves the conversation from "my model seems to learn symmetries" to "my model is optimizing a representation-theoretic energy functional."

The code is reportedly available at github.com/Y-Agent/nn-group-representation-learning. If you want to validate these spectral patterns on your own custom group structures, start there. Don't expect to use this to debug a 70B parameter LLM tomorrow. Instead, use it to understand the fundamental mechanics of how symmetry-aware layers actually organize their weights.

Figures from the paper