Abstract

Generalization in deep learning is difficult to explain because neural networks do not merely choose a predictor from a fixed hypothesis class. They also learn the representation in which prediction becomes possible. This paper studies that problem in a controlled setting: modular arithmetic tasks where transformers exhibit grokking, a delayed transition from memorization to generalization.

The paper argues that grokking can be understood as convergence toward a representation determined by the task itself. For modular arithmetic, this task-derived structure can be written down exactly using quotient groups and irreducible representations. Generalization begins when the learned representation aligns with a sufficient canonical subspace, and the remaining deviation away from that subspace predicts excess loss.

Main Idea

The key object is the canonical representation of a task. A deterministic task partitions its input space into equivalence classes: two inputs are equivalent when they have the same correct label. This gives a canonical quotient of the input space before any model is trained. For modular arithmetic, this quotient is especially explicit because the target functions come from cyclic group operations. Addition and subtraction reduce to the additive group (\mathbb Z_p); multiplication and division reduce to the multiplicative group (\mathbb Z_p^\times), or equivalently to addition in exponent space after choosing a generator.

Canonical decomposition and representation of modular addition for p = 7.
Canonical decomposition. Modular equivalence first collapses inputs into quotient classes. Those classes can then be represented geometrically, here as rotations of \(\mathbb Z_7\) in the plane.

Once the quotient is known, representation theory tells us what the corresponding vector-space structure should be. For modular addition, the quotient group (\mathbb Z_p) can be represented by rotation blocks, the real irreducible representations of the cyclic group. A single non-trivial block is already faithful to the quotient structure, while direct sums of blocks give richer canonical representations. This is the paper’s bridge from the algebra of the task to the geometry a neural network can learn.

From Geometry to Loss

The next step is to connect this geometry to loss. Cross-entropy does not only require the correct class to win; it requires margins between the correct logit and all incorrect logits. The paper therefore defines a sufficient representation: a canonical representation whose margins are large enough to reach a chosen loss threshold. In the modular arithmetic case, each irreducible block contributes a predictable amount to every class margin, and these contributions add across blocks.

Additive margin contributions from irreducible blocks for modular addition.
Margins from irreducible blocks. Each rotation block contributes a geometric margin. Direct sums of blocks add these contributions, which turns the representation-theoretic decomposition into a prediction about cross-entropy loss.

This gives a concrete prediction: for a target loss, there is a family of sufficient irrep subspaces, and optimization under weight decay should prefer the cheapest one. Increasing the amplitude of a few modes improves margins but costs parameter norm; activating more modes spreads the required amplitude but adds architectural cost. The paper models this tradeoff and predicts the number of irreducible modes the model should converge toward.

What the Experiments Show

The experiments test these predictions on modular addition, subtraction, multiplication, and division. During training, the learned logits are projected into the predicted sufficient subspace. Generalization begins when the relevant irreducible modes become active and the representation concentrates on them. The spectrum below illustrates this collapse: rather than using arbitrary directions, the model reallocates mass toward a small set of task-aligned modes.

Irrep amplitude spectrum during grokking.
Spectrum collapse during training. As grokking occurs, representation mass concentrates on a small subset of irreducible modes. This is the empirical signature of convergence toward the predicted sufficient representation.

The paper then defines representational deviation as the component of the learned representation outside the sufficient subspace. This deviation predicts excess loss: once the sufficient component is fixed, the remaining loss contribution comes from the off-subspace part. Empirically, this deviation term closely tracks the generalization gap during training.

Why It Matters

The final experiment turns the theory into an intervention. Instead of using only standard weight decay, the model is trained with a penalty on deviation from the canonical sufficient subspace. This regularizer accelerates grokking and can make models generalize in data regimes where standard weight decay fails. The main conclusion is that, at least for modular arithmetic, generalization is governed by alignment with a task-derived canonical representation rather than by memorization or parameter norm alone.

More broadly, the paper gives a precise version of a common intuition in representation learning: models generalize when they discover the right latent structure. In this setting, the “right” structure is not inferred post hoc from the model. It is derived from the task before training, then used to predict and intervene on the learning dynamics.

Citation

Moullec, Matthieu, and Andreas Vlachos. 2026. “The Canonical Representation of a Task.” In review at TMLR.

@misc{moullec2026canonical,
  author = {Matthieu Moullec and Andreas Vlachos},
  year = {2026},
  title = {The Canonical Representation of a Task},
  note = {In Review, TMLR}
}

grokking representation theory generalization modular arithmetic interpretability