Technical Summary: Grokking Finite-Dimensional Algebra

Problem Statement

The paper investigates the "grokking" phenomenon—the sudden transition from memorization to generalization observed during the training of overparameterized neural networks—within the context of learning multiplication in Finite-Dimensional Algebras (FDAs). While prior research has largely focused on grokking in the context of finite group operations (e.g., modular arithmetic), this work extends the analysis to more general algebraic structures. These include non-associative, non-commutative, and non-unital algebras, which better model complex systems found in physics, chemistry, and natural language processing where order and structural constraints (like Lie brackets or chemical reaction pathways) are critical. The core challenge is to understand how specific algebraic properties (associativity, commutativity, unitality) and structural features of the algebra's representation (sparsity, rank) influence the emergence, timing, and dynamics of grokking.

Methodology

The authors propose a unified tensorial framework to study FDAs, treating the learning of algebraic multiplication as a problem of recovering a structure tensor.

1. Tensorial Formulation: Any $n$-dimensional algebra over a field $\mathbb{F}$ is identified with a structure tensor $C \in \mathbb{F}^{n \times n \times n}$. The multiplication of two elements $u, v$ is expressed as a bilinear product $u \cdot v = C \times_1 u \times_2 v$. This formulation allows algebraic properties (associativity, commutativity, unitality) to be characterized as specific tensor identities.
2. Unification with Groups: The paper demonstrates that learning group operations is a special case of learning FDAs. The group algebra $\mathbb{F}[G]$ is shown to be an $n$-dimensional FDA where the structure tensor entries are determined by the group's Cayley table.
3. Learning Regimes:
   • Real Fields ($\mathbb{R}$): Learning is framed as a linear inverse problem or matrix factorization task. The paper connects FDA learning to low-rank matrix recovery, noting that over $\mathbb{R}$, structured algebras (e.g., unital) form a measure-zero subset, making principled empirical study difficult without specific constraints.
   • Finite Fields ($\mathbb{F}_p$): This setting is the primary focus. Here, algebraic properties have positive density. The task is modeled as learning an embedding matrix $E$ for algebra elements (treated as a vocabulary) followed by a classifier. Grokking is hypothesized to occur when the model learns an embedding geometry that aligns with the algebraic structure, allowing a linear readout to generalize.
4. Experimental Setup: The authors train models (MLPs, LSTMs, Transformers) to predict the product of two elements in an algebra over $\mathbb{F}_7$. They systematically vary:
   • Algebraic Properties: Associativity ($a$), Commutativity ($c$), and Unitality ($u$), creating categories like $ac\bar{u}$ (associative, commutative, non-unital).
   • Structural Properties: The rank of the mode-3 unfolding of the structure tensor ($r^{(3)}$) and its sparsity ($s$).
   • Data Availability: The ratio of training data to the total possible pairs ($r = |D_{train}|/|D|$).
5. Representation Probing: To verify if models learn algebraic representations, the authors perform linear probing. They check if the model's internal features $f(u)$ and $f(v)$ can be combined via a learned bilinear map to reconstruct the feature of the product $f(u \cdot v)$, measuring the alignment between learned features and the algebra's structure.

Key Contributions

• Unified Framework: The paper establishes an equivalence between the space of structure tensors $\mathbb{F}^{n \times n \times n}$ and all $n$-dimensional FDAs, allowing the study of algebraic properties through tensor constraints rather than abstract algebraic definitions.
• Generalization of Grokking: It extends the grokking phenomenon beyond groups to non-associative, non-commutative, and non-unital algebras, showing that group learning is a special case of this broader framework.
• Theoretical Connections:
  • Over $\mathbb{R}$, FDA learning is connected to matrix factorization with implicit low-rank bias.
  • Over finite fields, grokking is linked to the necessity of learning discrete, algebra-consistent embeddings.
• Systematic Empirical Study: The authors provide the first systematic investigation of how algebraic constraints and tensor structural features affect grokking delay and sample complexity.

Results

• Algebraic Properties and Grokking Delay:
  • Non-unital Associative Algebras: These are the easiest to learn, exhibiting the shortest grokking delays. The absence of a unit constraint allows the model to exploit "shortcuts" in the representation space (e.g., the zero representation or left regular representation), simplifying the optimization landscape.
  • Unital Algebras: The requirement to represent an identity element ($\sum \lambda_i R_i = I_m$) significantly restricts the solution space, often delaying convergence and increasing grokking delay.
  • Commutativity: In non-unital regimes, commutativity further reduces grokking delay by reducing the number of distinct equations the model must satisfy (symmetry in inputs).
• Structural Properties:
  • Rank and Sparsity: Higher rank and higher sparsity of the structure tensor correlate with increased test loss and longer grokking delays. The rank reflects the intrinsic multilinear complexity of the operation, while sparsity is treated as a task-dependent difficulty parameter in the canonical basis.
• Representation Alignment: Generalization coincides with the emergence of algebra-consistent representations. Before grokking, the alignment between model features and the algebraic structure is low; it sharply increases as the model transitions to generalization. In low-data regimes, test accuracy is directly correlated with the quality of this representation alignment.
• Model Architecture: Results were consistent across MLPs, LSTMs, and Transformers, though specific layer behaviors (e.g., embedding layers in Transformers) showed nuances in how structure emerges.

Significance and Claims

The paper claims to provide a unified framework for understanding grokking across diverse algebraic structures, bridging classical algebra with modern learning theory. It posits that:

1. Mathematical Structure Governs Dynamics: The specific algebraic constraints (associativity, unitality) and structural properties (rank, sparsity) of the task fundamentally dictate the difficulty of generalization and the timing of the grokking transition.
2. Grokking as a Phase Transition: The phenomenon is characterized as a representational phase transition where the model shifts from brittle memorization to learning an internal representation that respects the algebraic multiplication rules.
3. Limitations of Group-Centric Views: By moving beyond groups, the work reveals that properties often assumed to simplify learning (like unitality) can actually hinder it by constraining the solution space, challenging previous intuitions derived solely from group operation studies.

The authors remain modest regarding future implications, noting that while their framework recovers group tasks as a special case, the computational complexity of exploring the space of structure tensors grows rapidly with dimension, limiting the scope of current empirical evaluations to moderate dimensions ($n=2$ over $\mathbb{F}_7$). They suggest that tighter theoretical links between tensor properties (like coherence) and sample complexity over finite fields remain an open direction for future work.