PaperlayerView exact revision on arXiv

Source-linked AI summary

Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets

Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, Vedant Misra

arXiv:2201.02177v1cs.LG

TL;DR

The paper asks how overparameterized neural networks generalize beyond memorizing small training sets. It studies this question with tokenized algorithmic datasets and finds late, sometimes abrupt generalization, while smaller datasets require substantially more optimization and some operations remain within-budget failures.

  • Problem

    Generalization beyond memorization is poorly understood, especially for overparameterized neural networks, and small algorithmic datasets offer a controlled setting to study it.

  • Method

    The authors train a small transformer on tokenized equations from binary-operation tables and vary dataset size, operations, and optimization interventions.

  • Results

    Neural networks sometimes grok patterns after severe overfitting, while smaller datasets require rapidly increasing optimization time and weight decay substantially improves data efficiency.

  • Takeaways & Limitations

    Small algorithmic datasets can serve as reproducible testbeds for investigating late generalization, optimization effects, and learned mathematical structure.

  • Takeaways & Limitations

    Some operations did not generalize within the allowed optimization budget even with up to 95% of the data, and the broader transfer of the compute–data trade-off remains for future work.

Abstract

from arXiv · show

In this paper we propose to study generalization of neural networks on small algorithmically generated datasets. In this setting, questions about data efficiency, memorization, generalization, and speed of learning can be studied in great detail. In some situations we show that neural networks learn through a process of "grokking" a pattern in the data, improving generalization performance from random chance level to perfect generalization, and that this improvement in generalization can happen well past the point of overfitting. We also study generalization as a function of dataset size and find that smaller datasets require increasing amounts of optimization for generalization. We argue that these datasets provide a fertile ground for studying a poorly understood aspect of deep learning: generalization of overparametrized neural networks beyond memorization of the finite training dataset.

1 INTRODUCTION

The paper uses small algorithmically generated datasets to study generalization beyond memorization, including cases where validation improves long after training has overfit. These datasets also expose how data size and optimization affect generalization.

  • Motivation: Small binary-operation tables provide controlled testbeds for studying overparameterized-network generalization beyond finite-dataset memorization.Experiments can be quickly reproduced on a single GPU.
  • Dataset setting: The datasets encode equations a ◦ b = c using abstract symbols without internal structure, requiring networks to infer properties from interactions.This is analogous to filling missing entries in an operation table.
  • Contributions: Neural networks can generalize to empty slots in varied binary-operation tables, sometimes after severe overfitting through a late transition called grokking.Validation accuracy can rise from chance toward perfect generalization well after training accuracy becomes nearly perfect.
  • Contributions: The paper measures data-efficiency curves, optimization-time requirements, and the effects of optimization details across algorithmic tasks.It also visualizes learned symbol embeddings for recognizable mathematical structure.

2 METHOD

The experiments train a small transformer on tokenized equations representing binary operations. Each operand, operator, equality sign, and result is encoded as a separate token.

  • Model and data: All experiments use a small transformer trained on equations of the form a ◦ b = c.The architecture, operations, hyperparameters, and tokenization are described in Appendix A.1.
  • Tokenization: Each of “a”, “◦”, “b”, “=”, and “c” is represented as a separate token.

3 EXPERIMENTS

Experiments show that generalization can emerge long after overfitting, while reduced data preserves eventual performance but sharply increases optimization time. Optimization choices, operation structure, and task complexity affect whether generalization is reached.

  • 3.1 GENERALIZATION BEYOND OVERFITTING: Validation accuracy may begin rising only after 1000 times more optimization steps than training accuracy needs to approach optimal.This late improvement is accompanied by double descent of validation loss in the modular-division example.
  • 3.1 GENERALIZATION BEYOND OVERFITTING: The late-generalization behavior is typical near the minimal dataset sizes that generalize within the allotted optimization budget.For larger datasets, training and validation curves track each other more closely.
  • 3.1.1 LEARNING TIME CURVES: 100% converged performance can remain constant across a range of training-set sizes while optimization time increases quickly as data decreases.
  • 3.3 ABLATIONS AND TRICKS: Different optimization algorithms produce different generalization levels within 10^5 steps, with weight decay improving generalization most on S5.Suboptimal hyperparameters can severely limit generalization, although some methods generalize at high data percentages.
  • 3.1.1 LEARNING TIME CURVES: A 1% data reduction near 25–30% of the dataset increases median time to generalization by 40–50%.This measures median optimization steps until validation performance first reaches 99% for the product in abstract group S5.
  • 3.2 GROKKING ON A VARIETY OF PROBLEMS: Operations with symmetric operands tend to require less data for generalization than closely related nonsymmetric operations.The authors suggest this effect may partly depend on transformer positional embeddings.
  • 3.2 GROKKING ON A VARIETY OF PROBLEMS: Some operations fail to generalize within the optimization budget even with up to 95% of data, leaving models effectively memorizing the training set.
  • 3.2 GROKKING ON A VARIETY OF PROBLEMS: Generalization can occur for an operation that mixes additive and multiplicative rules depending on the parity of one operand.This demonstrates generalization beyond operations cleanly interpretable through group or ring operations.

4 DISCUSSION

The paper presents small algorithmic datasets as a setting where unusual generalization dynamics and learned mathematical structure can be examined. It also identifies a compute–data trade-off that remains open for other datasets.

  • Small binary-operation tables exhibit striking effects including double descent, late generalization, and improved generalization from weight decay.These datasets are proposed as a useful setting for investigating generalization, including future tests of minima-flatness measures.
  • Embedding visualizations reveal natural mathematical structure, with modular-arithmetic representations tending toward circles or cylinders.The authors speculate that such visualizations may eventually help build intuitions about novel mathematical objects.
  • Reducing training-dataset size sharply increases the optimization steps needed to reach a given performance level.The authors identify testing whether this compute–data trade-off appears on other datasets as future work.

A.1.1 BINARY OPERATIONS

The experiments use binary-operation tables over discrete abstract symbols, covering modular arithmetic, polynomial operations, and permutation-group operations. Training examples are equation tokens sampled from the complete operation table.

  • The study evaluates binary operations including modular addition, subtraction, division, polynomial expressions, and several operations on S5.The polynomial examples use p = 97, while the S5 operations include composition, conjugation, and a nonstandard product.
  • One mixed operation uses modular division when y is odd and modular subtraction when y is even.
  • Each operation is represented by equations of the form ⟨x⟩⟨op⟩⟨y⟩⟨=⟩⟨x ◦y⟩, with one token for each element symbol.This converts completion of the operation table into prediction of missing equation answers.
  • For each operation, a random fraction of all available equations forms the training set, while the remaining equations form validation data.

A.1.2 MODEL AND OPTIMIZATION

The experiments use a small decoder-only transformer and compare optimization configurations spanning batch size, noise, dropout, weight decay, and training budgets. Runs use repeated random seeds and extended optimization for late-generalization analyses.

  • The model is a 2-layer, width-128, 4-head decoder-only transformer with about 4 · 105 non-embedding parameters.Causal attention masking is used, and loss and accuracy are calculated only on the answer portion of each equation.
  • Most experiments use AdamW with learning rate 10−3, weight decay 1, β1 = 0.9, β2 = 0.98, warmup, and minibatches capped at 512 or half the training set.
  • Optimization variants include full-batch Adam, gradient noise, residual dropout, weight decay toward the initialization, and Gaussian weight noise.The gradient-noise variant adds unit-Gaussian noise to each update direction.
  • The Section 3.1 analyses use up to 106 optimization steps without weight decay to emphasize how late generalization can begin.Section 3.1.1 uses a 5·105-step budget to capture increases in time to perfect generalization.
  • Experiments generally use 3 random seeds per dataset size, while Section 3.1.1 aggregates results over 7 random seeds.

A.2 ADDITIONAL FIGURES

The additional figures show loss dynamics for modular division and an example of a binary-operation table solved by the network. Modular-division validation loss rises before beginning a second descent.

  • Figure 4 presents the loss curves corresponding to the accuracy curves in Figure 1.
  • Figure 5 provides an example of a binary-operation table that the network can solve.
  • For modular division, validation loss increases from 102 to about 105 optimization steps before beginning a second descent.

A.3 RELATED WORK

The paper studies generalization dynamics on small algorithmically generated datasets, emphasizing behavior beyond memorization and fixed performance snapshots. It distinguishes its generalization-focused “grokking” phenomenon from related optimization and complexity findings.

  • The paper examines training and generalization dynamics on small algorithmic datasets, extending earlier benchmark work on symbolic and algorithmic reasoning.
  • Prior algorithmic-dataset studies often report point estimates for particular architectures or training methods, whereas this work tracks generalization after complete training-set memorization.
  • The paper distinguishes its grokking phenomenon, which concerns delayed generalization, from earlier grok-like curves attributed to optimization difficulty.
  • Compared with procedurally generated mathematics tasks, the paper’s simpler binary-operation problems are intended to make late-generalization phenomena observable with fewer samples.
  • Earlier work found flatness-based measures most predictive of generalization, motivating the paper’s conjecture that SGD noise may favor flatter, simpler solutions.
  • The paper relates its findings to work showing that neural networks can interpolate arbitrary labels yet generalize with meaningful labels and suitable optimization.
  • It also connects late validation-loss descent to double descent, while emphasizing a second descent in validation loss rather than accuracy.

A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS

The outlier experiments test whether networks denoise mislabeled equations or memorize them as exceptions. Increasing outliers reduces the conditions under which models generalize, although small numbers have limited impact.

  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: The experiments replace k training answers with answers from other randomly sampled equations, leaving the remaining training and validation equations unchanged.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: The experiment contrasts denoising the outliers with interpolating them as exceptions that require a more complicated function.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: 100% training accuracy is reached in every experiment, and changing k does not substantially alter when training interpolation occurs.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: Increasing k narrows the range of training-data percentages for which optimization converges to models that generalize.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: Up to 1000 outliers have little effect on data efficiency, whereas larger numbers substantially hinder generalization performance.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: The results support the view that the network and optimizer can memorize all training labels, so generalization requires explaining why a simpler pattern is selected.
  • A.4 GENERALIZATION WITH MEMORIZING SEVERAL OUTLIERS: A preliminary investigation found that sharpness of the trained network’s minimum appears predictive of generalization on one S5 composition dataset.
Loading 2201.02177v1…