Acceleration of Grokking in Learning Arithmetic Operations via Kolmogorov-Arnold Representation

TL;DR

The paper addresses the grokking phenomenon observed when transformers learn arithmetic operations, proposing a principled KA-representation framework to accelerate generalization and transfer knowledge across related tasks. By linking KA embeddings with transformer components, it shows that the SUM (Sigma) and embedding (Phi) parts can be shared across operations, enabling decoder-block and embedding transfers to speed grokking. The authors demonstrate improved learning speed across commutative, abelian, and anti-abelian settings, and extend to tasks such as composition of operations and systems of equations, including learning under limited token sets. The work provides both theoretical (KA-based decomposition and representation theory) and empirical contributions that yield more data-efficient, transferable learning for algorithmic reasoning in neural networks.

Abstract

We propose novel methodologies aimed at accelerating the grokking phenomenon, which refers to the rapid increment of test accuracy after a long period of overfitting as reported in~\cite{power2022grokking}. Focusing on the grokking phenomenon that arises in learning arithmetic binary operations via the transformer model, we begin with a discussion on data augmentation in the case of commutative binary operations. To further accelerate, we elucidate arithmetic operations through the lens of the Kolmogorov-Arnold (KA) representation theorem, revealing its correspondence to the transformer architecture: embedding, decoder block, and classifier. Observing the shared structure between KA representations associated with binary operations, we suggest various transfer learning mechanisms that expedite grokking. This interpretation is substantiated through a series of rigorous experiments. In addition, our approach is successful in learning two nonstandard arithmetic tasks: composition of operations and a system of equations. Furthermore, we reveal that the model is capable of learning arithmetic operations using a limited number of tokens under embedding transfer, which is supported by a set of experiments as well.

Figures (6)

• Figure 1: Illustration of shared structures within KA representation for binary operations.
• Figure 2: Schematic diagram for weight transfer. Weights for learning addition (middle) are transferred to learn the composited operations (left) and other commutative binary operations (right).
• Figure 3: Schematic diagram for learning a system of equations via weight transfer.
• Figure 4: PCA visualization for addition (top) and multiplication (bottom). We see that $\theta_{i,i+1}$'s, the angles between token $i$ and $i+1$, are maintained similarly. A similar feature is observed in the multiplication hinting that tokens form a multiplicative group.
• Figure 5: The plots show the effectiveness of decoder block transfer in accelerating grokking compared to the baseline.

Theorems & Definitions (19)

• Definition 1: cyclic group
• Theorem 1: fundamental theorem of finite abelian group
• Theorem 2: Kolmogorov-Arnold representation braun2009constructive
• Example 1
• Theorem 3: permutation-invariant representation zaheer2017deep
• Corollary 1
• proof
• Theorem 4: abelian group representation
• proof
• Corollary 2: cyclic group representation