The Grokking Trap: How Batch Size Changes the Story of Faster AI Learning

The study compares models trained with small batches, large batches, and batches that grow during training. The finding is that a method can look faster when researchers count training steps without being more efficient in the number of examples it processes.

An AI model can memorise its training examples and still struggle with unfamiliar problems that follow the same rules. Thousands of training steps later, its performance can suddenly improve, a pattern researchers call "grokking." Understanding what helps that breakthrough happen means looking closely at how the model learns and how researchers measure its progress.

Roland Fang of Zhejiang University explores this question in "Revisiting dynamic batch scheduling for grokking: the role of batch size, sample budget, and regularization," published in Frontiers in Artificial Intelligence. The study compares models trained with small batches, large batches, and batches that grow during training. The finding is that a method can look faster when researchers count training steps without being more efficient in the number of examples it processes.

Batch size describes how many examples a model processes before updating its internal settings. A model using batches of 512 handles sixteen times as many examples per update as one using batches of 32. Counting those updates alone can give an incomplete picture of how much training each model needed.

How the Study Tested AI's Learning Breakthrough

Fang tested four tasks: modular addition, subtraction, division, and permutation group operations. Modular arithmetic involves numbers that wrap around a fixed limit, similar to hours on a clock, and the permutation task involves combining rearrangements of five items. These controlled problems make it possible to check whether a model can apply learned patterns to combinations it did not train on.

The experiments used a small transformer with two decoder blocks, four attention heads, and a model width of 128. Half the input combinations were used for training and half for validation, with grokking recorded when validation accuracy first reached 99%. Each main experiment was repeated with ten random seeds, giving the researcher a way to see how much results changed across runs.

Five training approaches were compared over 7,500 updates: fixed batches of 32 or 512, plus three schedules that increased batch size from 32 to 512. Those schedules increased batches evenly, exponentially, or through a series of jumps. The standard setup used the AdamW optimiser, a learning rate of 0.001, weight decay of 0.5, and a short learning-rate warmup.

Why Large Batches Looked Faster

The fixed batch of 512 performed strongly when progress was measured in updates. On modular division, all ten runs reached the grokking threshold, taking an average of 3,245 updates and finishing with 99.6% mean accuracy. The small-batch approach reached that threshold in only four runs, with successful runs averaging 4,450 updates.

Large batches reached grokking in every run for addition and subtraction, averaging 2,135 and 3,610 updates, respectively. Growing batches offered useful results, including 99.9% final accuracy for exponential scheduling on addition, but no growing-batch method consistently beat the large fixed batch. The permutation task remained difficult, with the large-batch approach reaching grokking in four runs and each growing-batch approach succeeding in just one.

Across 7,500 updates, batches of 512 processed 3.84 million examples, compared with 240,000 for batches of 32, counting repeated encounters with training examples. Fang reran the comparisons with every approach limited to 240,000 processed examples, equivalent to 960,000 input tokens.

Under that shared budget, the large fixed batch never reached grokking on any task and received only about 469 updates. Its division accuracy fell to 1.6%, compared with 89.7% for the small fixed batch, and exponential scheduling achieved 98.8% accuracy on addition with nine successful runs. No approach reached grokking on the permutation task under this budget, showing how strongly the outcome depended on both the task and the training allowance.

The Settings That Changed the Results

Batch size was only part of the explanation, with additional experiments examining weight decay, model width, and the largest permitted batch. Weight decay discourages model weights from growing too large and can help reduce reliance on memorised examples. None of the tested strategies reached grokking without weight decay in the division and addition experiments.

Stronger weight decay did not guarantee better performance at the end of training. In the focused division tests, increasing it to 1.0 helped the large-batch model reach grokking after 1,720 updates, but its final mean accuracy dropped to 76.9%. Reaching 99% accuracy once did not mean a model would maintain that performance.

Increasing model width brought no consistent improvement, and the widest model tested with exponential scheduling finished with particularly variable accuracy. Limiting that schedule to batches of 128 or 256 produced final division accuracies of 99.7% and 99.9%, compared with 84.3% when batches grew to 512. Separate checks found that removing warmup barely changed the time to grokking, reducing model width delayed it by 28%, and starting the linear schedule at 64 improved its performance measured in updates.

What This Means for AI Training Research

Randomness made a substantial difference, even when the training approach stayed the same. Final division accuracy under the stepped schedule ranged from 3.1% to 100% across runs, making a single successful experiment a poor basis for a broad claim. The results fit earlier explanations involving competition between memorisation and generalisation inside a network, but this study did not directly establish that mechanism.

These findings come from small models solving controlled algorithmic problems, which limits what they establish about larger AI systems and everyday datasets. The shared example budget was tied to the small-batch baseline, and larger budgets or longer training could change the rankings. The additional settings tests used five seeds per configuration and did not explore every learning rate, optimiser, or training-data split.

Fang recommends reporting updates, processed examples, and processed tokens together when describing how quickly a model reaches grokking. Future research could test several training and computing budgets, more complex tasks, and batch sizes that adjust in response to learning progress. For readers assessing claims about faster AI training, the useful question is how much work sits behind the reported improvement.