Grokking Finite-Dimensional Algebra

Pascal Jr Tikeng Notsawo; Guillaume Dumas; Guillaume Rabusseau

cs.LGFeb 23, 2026

Grokking Finite-Dimensional Algebra

Authors: Pascal Jr Tikeng Notsawo, Guillaume Dumas, Guillaume Rabusseau

Organizations: Université de Montréal · Mila, Montréal · CHU Sainte-Justine Research Center, Montréal · CIFAR AI Chair

Abstract

This paper investigates the grokking phenomenon, which refers to the sudden transition from a long memorization to generalization observed during neural networks training, in the context of learning multiplication in finite-dimensional algebras (FDA). While prior work on grokking has focused mainly on group operations, we extend the analysis to more general algebraic structures, including non-associative, non-commutative, and non-unital algebras. We show that learning group operations is a special case of learning FDA, and that learning multiplication in FDA amounts to learning a bilinear product specified by the algebra's structure tensor. For algebras over the reals, we connect the learning problem to matrix factorization with an implicit low-rank bias, and for algebras over finite fields, we show that grokking emerges naturally as models must learn discrete representations of algebraic elements. This leads us to experimentally investigate the following core questions: (i) how do algebraic properties such as commutativity, associativity, and unitality influence both the emergence and timing of grokking, (ii) how structural properties of the structure tensor of the FDA, such as sparsity and rank, influence generalization, and (iii) to what extent generalization correlates with the model learning latent embeddings aligned with the algebra's representation. Our work provides a unified framework for grokking across algebraic structures and new insights into how mathematical structure governs neural network generalization dynamics.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Grokking (non-trivial $\Delta t_{\text{gen}}=t_{4}-t_{2}$ ) on a 2-layer MLP with $\mathrm{ReLU}$ activation trained to learn multiplication in the algebra of complex numbers over the field $\mathbb{F}_{7}=\mathbb{Z}/7\mathbb{Z}$ , where $(a+b\mathbf{i})(c+d\mathbf{i})=ac+bd\mathbf{i}^{2}+(ad+bc)\mathbf{i}$ and $\mathbf{i}^{2}=-1\equiv 6\text{ mod }7$ . (b) Histogram of grokking delay $\Delta t_{\text{gen}}$ under single modifications of the multiplication rules of complex numbers in $\mathbb{F}_{7}$ , where each perturbation consists in changing the outcome of one multiplication relation (e.g., replacing $\mathbf{i}^{2}=-1$ with $\mathbf{i}^{2}=1$ ).

Panel of Figure 2 — Figure 2: Representation quality and generalization performances as a function of training steps. Before grokking $\mathcal{A}_{\text{rep}}$ remains relatively low and then sharply increases as the model groks.

Panel of Figure 3 — Figure 3: Generalization accuracy as a function of representation quality for different training data size ( $r$ ) and model layers ( $0$ for first layer, $1$ for the second, etc.) for a MLP and a Transformer Encoder: $\mathcal{A}_{\text{test}}$ increases with $\mathcal{A}_{\text{rep}}$ in a low-data regime.

Figure 4: Evolution of the test accuracy

\mathcal{A}_{\text{test}}

during training for different training data fraction

r\in\{0.2,0.3\}

Figure 5: Test loss

\mathcal{L}_{\text{test}}

and grokking step

t_{4}

as a function of training data fraction

r

Figure 7: Generalization delay

\Delta t_{\text{gen}}

as a function of training data fraction

r\in(0,1)

for

\mathcal{M}_{2}(\mathbb{F}_{3})

and

\text{UT}_{2}(\mathbb{F}_{5})

, the algebra of

2\times 2

matrices (

n=4

) over

\mathbb{F}_{3}

and th

2\times 2

upper triangular matrices (

n=3

) over

\mathbb{F}_{5}

, respectively.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Grokking on 2 layers Tansformer trained on the algebra of complex numbers in

\mathbb{Z}/7\mathbb{Z}

using a training data fraction of

r=0.5

Figure 9: Grokking on 3 layers LSTM trained on the algebra of complex numbers in

\mathbb{Z}/7\mathbb{Z}

using a training data fraction of

r=0.3

Figure 10: Generalization accuracy as a function of representation quality for different training data size ( $r$ ) and model layers ( $0$ for first layer, $1$ for the second, etc.) for a LSTM: $\mathcal{A}_{\text{test}}$ increases with $\mathcal{A}_{\text{rep}}$ in a low-data regime.

Figure 11: Training and test loss and accuracy as a function of the training data fraction.

Figure 12: Phase transition times as a function of the fraction of data used to train the model:

t_{1}

is the step a which the training accuracy

\mathcal{A}_{\text{train}}

reaches a value strictly greater than

5\%

for the first time,

t_{2}

the step at which

\mathcal{A}_{\text{train}}

reaches

99\%

for the first time,

t_{3}

the step at which the test accuracy

\mathcal{A}_{\text{test}}

reaches a value strictly greater than

5\%

for the first time, and

t_{2}

the step at which

\mathcal{A}_{\text{test}}

reaches

\approx 99\%

for the first time.

Figure 13: Test loss as a function of training data fraction

r

for

\rank(\mathbf{\boldsymbol{\mathcal{C}}}^{*}_{(1)})

(left) and

\rank(\mathbf{\boldsymbol{\mathcal{C}}}^{*}_{(2)})

(right). The sizes of the dots are proportional to

t_{4}

Figure 14: Training and test loss and accuracy as a function of training steps for different values of the training data fraction

r

Grokking

Grokking is delayed generalization in which neural networks memorize training data before abruptly learning a rule that works on unseen examples.

See all 77 papers

Algebraic Representability as the Limiting Regime of Grokking: An Exactly Solvable Model with Holomorphic Activations
Jul 15, 2026Chon-Fai Kam, Xavier Cadet, Miloud Bessafi +1Feedforward Neural Networks Neural Network Expressivity
University Paris City & University of Reunion, Paris, France · Dipartimento di Fisica e Chimica Emilio Segrè, Università degli Studi di Palermo, Palermo, Italy · Thayer School of Engineering, Dartmouth College, Hanover, NH 03755, USA +2
From Symmetry to Invariance: Learning Galois Equivalent Representations in Finite Fields
Aug 23, 2026Zheng Zhang, Na ZhangInvariant Representation Learning Equivariant Representation Learning
Department of Mathematics, Towson University, 7800 York Rd, Towson, MD 21204, USA
Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic
Mar 24, 2026Anand SwaroopGrokking Neural Network Interpretability

Jul 15, 2026J/K move · Enter open · S save

Algebraic Representability as the Limiting Regime of Grokking: An Exactly Solvable Model with Holomorphic Activations

Chon-Fai Kam, Xavier Cadet, Miloud Bessafi, Frederic Cadet

University Paris City & University of Reunion, Paris, France · Dipartimento di Fisica e Chimica Emilio Segrè, Università degli Studi di Palermo, Palermo, Italy · Thayer School of Engineering, Dartmouth College, Hanover, NH 03755, USA+2

Feedforward Neural Networks Neural Network Expressivity Modular Arithmetic Grokking

Open paper AskPDF arXiv

Neural networks trained on modular arithmetic exhibit grokking, a delayed transition from memorisation to generalisation known to depend on model capacity: too little and the network memorises slowly or not at all, too much and it generalises almost immediately. What happens at the extreme of this spectrum, when the architecture's expressible function class collapses to a finite-dimensional algebraic variety? We study two-layer networks with a holomorphic monomial activation sigma(z)=z^k, trained on modular tasks encoded via roots of unity. Here the network output, regardless of hidden width, is confined to a (k+1)-dimensional subspace of characters of (Z_p)^2, an O(k/p^2) slice of the full function space. We give a complete algebraic characterisation of this subspace: a task is representable if and only if its discrete Fourier support lies on the diagonal u+v = k (mod p), which for linear-phase targets reduces to the arithmetic criterion m+n=k. This is not merely a constraint on eventual generalisation but on memorisation itself: because the outputs are algebraically confined, a non-representable target cannot be fit even on the training set, and we prove a positive lower bound on the training loss, independent of width. Across 585 runs the algebraic prediction matches the observed outcome with 99.8% accuracy, with no memorisation regime and no grokking; outcomes split cleanly into instant success and outright failure. This binary behaviour is the limiting case of the capacity-grokking relationship: when the expressible class shrinks to a fixed algebraic object, the question of when a network will grok dissolves into whether it can represent the target at all. A bottleneck ablation connects this extreme to standard networks, tracing a continuous path from representational failure, through memorisation without generalisation, to grokking with a shrinking gap as capacity grows.

Abstract

Grokking

Explore similar work

Algebraic Representability as the Limiting Regime of Grokking: An Exactly Solvable Model with Holomorphic Activations

From Symmetry to Invariance: Learning Galois Equivalent Representations in Finite Fields

Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic