Large Batches Improve Grokking Speed, Not Data Efficiency

A new study published in Frontiers in Artificial Intelligence challenges the prevailing assumption that larger batch sizes inherently improve the optimization dynamics of neural networks during a phenomenon known as grokking. The research, led by Roland Fang of Zhejiang University, revisits dynamic batch size scheduling to determine whether the apparent benefits of larger batches stem from superior optimization or simply from processing more training examples per update. The findings indicate that while large batches significantly reduce the number of update steps required to achieve generalization, they do not necessarily improve sample efficiency when the total number of processed examples is held constant.

Grokking is defined in the literature as the sudden generalization of neural networks long after they have overfit the training data. This phenomenon has inspired various training strategies aimed at accelerating its onset, with batch size being a fundamental hyperparameter that influences both optimization dynamics and generalization. The study evaluates three dynamic batch size schedules—linear, exponential, and step-wise increases from 32 to 512—against fixed batch sizes of 32 and 512. The experiments were conducted across four algorithmic tasks: modular division, modular addition, modular subtraction, and permutation group operations, each repeated with 10 random seeds to account for variability.

The experimental design employed a transformer-based model architecture with a model dimension of 128, utilizing the AdamW optimizer with a learning rate of 10^-3 and a weight decay of 0.5. The models were trained for 7,500 update steps, with the grokking point defined as the training step at which the model first achieves 99% validation accuracy. By reporting both conventional update-step metrics and sample-normalized metrics based on the total number of processed training examples, the researchers aimed to disentangle update-step efficiency from sample efficiency, a distinction that is critical because large-batch runs process far more examples per update than small-batch runs.

When measured by update steps, the fixed batch size of 512 demonstrated highly competitive performance, often outperforming the dynamic strategies. On the modular division task, the fixed-512 baseline reached a final accuracy of 99.6% with all 10 runs grokking at a mean of 3,245 steps. Similarly, on modular subtraction, it achieved 100.0% accuracy with all 10 runs grokking. These results suggest that, in terms of raw update-step efficiency, larger batches provide a significant advantage, matching or exceeding the performance of dynamic schedules that gradually increase batch size during training.

However, the interpretation changes substantially when the results are normalized by the total number of examples processed. Under a matched budget of 240,000 training examples, which corresponds to the number of examples processed by the fixed batch size 32 baseline in 7,500 steps, the fixed-512 baseline performed poorly and never grokked on any of the four tasks. In contrast, smaller-batch and dynamic strategies remained competitive on the modular tasks. For instance, on modular division, the fixed-512 baseline obtained only 1.6% validation accuracy with zero grokked runs, while the fixed-32 baseline reached 89.7% accuracy with four out of 10 runs grokking.

The study further reveals that grokking is a highly stochastic phenomenon, with reliability varying substantially across random seeds. The small-batch baseline only grokked in 4 out of 10 runs for modular division and subtraction, highlighting the importance of multi-seed evaluation in grokking research. The high variance observed across seeds, particularly for the step-wise strategy, underscores that single-seed evaluations can be highly misleading and that conclusions drawn from small numbers of seeds may be unreliable.

A focused hyperparameter sensitivity study was conducted to examine the role of weight decay, model size, and maximum batch size in grokking reliability. The results show that grokking depends strongly on weight decay: without weight decay, none of the tested strategies grokked on either modular division or modular addition. Increasing weight decay to 0.5 produced reliable grokking for fixed-512 and exponential scheduling, while weight decay of 1.0 accelerated grokking in update steps but reduced final accuracy, indicating that overly strong regularization can harm final performance even when the 99% threshold is temporarily reached.

The model-size sweep demonstrated that the batch-size effect is not independent of capacity. For fixed-512 on modular division, all tested model widths grokked in 5 out of 5 runs, but the mean grokking step varied depending on the model dimension. The maximum-batch sweep suggested that increasing the final batch size is not monotonically beneficial for the exponential schedule, with maximum batch sizes of 128 and 256 achieving higher final accuracies than 512, even though all settings reached the grokking threshold in all runs.

The authors conclude that batch size magnitude is a major confound in grokking experiments, and that claims about acceleration should report both update-step and sample-normalized metrics. Larger batches improve update-step efficiency but are not necessarily more sample-efficient, and their effect is conditional on regularization and capacity. The study highlights the need for future research to evaluate grokking using both update-based and sample-normalized budgets to accurately assess the impact of training strategies.

Limitations of the study include the specific choice of the matched sample budget, which corresponds to the original fixed-32 baseline and may not reflect trade-offs at larger budgets. The experiments were conducted on algorithmic tasks, and the observed effects may not generalize to other types of machine learning problems or datasets. Additionally, the high variance across seeds and the limited training duration of 7,500 update steps may affect the reliability of conclusions, particularly for the permutation task, which remained challenging for all methods.

This research contributes to the understanding of deep learning dynamics by clarifying the role of batch size in grokking. It emphasizes that dynamic batch scheduling should be evaluated against both large-batch baselines and matched-sample-budget controls. The findings urge the community to report examples- and tokens-to-grokking alongside update steps to provide a more comprehensive picture of training efficiency.

The study was published in Frontiers in Artificial Intelligence, Volume 9, in 2026. The authors acknowledge the importance of multi-seed evaluation and the need for further investigation into the interaction between batch size, weight decay, and learning rate. Future work could explore multiple matched sample and compute budgets to map the trade-off between update-step efficiency and sample efficiency, and extend the study to more complex tasks and datasets to assess the persistence of these trade-offs.

Earlier mechanistic investigations demystified grokking by reverse-engineering algorithmic models. Specifically, research on one-layer transformers trained on modular addition demonstrated that networks implement discrete Fourier transforms to map addition into circular rotations. This process unfolds across three phases: memorization, circuit formation, and cleanup, showing that grokking is gradual amplification rather than an abrupt shift. Theoretical frameworks place grokking within a narrow goldilocks zone between memorization and confusion, requiring constrained data alongside sufficient regularization. Furthermore, network pruning studies show competition between a decaying memorizing subnetwork and an expanding generalizing subnetwork, while norm-separation theories propose that grokking delay scales inversely with the product of learning rate and weight decay.

To evaluate these dynamics empirically, the experimental pipeline implemented a transformer architecture comprising token and positional embedding layers, two decoder blocks, and a final linear prediction layer. Each decoder block integrated a four-head self-attention mechanism, layer normalization, and a feed-forward neural network with a hidden dimension of 512, keeping the core model dimension at 128. For the modular arithmetic problems of division, addition, and subtraction, the experiments employed a prime modulus of 97 with integer operands from 0 to 96, split equally into training and validation sets. Input sequences consisted of four discrete tokens covering two operands, an operation symbol, and a separator symbol, trained using custom PyTorch iterable datasets that adjusted batch sizes at epoch boundaries and dropped trailing samples to preserve uniform batch sizes.

A dedicated ablation study on modular division isolated how individual architectural and training components influence grokking onset and stability across ten random seeds. Eliminating the 25-step linear learning rate warmup produced minimal variation in average grokking speed, reaching the generalization threshold in 4,015 steps compared to 4,040 steps for the default exponential schedule, though final validation accuracy dropped from 92.0 percent to 90.8 percent as one seed failed entirely at 11.0 percent accuracy. Shrinking model capacity by reducing the model dimension from 128 to 64 caused substantial performance degradation, delaying grokking onset by 28 percent to 5,179 steps, decreasing final validation accuracy to 88.8 percent, and lowering success rates to seven out of ten runs. Conversely, raising the starting batch size of the linear schedule from 32 to 64 accelerated grokking to 3,850 steps, lowered run-to-run variance, and boosted final validation accuracy to 99.3 percent.

The experimental results also highlighted sharp performance disparities across specific mathematical domains, particularly on the permutation group operation task involving compositions of a five-element set. Under the step-based training protocol, fixed 32-sample batches achieved an average accuracy of only 2.4 percent without a single grokking run, while dynamic linear, exponential, and step-wise schedules each produced only a single grokked run across ten seeds, reaching accuracies between 12.3 percent and 22.0 percent. The top-performing fixed-512 configuration reached 53.6 percent accuracy and grokked in four of ten runs at 6,525 steps. Under the matched 240,000-example budget, every strategy failed to grok permutations, lingering between 0.3 percent and 2.4 percent accuracy.

Subject of Research: The impact of batch size scheduling on the grokking phenomenon in neural networks, specifically regarding update-step versus sample efficiency.

Article Title: Revisiting dynamic batch scheduling for grokking: the role of batch size, sample budget, and regularization

Article References: Fang, R. (2026). Revisiting dynamic batch scheduling for grokking: the role of batch size, sample budget, and regularization. Frontiers in Artificial Intelligence, 9, Article 1860893. https://doi.org/10.3389/frai.2026.1860893

Image Credits: AI Generated

DOI: 10.3389/frai.2026.1860893

Keywords: grokking, batch size scheduling, neural network optimization, sample efficiency, update-step efficiency, weight decay, transformer architecture, algorithmic tasks, modular arithmetic, permutation group operations, random seed variability, hyperparameter sensitivity, generalization dynamics, deep learning training, Frontiers in Artificial Intelligence