“Grokked Transformers Are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization”, Boshi Wang, Xiang Yue, Yu Su, Huan Sun2024-05-23 (GPT-2, GPT-4 nonfiction, PaLM 2, grokking (NN); backlinks)⁠:

[code, Twitter] We study whether transformers can learn to implicitly reason over parametric knowledge, a skill that even the most capable language models struggle with. Focusing on two representative reasoning types, composition and comparison, we consistently find that transformers can learn implicit reasoning, but only through grokking, ie. extended training far beyond overfitting.

The levels of generalization also vary across reasoning types: when faced with out-of-distribution examples, transformers fail to systematically generalize for composition but succeed for comparison. We delve into the model’s internals throughout training, conducting analytical experiments that reveal: (1) the mechanism behind grokking, such as the formation of the generalizing circuit and its relation to the relative efficiency of generalizing and memorizing circuits, and (2) the connection between systematicity and the configuration of the generalizing circuit.

Our findings guide data and training setup to better induce implicit reasoning and suggest potential improvements to the transformer architecture, such as encouraging cross-layer knowledge sharing.

Furthermore, we demonstrate that for a challenging reasoning task with a large search space, GPT-4-Turbo and Gemini-1.5-Pro based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation, while a fully grokked transformer can achieve near-perfect accuracy, showcasing the power of parametric memory for complex reasoning.

[Possible implication: scaling up LLMs may never trigger grokking without data-pruning if the scaled-up datasets simply maintain the ratio of memorable:learnable datapoints, leaving them in the inferior-generalizing regimes, requiring vastly larger sample sizes? See also Huang et al 2024 on the harm a memorization task can do.]

…On two representative reasoning types, composition and comparison, we consistently observe the ubiquitous role of grokking in transformer’s acquisition of implicit reasoning. Further experiments reveal that the speed of grokking correlates with the ratio between inferred and atomic facts, and depends little on the absolute size of training data. This suggests a simple correction of prior explanations of grokking that the training data distribution, rather than size, may be the actual critical factor behind grokking. Moreover, the systematicity level varies across reasoning types—in the OOD scenario, the model fails to systematically generalize for composition, but succeeds for comparison.

…Inferred/atomic ratio ϕ correlates with generalization speed.

The Figure 2(a) shows the ID accuracy across different ϕ. We omit the other splits since for all settings, the training performance saturates quickly and the OOD accuracy remains at zero as earlier.³ It could be seen that the ratio ϕ strongly correlates with the speed of generalization. A very large ratio can push generalization to improve at a similar pace as the model fits the training data, reducing the need for extended training.⁴

Training data distribution, instead of training data size, qualitatively influences generalization behavior. When ϕ increases and |ℰ| holds constant, the size of training data also gets larger. Prior studies hypothesize that training data size plays a central role in order for grokking to happen. In particular, previous work connects grokking with the notion of critical data size (CDS)^{33, 61, 78, 21}, where it is hypothesized that CDS marks the shift from memorization to generalization (via grokking), and the speed of generalization improves as the training data further scales.

However, results from our controlled experiments seem to contradict such a hypothesis. Figure 2(b) shows the results of varying |ℰ| with a fixed ϕ = 9.0, where we change the horizontal axis from optimization step to epoch for better visualization.⁵ When fixing the ratio ϕ, the training data size does not qualitatively affect the model’s generalization. Specifically, scaling the data affects neither the relative speed of ID generalization and training improvement (as seen by the rather constant “gap” between train_inferred_ID and test_inferred_ID curves), nor the systematicity level (OOD performance stays zero). We also run the experiments across different ϕ and find the results to be consistent.

This suggests that critical data “distribution”, not size, may be the actual deciding factor behind grokking and generalization. In addition, we find that scaling up the model size also does not qualitatively change the generalization behaviors observed here (Appendix B), and the main pattern is that larger models converge in fewer optimization steps, which shares with prior findings^{60, 28}.

---

…So what happens during grokking, why does it happen, and why do transformers exhibit different levels of systematicity in generalization?

To answer these questions, we analyze the model internals throughout grokking.

We find distinct generalizing circuits for the two reasoning types, and strong patterns in the evolution of the circuits that explain the underlying mechanisms behind grokking and variations in systematicity.

Notably, for the comparison task, the transformer gradually forms a parallel circuit which allows the model to store/access the ID/OOD atomic facts in the same region, enabling systematicity to happen.

…This also indicates that, before grokking, the model is very likely mostly memorizing the examples in train_inferred_ID by directly associating (h,r₁,r₂) with t, without going through the first hop.

---

…Why does grokking happen? These observations suggest a natural explanation of why grokking happens through the lens of circuit efficiency. Specifically, as illustrated above, there exist both a memorizing circuit C_mem and a generalizing circuit C_gen that can fit the training data. While C_mem is learned first (which causes training performance to saturate quickly), C_gen is relatively more efficient, in the sense that it could fit the data with a lower complexity. To see this, we can compare the amount of facts C_mem and C_gen need to store (denoted as N_mem and N_gen) as a proxy for their complexity.⁷ C_mem stores both atomic facts and inferred facts in the weights. C_gen (Figure 4(a)) stores the atomic facts in the lower layers, and another copy of the atomic facts that appear as the second hop in the inferred facts in the upper layers.

As the inferred/atomic ratio ϕ increases, N_mem would increase rapidly while N_gen increases slowly and is always bounded by two times the total amount of atomic facts, and hence, the relative efficiency of C_gen increases. In the long run, the model will be incentivized to transition from C_mem to C_gen due to implicit bias of the optimization⁵³ and explicit regularization such as weight decay which prefers more efficient circuits, and the transition would happen faster as ϕ increases.

This also explains why the training data size does not affect the speed of grokking, since solely increasing the size does not change the relative efficiency of C_mem and C_gen. The explanation also implies that a larger regularization factor should accelerate grokking (and vice versa), which we confirm by varying the degree of weight decay (Appendix E.1).

…Optimization is done by AdamW with learning rate 10⁻⁴, batch size 512, weight decay 0.1 and 2,000 warm-up steps with a linear schedule. More details are included in Appendix A.

[commentary on relationship to Berglund et al 2023 & Treutlein et al 2024.]

Backlinks: