Sequential Group Composition: A Window into the Mechanics of Deep Learning

AI summary

Overview

Research area: machine learning theory — specifically mechanistic interpretability, learning dynamics, and computational expressivity of neural networks trained on structured sequence tasks. Technical level: Advanced (requires familiarity with group theory, representation theory, and Fourier analysis over finite groups). Scope: the paper introduces and analyzes a synthetic "sequential group composition" task to characterize, from first principles, what features two-layer networks learn and in what order, and how depth and recurrence change the efficiency of those solutions.

What This Paper Is About

The paper asks how neural networks trained on sequences acquire the ability to perform structured operations such as arithmetic, geometric, and algorithmic computation. To make this question tractable, the authors define the sequential group composition task, in which a network receives a sequence of elements from a finite group (Abelian or non-Abelian) encoded as real vectors and must predict their ordered cumulative product. The goal is to isolate the separate roles of group structure, encoding statistics, and sequence length in shaping learning and expressivity.

Key Contributions

1. A new tractable task. The authors introduce sequential group composition, a regression problem mapping a sequence $(g_1, \ldots, g_k) \in G^k$ to $\prod_{i=1}^{k} g_i \in G$. They prove (Lemma 3.5) that for a nontrivial, mean-centered encoding there is no linear map from $\mathbb{R}^{k|G|}$ to $\mathbb{R}^{|G|}$ that solves it, so the task necessarily requires nonlinear interactions between inputs.

2. An exact theory of feature learning for two-layer networks. For two-layer networks with polynomial activation trained from vanishing initialization, the authors prove that learning proceeds through the Alternating Gradient Flows (AGF) framework of Kunin et al. (2025), acquiring one irreducible representation (irrep) of the group at a time. The order is determined by the group Fourier statistics of the encoding vector $x$ via the argmax in Theorem 4.1.

3. A characterization of when two-layer networks suffice. Theorem 4.3 gives a loss lower bound for newly activated neurons and shows it is achievable when $N \geq (k+1)2^{k}n_{\rho_*}^{k+1}$ neurons activate together. This yields a sufficient width condition $H \geq (k+1)2^{k}\sum_{\rho \in \mathcal{J}(G)} n_{\rho}^{k+1}$, which for the cyclic group with monomial activation reduces to $H \geq (k+1)2^{k-1}\lfloor |G|/2 \rfloor$.

4. Constructive efficiency gains from depth. The authors construct recurrent networks that compose elements sequentially in $k$ steps and multilayer networks that compose adjacent pairs in parallel in $\log k$ layers, both exploiting associativity of the group operation. These results are constructive rather than dynamical — the paper shows efficient solutions exist but does not prove gradient descent finds them.

Main Findings

• Learning proceeds one irrep at a time. Under vanishing initialization, two-layer networks decompose the task into irreps of $G$ and acquire them greedily. During utility maximization, a single neuron specializes to the real part of a single irrep plus its conjugate; during cost minimization, groups of aligned neurons jointly minimize the loss.

• The learning order is set by Fourier statistics of the encoding. Theorem 4.1 shows the next irrep learned is the one maximizing $|\widehat{x}[\rho]|{\text{op}}^{k+1}/(C\rho n_\rho)^{(k-1)/2}$, where $C_\rho = 1$ for real irreps and $C_\rho = 2$ otherwise. Utility maximizers are real parts of complex linear combinations of the matrix entries of $\rho_*$, formalizing the intuition that neurons "align" to an irrep.

• A systematic bias toward low-dimensional irreps. The utility formula favors smaller $n_\rho$ as sequence length $k$ grows. For a one-hot encoding, $|\widehat{x}[\rho]|_{\text{op}} = 1$ for all $\rho$, yet the ordering preference remains. The authors state this establishes a form of "strong universality" hypothesized in Chughtai et al. (2023) and explains why the ordering was hard to see empirically: the effect is subtle at $k=2$ but pronounced as sequence length increases.

• Two-layer networks need exponential width. Figure 4 reports results for two-layer networks with quadratic activations on the cyclic group across 2205 training runs, spanning 21 group sizes, 21 hidden widths, and 5 sequence lengths. Heatmap colors indicate normalized training loss (final loss divided by initial loss), with training run until 99.9% loss reduction or $10^{6}$ optimization steps. The empirical boundary matches the theoretical lower bound $H \geq (k+1)2^{k-1}\lfloor |G|/2 \rfloor$. The paper notes (Remark B.7) that the exponential factor is optimal in general.

• An intermediate capacity regime exists. When the network lacks enough neurons for exact learning, it can still achieve strong performance by finding partial solutions. The authors note these regimes are often associated with unstable dynamics, potentially related to Martinelli et al. (2025) on pairs of neurons approximating gated linear units at the "edge of stability."

• The exponential cost comes from canceling extraneous terms. The per-neuron function splits into a sigma-pi-sigma interaction term $f^{(\times)}$ involving all $k$ inputs, and an additive term $f^{(+)}$ that is extraneous to the task. Multiple neurons must coordinate to cancel $f^{(+)}$, and the Waring-type decomposition of the monomial $z_1 \cdots z_k$ requires $2^k$ neurons to implement a single sigma-pi-sigma unit.

• Depth changes the scaling qualitatively. RNNs with quadratic nonlinearity can compose elements sequentially in $k$ steps, and multilayer networks can compose adjacent pairs in parallel in $\log k$ layers, both leveraging associativity. The paper reports partial empirical evidence (Figure 5) that trained recurrent networks recover structure consistent with the construction.

• Representations and encoding are tied together. Encoding is defined by the orbit of a fixed vector $x \in \mathbb{R}^{|G|}$ under the regular representation, $g \mapsto \lambda(g)^{\intercal} x$; setting $x = e_1$ recovers standard one-hot encoding. The analysis assumes $x$ is mean-centered, that each $\widehat{x}[\rho]$ is either invertible or zero, and that the quantities in Equation 13 are distinct for nonzero coefficients.

Methodology in Plain English

The authors take an idealized but general problem — composing a sequence of transformations drawn from a finite group — and analyze it with the tools of abstract harmonic analysis. Because any finite group's structure can be encoded in unitary matrices, and any vector encoding can be decomposed with a group Fourier transform into irreducible representations, the composition task can be broken into independent frequency-like pieces.

To study learning rather than just expressivity, the authors place a two-layer network with a monic, origin-passing polynomial activation of degree $k$ in the vanishing initialization regime, where parameters start near the origin and evolve under a time-rescaled gradient flow. This puts them inside the AGF framework, which describes training as an alternation between a phase where dormant neurons compete to align with the most informative direction in the residual (utility maximization) and a phase where newly activated neurons jointly minimize the loss (cost minimization). By solving these two optimization problems in closed form for the group composition task, the authors derive exactly which irrep is learned at each step, the resulting loss bound, and the number of neurons needed to saturate it.

For depth, the authors switch to an explicit construction approach: they write down recurrent and multilayer architectures and show by construction that associativity lets them compose intermediate results instead of requiring one giant hidden layer to coordinate all $k$ inputs at once. They complement the theory with experiments on cyclic and dihedral groups, including a large sweep over group sizes, widths, and sequence lengths to map the empirical width threshold.

Why This Matters

This paper shifts the study of learned structure from post-hoc inspection of trained models to a derivation from first principles. Rather than inferring circuits from empirical patterns, it predicts which Fourier features appear, in what order, and why — and it ties the number of neurons required to the algebraic structure of the task.

The framework unifies a body of prior empirical work, including the widely studied modular addition task (the binary case $k=2$ over the cyclic group $C_p$), and gives a mechanistic explanation for why networks acquire low-dimensional representations first.

Real-world applications motivated by the paper's framing:

• Arithmetic and algorithm learning in language models. The paper cites work showing related Fourier features in large pre-trained models performing arithmetic, and the task generalizes the classical discrete Fourier transform, which is recovered when $G = C_p$.
• Robotics and pose tracking. Tracking a body's trajectory through physical space requires composing rigid motions or integrating successive displacements — a composition of group elements.
• Combinatorial puzzles and planning. Solving puzzles such as the Rubik's Cube amounts to composing a sequence of moves, each a group element.
• Symmetry-aware modeling of scientific data. Reindexing atoms of a molecule leaves its physical properties unchanged, and translating or reflecting an image preserves the scene.

Industry relevance: the results bear on architecture selection for sequence tasks where compositional structure matters. The paper's central practical lesson — that transformers-style parallel composition achieves $\log k$ depth while recurrent models do it in $k$ sequential steps and shallow feedforward networks need exponential width — is relevant to anyone choosing between feedforward, recurrent, and attention-based designs for algorithmic or structured reasoning workloads. It also suggests that feature-learning order in such models is predictable from the statistics of the input encoding, which has implications for data representation choices.

Future Directions

• From construction to dynamics with depth. The paper's efficiency results for RNNs and multilayer networks are constructive only; whether and how gradient descent converges to these efficient solutions is left open. The authors note partial empirical evidence in Figure 5 but no dynamical theory.

• Closing the gap in the intermediate capacity regime. Networks below the exact-learning width still achieve strong performance through partial solutions, and the authors connect these regimes to unstable "edge of stability" dynamics. A theory of these partial solutions is not established.

• Generalizing beyond polynomial activations. The main analysis uses a monic, origin-passing polynomial of degree $k$. Section C.2 discusses how smooth non-polynomial activations can access analogous interaction terms through Taylor expansions near the origin, but the full analysis is not extended there.

• Broader implications and limitations. Section 6 is described as discussing broader implications, limitations, and next steps; Appendix E collects technical remarks on how the work relates to previous group-theoretical constructions.

Target Audience

This paper is for machine learning theorists, mechanistic interpretability researchers, and mathematically inclined practitioners interested in why neural networks learn structured computations. Readers benefit from background in group theory, representation theory, and Fourier analysis. It will be especially relevant to those studying feature learning dynamics, grokking-adjacent phenomena, algorithmic reasoning tasks such as modular arithmetic and group composition, and the expressivity differences between feedforward, recurrent, and attention-based architectures.

Authors’ abstract

How do neural networks trained over sequences acquire the ability to perform structured operations, such as arithmetic, geometric, and algorithmic computation? To gain insight into this question, we introduce the sequential group composition task. In this task, networks receive a sequence of elements from a finite group encoded in a real vector space and must predict their cumulative product. This task can be order-sensitive and cannot be solved by a linear model. Our analysis isolates the roles of the group structure, encoding statistics, and sequence length in shaping learning. We prove that two-layer networks from vanishing initialization learn this task one irreducible representation of the group at a time in an order determined by the Fourier statistics of the encoding. To perfectly learn the task, these networks require a hidden width exponential in the sequence length $k$. In contrast, we construct deeper architectures that exploit associativity to dramatically improve this scaling: recurrent neural networks can compose elements sequentially in $k$ steps, while multilayer networks can compose adjacent pairs in parallel in $\log k$ layers. Overall, the sequential group composition task offers a tractable window into the mechanics of deep learning.