Slower Generalization, Faster Memorization: A Sweet Spot in Algorithmic Learning
For researchers studying grokking and generalization in structured-output tasks, this work reveals a counterintuitive divergence between dataset size and convergence speed, challenging common assumptions about the benefits of larger datasets.
In a structured-output task (Needleman-Wunsch matrix generation), small Transformers achieve fastest validation exact-match accuracy at an intermediate dataset size, not the largest, contradicting the intuition that more data accelerates generalization. Past this sweet spot, generalization requires more gradient updates, while larger datasets can speed up training accuracy in the regime where partial competence first appears.
Critical-data-size accounts of grokking suggest a natural post-threshold intuition: once training data is sufficient to identify the underlying rule, additional data should accelerate validation convergence. We show that this intuition can fail in a controlled structured-output task. In Needleman--Wunsch (NW) matrix generation, small Transformers reach high validation exact-match accuracy fastest at an intermediate dataset size, not at the largest one. Past this dataset-size sweet spot, generalization remains achievable but requires more gradient updates. Conversely, in the regime where partial validation competence first appears, larger datasets can require fewer updates to reach high training accuracy, suggesting that emerging rule structure can accelerate fitting beyond example-wise memorization. A multiplication baseline does not show the same post-threshold slowdown. These results separate the critical data size for the onset of generalization from the dataset size that optimizes update-based convergence, and identify structured-output tasks where learning the rule and completing exact-fitting can diverge.