Emergence in non-neural models: grokking modular arithmetic via average gradient outer product

Emergence in non-neural models: grokking modular arithmetic via average gradient outer product

Speakers: Neil Mallinar, Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, Misha Belkin

Conference: ICML 2025

YouTube: https://slideslive.com/39043902

Overview

In this compelling talk, Neil Mallinar and his co-authors present groundbreaking research challenging the conventional understanding of generalization in machine learning, particularly the phenomenon known as grokking. Traditionally observed in neural networks, grokking describes a delayed generalization where a model first perfectly fits the training data with poor test performance, only to suddenly achieve high test accuracy after thousands of additional training steps. Mallinar's presentation, "Emergence in non-neural models: grokking modular arithmetic via average gradient outer product," investigates whether this emergent behavior is unique to neural networks and proposes a unified framework to understand it across diverse model architectures.

The core of their work revolves around the task of modular arithmetic, a synthetic problem known to induce grokking in neural networks. By introducing the Average Gradient Outer Product (AGOP) as a tool to analyze feature learning, the researchers demonstrate that non-neural models, specifically Kernel Recursive Feature Machines (RFM), can also exhibit grokking. This finding fundamentally expands our understanding of emergent phenomena in ML. Furthermore, they reveal a common underlying mechanism—the learning of block circulant feature transformations—that drives generalization in both neural and non-neural models, offering a model-agnostic perspective on this complex behavior.

The significance of this research extends beyond theoretical insights. By identifying the specific feature structures learned during grokking, the authors propose methods to proactively engineer inputs that eliminate the generalization delay, leading to immediate learning. Crucially, they introduce novel a posteriori metrics, circulant deviation and AGOP alignment, which can track the gradual emergence of these features even when traditional performance metrics like accuracy and loss remain stagnant. This work provides a powerful analytical lens for understanding deep learning's inner workings and offers practical implications for accelerating model development and improving diagnostic tools for emergent learning behaviors.

Background

▶ Watch: Introduction to grokking and delayed generalization phenomenon (0:00)

The landscape of machine learning has long been characterized by a delicate balance between fitting training data and generalizing to unseen examples. A common pattern observed during model training is that initial steps lead to similar train and test performance, indicating low generalization error. As training progresses, models might overfit, where train accuracy reaches 100% while test accuracy degrades, making it easy to identify an optimal stopping point. However, the phenomenon of grokking presents a stark departure from this traditional view. First extensively documented by Power et al., grokking describes a scenario where a model rapidly achieves perfect training accuracy, but test accuracy remains near zero for an extended period—sometimes thousands of training steps. Remarkably, after this prolonged stagnation, test accuracy sharply increases to near-perfect levels. This delayed generalization was initially considered a unique "emergent phenomenon" primarily associated with neural networks.

The task chosen for studying grokking in this research is modular arithmetic, specifically addition (A + B mod P) and multiplication (A * B mod P). This task is particularly interesting because it is simple enough to analyze, yet complex enough to induce grokking in neural networks. The problem can be conceptualized as imputing missing values in a Cayley table. For example, in modular addition mod 3, one might learn that 1 + 2 mod 3 = 0. Data is represented by one-hot encoding individual digits (A and B) and concatenating them into a 2P-dimensional input vector, with the output being a P-dimensional one-hot vector representing the label. Non-feature learning methods, such as standard kernels, typically fail to generalize on modular arithmetic, making it an ideal testbed for investigating feature learning and emergent generalization.

A significant challenge posed by grokking is the inadequacy of traditional performance metrics like train and test accuracy or loss. During the thousands of steps where test accuracy is near zero, these metrics offer no indication of progress, making it impossible to determine whether continued training is fruitful or merely deepening overfitting. This highlights a critical need for alternative measures that can track the underlying learning process, especially the emergence of meaningful features, even when superficial performance metrics are stagnant. To address this, the authors leverage the Average Gradient Outer Product (AGOP). AGOP is defined as the average of the outer product of the gradient of the predictor function (F) with respect to its input data (Xi). For a D-dimensional input, AGOP yields a D x D matrix. Prior work has established that AGOP captures the features learned by neural networks, effectively modeling the covariance of the weights within the network during training. Intuitively, AGOP reveals how perturbations to specific input features most significantly affect the predictor's output. It can also be thought of as a form of supervised Principal Component Analysis (PCA) through the learned function. Critically, AGOP's power lies in its ability to completely decouple feature learning from the specific functional form of the predictor F, requiring only F and its differentiability with respect to the input. This decoupling allows AGOP to enable feature learning in more general models, including non-neural architectures.

Key Findings

▶ Watch: Introducing Average Gradient Outer Product (AGOP) for features (3:00)

The research presented in this talk yields several profound insights into the nature of grokking and feature learning:

1. Grokking is Not Unique to Neural Networks: A primary finding is the demonstration that non-neural models can indeed grok modular arithmetic. Specifically, the Kernel Recursive Feature Machine (RFM), when applied to modular addition and multiplication, exhibits the characteristic delayed generalization previously thought to be exclusive to neural networks. This expands the scope of emergent phenomena beyond deep learning architectures.

1. Grokking as Gradual Feature Learning: The authors argue that grokking is a manifestation of gradual feature learning. By tracking the AGOP matrix during the training of kernel RFM, they observe a continuous, monotonic emergence of structured features, even during the initial phases where accuracy and loss metrics remain flat and uninformative. This challenges the notion that grokking is a sudden, unpredictable event, re-framing it as the observable outcome of a slower, underlying process of feature discovery.

1. Unified Understanding Through Block Circulant Feature Transformations: AGOP serves as a key analytical tool, revealing a common feature structure learned by both neural and non-neural models for modular arithmetic tasks. For modular addition, both kernel RFM and neural networks implicitly learn block circulant feature transformations. For modular multiplication, a similar block circulant structure is discovered after reordering coordinates using the discrete logarithm. This indicates a universal computational strategy for solving these tasks, transcending specific model architectures.

1. Eliminating Grokking with A Priori Feature Engineering: A significant practical implication is the discovery that by applying these learned feature transformations a priori to the input data, the grokking delay can be completely eliminated. When inputs are transformed using random circulant features, both standard kernels and neural networks immediately generalize to modular arithmetic tasks without any delay. This not only confirms the importance of these specific features but also provides a general prescription for accelerating learning in scenarios where the underlying feature structure is understood.

1. Novel A Posteriori Measures for Progress Tracking: Recognizing the limitations of traditional accuracy and loss metrics during grokking, the research proposes two a posteriori measures to track feature learning progress:

• Circulant Deviation: A task-specific metric for modular arithmetic that quantifies how far the AGOP matrix is from a perfect circulant structure. It monotonically drops to zero as features emerge.
• AGOP Alignment: A more general measure that calculates the Pearson correlation between the current AGOP and the final AGOP obtained from a fully generalized model. This metric smoothly and monotonically increases to one, accurately reflecting the progress of feature learning even when accuracy is stagnant.

1. Grokking is Not Tied to Gradient Descent: The observation of grokking in kernel RFM, which does not rely on gradient descent-based optimization in the same manner as neural networks, suggests that the phenomenon is not inherently tied to the specifics of gradient-based optimization methods. Instead, it appears to be a broader characteristic of feature learning itself.

Technical Deep Dive

▶ Watch: Kernel RFM's initial overfitting before generalization (5:00)

The core of this research unpacks the mechanisms behind grokking, moving beyond simply observing the phenomenon to understanding its underlying technical drivers. The exploration begins with the setup of modular arithmetic as the learning task. For A + B mod P, the inputs A and B are individually one-hot encoded into P-dimensional vectors. These are then concatenated to form a 2P-dimensional input vector. The model's goal is to output a P-dimensional vector whose argmax corresponds to the correct one-hot encoded label (A+B mod P). For example, for A+B mod 3, the input vector would be 6-dimensional, and the output 3-dimensional.

The key to extending grokking to non-neural models lies in the Recursive Feature Machine (RFM) algorithm, which the authors apply in a kernelized setting, hence Kernel RFM. RFM operates iteratively:

1. Initialization: Start with some initial features, often the identity matrix (M-naught).
2. Estimation: Fit an estimator (e.g., kernel ridgeless regression) to the training data after it has been transformed by the current estimate of the features. This step ensures the model is learning from the current feature representation.
3. Feature Update: Compute the Average Gradient Outer Product (AGOP) over the current predictor and the training data. This AGOP matrix then becomes the updated estimate of the features for the next iteration.
4. Iteration: Repeat steps 2 and 3 until a convergence criterion is met.

The AGOP itself is a crucial technical component. For a D-dimensional input vector, the AGOP is a D x D matrix, representing the average of the outer product of the gradient of the predictor with respect to the input data. Formally, if F is the predictor and Xi are input data points, AGOP is E_i [ (∇_Xi F(Xi)) (∇_Xi F(Xi))^T ]. This matrix effectively captures how the model's sensitivity to input changes is distributed across different input dimensions. In the context of neural networks, prior work has shown AGOP to model the covariance of the weights. For Kernel RFM, AGOP provides a mechanism to dynamically learn and refine feature transformations.

When Kernel RFM is applied to modular addition (e.g., mod 61), the initial features are the identity matrix, and kernel ridgeless regression is used. As expected, the train accuracy quickly reaches 100%, but test accuracy remains at zero. However, by examining the AGOP matrix at each iteration, a fascinating pattern emerges. While initially identity, AGOP quickly develops non-trivial structure, particularly in its off-blocks. Over successive iterations (e.g., 5 steps, then 10 steps), this structure becomes increasingly sharp and well-defined. This revealed structure is a block circulant matrix. A circulant matrix is characterized by constant values along all its "wrapped" diagonals. A block circulant matrix extends this concept to blocks of values.

The emergence of block circulant matrices is highly significant because they can be diagonalized by the Discrete Fourier Transform (DFT). The DFT is a natural and well-known transformation for modular addition, forming the basis of the Fourier multiplication algorithm for efficient modular arithmetic. The authors even provide a theoretical proof in their paper that a quadratic kernel, when equipped with circulant feature transformations and trained on all P-squared samples, exactly implements this algorithm. This deep connection between the learned AGOP structure and a known efficient algorithm for the task strongly supports the claim that AGOP is indeed capturing meaningful feature learning.

The investigation extends to modular multiplication mod P. Initially, the AGOP for multiplication appears noisy and uninterpretable. However, the authors leverage the mathematical property that the logarithm transforms multiplication into addition. For modular arithmetic, this corresponds to the discrete logarithm (DL). By reordering the coordinates (rows and columns) of each sub-block in the AGOP matrix according to the discrete logarithm mod P, the hidden block circulant structure is revealed. This elegant solution further solidifies AGOP's ability to uncover task-appropriate feature transformations, even when they are not immediately obvious.

Crucially, the authors also demonstrate that neural networks trained on these modular arithmetic tasks, when generalized, also implement these same block circulant feature transformations. This is shown by examining the covariance of the weights within the neural network and the AGOP of the neural network itself, revealing that they too converge to block circulant structures. This provides a unified understanding of grokking across both neural and non-neural architectures, highlighting AGOP as a model-agnostic tool for feature analysis.

Experimental Setup & Results

▶ Watch: Kernel RFM successfully exhibits delayed generalization (grokking) (6:00)

The experimental setup focused on the task of learning modular arithmetic, specifically addition (A + B mod P) and multiplication (A * B mod P). A prime number P = 61 was frequently used for demonstrations. The dataset was generated by taking all possible pairs of A and B (from 0 to P-1), and a random 50% of these examples were used for training, with the remaining 50% reserved for testing.

The primary models investigated were:

1. Kernel Recursive Feature Machine (RFM): This non-neural model was initialized with identity features and used kernel ridgeless regression as its base estimator.
2. Neural Networks: Specifically, a fully connected network and a So Transformer model were mentioned as exhibiting grokking behavior on these tasks in prior literature (Power et al.) and replicated in this work.

The initial results for Kernel RFM on modular addition mod 61 mirrored the classic grokking curves:

• Accuracy Curves: Train accuracy rapidly reached 100%, while test accuracy remained near zero for an extended period (e.g., the first 10 RFM iterations). After this delay, test accuracy sharply increased to 100%, demonstrating successful generalization.
• Loss Curves: Similarly, training loss quickly dropped to zero, while test loss remained high, only to sharply decrease much later.

The key to understanding the underlying process came from visualizing the Average Gradient Outer Product (AGOP) matrices at different stages of RFM training:

• Initial AGOP (Identity Features): The AGOP matrix was the identity, reflecting no learned features.
• Early Iterations (e.g., 5 steps): Even though test accuracy was still zero, the AGOP matrix began to show a non-trivial, structured pattern, particularly in its off-diagonal blocks. This indicated a nascent feature emergence.
• Later Iterations (e.g., 10+ steps): As training progressed and grokking occurred, the AGOP matrix evolved into a clear and sharp block circulant structure. This visual evidence strongly supported the hypothesis of gradual feature learning.

For modular multiplication mod P, the initial AGOP visualizations appeared noisy and uninterpretable, even after generalization. However, by applying a coordinate reordering based on the discrete logarithm (DL) mod P, the hidden block circulant structure was revealed within the noisy AGOP. This confirmed that the same underlying feature transformation principles applied, albeit requiring a specific transformation to unmask them.

A critical set of experiments involved ablations or interventions to test the functional importance of the discovered feature structures:

• Random Circulant Feature Transformations: The researchers took random circulant feature matrices and used them to transform the input data a priori.
• When these transformed inputs were fed into a standard kernel model, the model immediately generalized to modular arithmetic tasks, achieving 100% test accuracy without any grokking delay. This setup even showed better sample complexity compared to the iterative RFM.
• Similarly, when transformed inputs were fed into a neural network, the grokking delay was completely eliminated, and the network learned the task immediately. These results unequivocally demonstrated that the block circulant features are sufficient for generalization and that their a priori application can bypass the grokking phase.

To address the challenge of tracking progress during the grokking delay, two new a posteriori metrics were proposed:

• Circulant Deviation: This metric measured the variance of the wrapped diagonals within the AGOP matrix. As the AGOP became more circulant, this deviation dropped smoothly and monotonically to zero across the initial iterations where accuracy was stagnant, accurately reflecting feature emergence.
• AGOP Alignment: This more general metric computed the Pearson correlation between the current AGOP and the final AGOP obtained from a fully generalized model. It also showed a smooth, monotonic increase towards one during the early iterations, providing a model-agnostic indicator of feature learning progress.

These experimental results collectively provided robust evidence for the central claims: non-neural models grok, grokking is gradual feature learning, and AGOP is a powerful tool for understanding and leveraging these emergent features.

Practical Implications

▶ Watch: Significance of circulant features: Discrete Fourier Transform (6:45)

The findings presented in this talk carry substantial practical implications for practitioners, infrastructure teams, model builders, and deployers across the ML ecosystem.

Firstly, the demonstration that grokking is not exclusive to neural networks but also occurs in non-neural models like Kernel RFM significantly broadens our understanding of generalization. This suggests that emergent, delayed generalization might be a more pervasive phenomenon in complex learning systems than previously thought. Practitioners encountering seemingly inexplicable plateaus in test accuracy, followed by sudden improvements, should consider that their models might be in a grokking phase, regardless of whether they are using deep learning or other sophisticated feature-learning algorithms.

Secondly, the research highlights the critical inadequacy of traditional accuracy and loss metrics during periods of emergent learning. Relying solely on these a priori measures can lead to premature stopping, misdiagnosis of overfitting, or unnecessary discarding of models that are, in fact, undergoing crucial feature learning. This calls for the development and adoption of more sophisticated diagnostic tools. The proposed a posteriori measures, circulant deviation and AGOP alignment, offer a blueprint for such tools. While a posteriori (requiring a generalized model for comparison), they demonstrate the feasibility of tracking hidden learning progress. Infra teams and MLOps platforms could integrate similar feature-centric metrics to provide a richer view of model training dynamics, especially for tasks prone to emergent behaviors.

Thirdly, the discovery that both neural and non-neural models learn specific block circulant feature transformations for modular arithmetic, and that applying these features a priori can eliminate grokking, has profound implications for feature engineering and model initialization. If the underlying mathematical structure of a problem is understood, or if AGOP analysis can reveal optimal feature transformations, these can be explicitly incorporated into the input data or model architecture. This could drastically accelerate training convergence, reduce computational costs associated with prolonged training, and improve sample efficiency by bypassing the time-consuming grokking phase. This shifts the paradigm from purely black-box learning to a more informed, hybrid approach where domain knowledge or learned feature insights are integrated.

Moreover, AGOP emerges as a powerful model-agnostic analytical tool. Its ability to decouple feature learning from the functional form of the predictor allows for comparative analysis across diverse architectures. This means that insights gained from analyzing the features of a kernel model could potentially inform the design or understanding of a neural network, and vice-versa. This cross-model interpretability could be invaluable for researchers seeking to understand the common principles underlying effective learning across different paradigms. The speaker's comment during the Q&A that AGOP has been computed on MLPs, CNNs, and transformers for language and vision tasks, showing high correlation with weight covariance, suggests its broader applicability beyond modular arithmetic.

Finally, the finding that grokking is not tied to gradient descent-based optimization methods suggests that the phenomenon is more fundamental to the process of discovering robust features rather than a quirk of specific optimization landscapes. This implies that even with advancements in optimizers, the challenge of delayed generalization might persist if the underlying feature learning process itself is inherently gradual. Understanding these limitations is crucial for setting realistic expectations and designing more robust training strategies.

The primary tradeoffs and limitations involve the a posteriori nature of the proposed progress metrics. While effective for analysis, they require access to a fully generalized model to establish the "final AGOP" or the target circulant structure. Developing a priori indicators of feature emergence that don't rely on a "future" generalized state remains an open challenge. Additionally, computing the AGOP for very high-dimensional inputs can be computationally intensive (a D x D matrix). However, the ability to compute AGOP with respect to hidden layer activations, as mentioned in the Q&A, suggests potential strategies for managing complexity in deep architectures.

Key Takeaways

• Grokking is a General Phenomenon: The delayed generalization known as grokking is not exclusive to neural networks; non-neural models like Kernel Recursive Feature Machines (RFM) also exhibit this behavior on modular arithmetic tasks.
• Grokking Reflects Gradual Feature Learning: Instead of a sudden, inexplicable event, grokking is shown to be a manifestation of gradual, continuous feature learning that occurs even when traditional metrics like accuracy and loss remain stagnant.
• Average Gradient Outer Product (AGOP) is Key: AGOP serves as a powerful, model-agnostic tool to analyze and understand the features learned by both neural and non-neural models, revealing the underlying structural transformations.
• Unified Feature Structures for Modular Arithmetic: For modular arithmetic, both model types learn block circulant feature transformations (or their discrete logarithm-reordered equivalents for multiplication), indicating a common computational strategy.
• Eliminating Grokking with A Priori Features: By pre-applying learned (e.g., random circulant) feature transformations to input data, the grokking delay can be completely eliminated, leading to immediate generalization and improved sample efficiency.
• New Metrics for Emergent Learning: Traditional accuracy and loss are insufficient; novel a posteriori measures like circulant deviation and AGOP alignment are proposed to accurately track feature learning progress during the grokking phase.

About the Speaker(s)

The oral presentation was delivered by Neil Mallinar. He presented the work on behalf of his co-authors: Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, and Misha Belkin. The specific titles and affiliations of the speakers were not detailed in the transcript, but the collaborative nature of the research highlights a team effort in advancing the understanding of emergent phenomena in machine learning.