Omnigrok: Grokking Beyond Algorithmic Data

Ziming Liu; Eric J. Michaud; Max Tegmark

Source-linked AI summary

Omnigrok: Grokking Beyond Algorithmic Data

Ziming Liu, Eric J. Michaud, Max Tegmark

arXiv:2210.01117v2Oct 3, 2022cs.LGcs.AIphysics.data-anstat.MEstat.ML

TL;DR

Grokking raises questions about why generalization can occur long after training-set overfitting and whether the phenomenon extends beyond algorithmic data. This paper analyzes neural loss landscapes, identifies the LU mechanism, and reports grokking across images, language, and molecules while linking its dramaticness to representation learning.

Problem
The paper asks why generalization is delayed after overfitting and whether grokking occurs beyond algorithmic datasets.
Method
The paper analyzes reduced training and test loss landscapes against model weight norm, calling their L- and U-shaped mismatch the LU mechanism.
Results
Grokking occurs across a wide range of tasks, including image classification, sentiment analysis, and molecule property prediction, though signals are usually less dramatic than for algorithmic datasets.
Takeaways & Limitations
The paper concludes that grokking’s dramaticness depends on how much a task relies on learning representations.
Takeaways & Limitations
Clear grokking signals have not yet been observed for large language models, possibly because language representations are more complex and pre-training avoids learning them from scratch.

Abstract

from arXiv · show

Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive. We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the mismatch between training and test losses as the cause for grokking. We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively. This simple mechanism can nicely explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc. Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules. In the reverse direction, we are able to eliminate grokking for algorithmic datasets. We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.

1 INTRODUCTION

The paper frames grokking as delayed generalization caused by a mismatch between training and test loss landscapes, and uses this view to explain its occurrence across tasks.

The paper asks why generalization can be delayed after overfitting and whether grokking occurs beyond algorithmic datasets.
The authors analyze neural loss landscapes to address grokking’s origin and prevalence.
The LU mechanism describes reduced training and test losses versus weight norm as L-shaped and U-shaped, respectively.
Grokking is demonstrated on image classification, sentiment analysis, and molecule property prediction, although its signals are usually less dramatic than on algorithmic datasets.
Landscape analysis explains dependencies on data size, weight decay, and representation emergence.
The paper reports grokking on realistic datasets including MNIST, IMDb reviews, and QM9, while constrained weight norm can almost eliminate grokking.

2 THE LU MECHANISM FOR GROKKING

The LU mechanism models a Goldilocks weight-norm region where generalization is best: training loss falls and flattens, while test loss forms a U-shape. Regularization then drives delayed radial motion toward that region.

A spherical Goldilocks zone around weight norm wc contains solutions with better generalization than solutions outside it.
Reduced training loss decreases with weight norm before flattening, whereas reduced test loss is U-shaped and higher both below and above wc.
With an initially large weight norm, training rapidly reaches an overfitting solution, while weak regularization slowly moves weights toward wc and produces grokking.
Under weight decay γ, the generalization time scales as t ≈ ln(w0/wc)/γ ∝ γ^-1 when training loss is negligible after overfitting.
Increasing initialization scales can induce grokking, while small initializations below wc can generalize fast regardless of regularization.
For representations, wc(bad representation) > wc(good representation), and decoding requires O(p2) versus O(p) examples for bad versus good representations.

3 GROKKING FOR A TEACHER-STUDENT SETUP

In a teacher-student setup, reduced losses exhibit the LU shape, and initialization scale together with weight decay determines whether training fails to generalize, groks, or generalizes quickly.

The teacher-student experiment uses a shared 5-100-100-5 tanh MLP architecture with independently seeded teacher and student networks.
The reduced test loss is U-shaped with a minimum near α ≈1, while reduced training loss forms an L-shape as α increases.
Large initialization α = 2.0 yields no generalization without regularization, grokking with small regularization, and faster generalization with large regularization.
Small initialization α = 0.5 always generalizes fast regardless of regularization.
The time to reach 95% training accuracy is independent of weight decay, whereas the time to reach 95% test accuracy is inversely proportional to weight decay.

4 OMNIGROK: GROKKING FOR MORE INTERESTING TASKS

The paper transfers loss-landscape analysis beyond toy settings, reducing training and test losses to weight norm and data size. Across MNIST, IMDb reviews, and QM9, small datasets and large initializations expose LU-shaped mismatches that can produce delayed generalization.

Image classification: Reduced training and test losses are analyzed as functions of weight norm w and data size N after minimizing over angular weight directions.This reduction is used for MNIST, IMDb reviews, and QM9.
Image classification: Larger initializations can induce MNIST grokking, while larger training sets broaden the Goldilocks zone and reduce or eliminate it.The analysis connects delayed generalization to low training error, elevated test error, and later norm reduction through regularization.
Image classification: MNIST grokking was induced by reducing training data from 60k to 1k samples and scaling the initialization by α > 1.The resulting model was a depth-3 width-200 ReLU MLP trained with AdamW and MSE loss.
Sentiment analysis: On IMDb reviews, 1k examples with large initialization α = 6 produced delayed, weak generalization, whereas standard initialization showed no grokking.With weight decay 1, overfitting completed within 10^2 steps and generalization began around 10^3 steps.
Molecules: On QM9 molecular polarizability, small data and large initialization α = 3 made grokking possible, while large data produced similarly U-shaped training and test losses.The figure reports 100 versus 3000 examples for reduced losses and 200 training samples for the initialization comparison.

5 REPRESENTATION IS KEY TO GROKKING

Representation learning distinguishes dramatic grokking on algorithmic tasks from its weak or absent appearance on standard tasks. The LU landscape explains how mismatched training and test losses, weight norm dynamics, and representation quality produce or suppress grokking.

Landscape construction: The reduced landscape uses representation messiness m and weight norm w, with representation interpolating between linear and random representations.The reduced training and test losses are defined as functions of these two variables, with weight decay treated separately in the dynamics.
Algorithmic datasets: In algorithmic landscapes, small-weight regions have high training loss and fast dynamics, while large-weight regions have low training loss and slow dynamics.Training and test landscapes differ especially in the low-training-loss, large-weight region, creating the setting for delayed generalization.
Algorithmic datasets: Grokking arises from slow traversal through the large-weight region and along its boundary toward the generalizing solution.The total traversal time is t = (L + htanθ)/(ηDγ), so larger weight decay γ shortens the delay.
Algorithmic datasets: Larger training sets can accelerate grokking, while training sizes 25 and 35 in the toy analysis never reach generalization.For training sizes 45 and 55, the contours connect to the generalizing solution, with θ55 < θ45 and therefore t55 < t45.
Suppressing grokking: Constraining weight norm can nearly eliminate grokking by bringing train and test accuracy learning curves together.The transformer experiment reports this effect under constant-weight-norm optimization, while MNIST landscapes show no significant training–test mismatch and no grokking path.
Representation dependence: Algorithmic datasets exhibit strong dependence of test performance on representation quality, whereas MNIST depends weakly on representation quality.This difference is identified as the key distinction underlying grokking versus non-grokking behavior.

6 RELATION TO RELATED WORKS

The paper relates grokking to prior explanations involving representations, phase changes, and regularization, while distinguishing its weight-norm loss-landscape view from double descent.

Grokking was previously attributed to slow formation of good representations, lower-loss generalizable solutions, or a phase change from limited data and regularization.
The paper’s U-shaped test loss does not conflict with double descent because it is plotted against weight norm rather than parameter count.
Large initializations can overfit data easily, while weight decay is associated with a grokking timescale proportional to 1/λ.

7 CONCLUSIONS

The study explains grokking through mismatched training and test loss landscapes, extends the phenomenon beyond algorithmic datasets, and links its dramaticness to representation learning.

Grokking originates from a mismatch between training and test losses, termed the LU mechanism.
Grokking occurs across varied models and datasets, although its signature is usually most dramatic for algorithmic datasets.
The dramaticness of grokking depends on how much the task relies on learning representations.
Reduced landscape analysis characterizes interactions between data and models and supports analysis of representation learning.

A EXPERIMENT DETAILS

The experiments cover sentiment classification on IMDb, molecular property prediction on QM9, and MNIST classification using distinct neural architectures and training setups.

IMDb sentiment analysis uses 50k reviews, the 1000 most frequent words, length-500 padded token arrays, and a two-layer LSTM.The LSTM uses embedding dimension 64 and hidden dimension 128.
QM9 molecular experiments predict isotropic polarizability with a two-layer ReLU graph convolutional network and a linear output layer.Training uses Adam with learning rate 0.001, MSE loss, and a 50/50 train/test split.
MNIST experiments train width-200, depth-3 ReLU MLPs with MSE loss and AdamW at learning rate 0.001.The batch size is 200.

B REDUCED LOSS FOR MODULAR ADDITION WITH TRANSFORMERS

The appendix examines reduced loss landscapes for one-layer transformers trained on modular addition, confirming the LU shape across initialization and train-set conditions.

The experiment uses a one-layer transformer for modular addition with p = 113, dmodel = 128, four attention heads, and dmlp = 512.The model uses ReLU activations and learning rate 0.001 while constraining model weight norm.
The LU shape holds for α ∈[0.1, 4] across varied train-set fractions, although optimization issues may affect some cases.
Figure 9 presents reduced loss landscapes for transformers trained on modular addition, the original setting where grokking was observed.

C TIME TO GENERALIZE VERSUS WEIGHT DECAY

The paper tests whether generalization time scales inversely with weight decay in modular addition and MNIST. The inverse relation holds over task-dependent ranges, while extreme decay values introduce deviations.

The experiments measure training steps needed to reach a specified test-accuracy level while varying AdamW weight decay γ.They use a transformer on modular addition and a ReLU MLP on MNIST, with full training curves reported alongside generalization-time plots.
Across roughly two orders of magnitude, modular-addition generalization time follows t ∝γ^-1 for each initialization seed.Different seeds require different absolute times, but each seed fits the inverse relation well.
For MNIST, t ∝γ^-1 holds approximately when weight decay γ lies between 0.1 and 1.0.Very high decay appears to disrupt optimization, while very low decay generalizes faster than expected, possibly through implicit regularization.
Lower weight decay generally prolongs the overfitting period before eventual generalization.This qualitative pattern motivates testing the inverse relationship between generalization time and γ.

D SECTION 5.1 SETUP

The setup analyzes addition with a decoder and separates landscape, reduced-trajectory, and full-trajectory analyses by which parameters remain trainable. Reduced trajectories rely on scale separation and linear representation evolution, so they may deviate from full dynamics.

Setup: The decoder maps scalar input representations through an MLP with hard-coded addition and predicts 30-dimensional random output vectors.Each input symbol receives a scalar encoding, while the two-hidden-layer network uses a 1-200-200-30 architecture with ReLU activations.
Analysis setups: The paper distinguishes three setups with different trainable-parameter subsets: landscape analysis, reduced trajectory analysis, and full trajectory analysis.This organization clarifies which quantities are fixed or evolved in each analysis.
Landscape analysis: Landscape analysis fixes the representation and weight norm while optimizing only the decoder weight direction.For each fixed pair of weight norm w and representation messiness m, training loss is minimized over the direction.
Reduced trajectory analysis: Reduced trajectory analysis compresses training dynamics to two dimensions by assuming fast direction dynamics and linear representation evolution.These assumptions allow the optimized direction to track the current weight norm and representation state.
Reduced trajectory analysis: The reduced trajectory can deviate from the full trajectory because its assumptions may not hold, although it can still illuminate full-training behavior.The paper links decreasing weight norm with generalization in this analysis.

E MNIST EXPERIMENTS WITH CROSS ENTROPY LOSS

Cross-entropy loss preserves the qualitative LU pattern and delayed generalization on MNIST, but produces a broader generalizing weight range and less clean training dynamics than MSE.

Motivation: The cross-entropy result addresses whether MSE is uniquely responsible for grokking on MNIST.The reported results are qualitatively similar, with some quantitative differences.
Loss landscapes: Cross-entropy experiments retain the LU pattern, with training and test errors resembling L and U shapes for small datasets.For larger datasets, the test-error U becomes more L-like because the mismatch between reduced training and test error is smaller.
Loss landscapes: Cross entropy produces a broader Goldilocks zone than MSE, meaning generalization occurs across a wider weight range.This is the main quantitative difference identified when comparing the two losses.
Training dynamics: Delayed generalization remains observable with cross entropy, although it requires a higher initialization scale α than with MSE.The paper reports this as predicted by the reduced loss landscapes.
Training dynamics: With α = 100 and D = 200, test accuracy rises to 30-40%, plateaus, then increases to ≈75% while training accuracy remains 100%.The initial plateau is above random-guess accuracy, making the dynamics less clean than with MSE.