A team course project by Jiaju Wu, Kairui Li, and Kehan Huang for Selected Topics in Deep Learning Theory.

When does a model that fits its training examples begin to generalize to unseen combinations? We built a common experiment framework for studying this gap on finite tasks with exact algebraic answers.

Code and reproduction instructions · Course report · Report source

Task-design schematic reconstructed from the repository. It illustrates the evaluation setup; it does not represent measured learning curves.

What we built

The PyTorch framework brings several task families into a common final-answer prediction interface:

• Modular arithmetic: addition, subtraction, division, and polynomial operations.
• Permutation tasks: multiplication, conjugation, and related operations in S₅.
• Longer inputs: K-ary modular summation, whose input space grows with sequence length.

A small causal Transformer is the main model. MLP, LSTM, and GRU implementations provide architecture comparisons. The shared training loop makes it possible to vary the optimizer, regularization, and fraction of observed input combinations while retaining the same prediction target.

These are team outputs. The repository and report credit all three authors; they do not establish an individual ownership breakdown for each component.

The central experiment

For modular addition with modulus 97, each example contains operands and an operation marker, followed by an answer. Supervision is applied to the final answer token. We train on a subset of operand pairs and measure validation accuracy on held-out pairs.

The report compares training fractions including 0.3, 0.5, and 0.7. Its representative curves show why training accuracy alone is insufficient: fitting observed pairs can precede improvement on held-out combinations. The size and timing of this gap depend on the training setting.

The architecture and task extensions make the framework useful beyond one modular-addition demonstration. Having support for a task in the code, however, is not by itself evidence that every setting exhibits grokking.

What the evidence supports

The project is an empirical study of delayed generalization and an experiment framework for comparing it across settings. A “memorization followed by rule learning” account is a useful interpretation, but accuracy curves alone do not identify the internal algorithm or prove that a model has acquired a particular representation.

Splits also matter. For a commutative operation, seeing one operand order can help on the reversed order. Stronger claims about learning the full algebraic rule require checking these relationships as well as the final validation score.

Reproducing and extending the study

The repository contains task generation, model implementations, training configuration, and report material. The report source is the reference for the experiments described here. A next step is to compare matched runs across seeds and split designs before attributing differences to any single optimization mechanism.