Grokking in the Wild: Data Augmentation for Real-World Multi-Hop Reasoning with Transformers

TL;DR

This work tackles the gap in real-world multi-hop factual reasoning by leveraging grokking—where neural networks suddenly generalize after learning underlying patterns—through data augmentation that increases the ratio of inferred to atomic facts, $\phi_r$, for each relation. By synthesizing and balancing KG data to exceed grokking thresholds $\phi_G$, the authors enable Transformer models to form internal generalizing circuits, achieving near-perfect OOD accuracy on structured tasks and substantial gains on 2WikiMultiHopQA, even when some synthetic data is factually imperfect. The approach demonstrates that raising $\phi_r$ can unlock implicit multi-hop reasoning in real-world factual datasets, offering a path toward more robust and interpretable reasoning in large language models, while also highlighting limitations related to domain coverage, factuality, and training cost. Overall, the work provides a principled data-centric mechanism to induce emergent reasoning capabilities in Transformers and suggests directions for broader domain applications and deeper mechanistic analyses.

Abstract

Transformers have achieved great success in numerous NLP tasks but continue to exhibit notable gaps in multi-step factual reasoning, especially when real-world knowledge is sparse. Recent advances in grokking have demonstrated that neural networks can transition from memorizing to perfectly generalizing once they detect underlying logical patterns - yet these studies have primarily used small, synthetic tasks. In this paper, for the first time, we extend grokking to real-world factual data and address the challenge of dataset sparsity by augmenting existing knowledge graphs with carefully designed synthetic data to raise the ratio $φ_r$ of inferred facts to atomic facts above the threshold required for grokking. Surprisingly, we find that even factually incorrect synthetic data can strengthen emergent reasoning circuits rather than degrade accuracy, as it forces the model to rely on relational structure rather than memorization. When evaluated on multi-hop reasoning benchmarks, our approach achieves up to 95-100% accuracy on 2WikiMultiHopQA - substantially improving over strong baselines and matching or exceeding current state-of-the-art results. We further provide an in-depth analysis of how increasing $φ_r$ drives the formation of generalizing circuits inside Transformers. Our findings suggest that grokking-based data augmentation can unlock implicit multi-hop reasoning capabilities, opening the door to more robust and interpretable factual reasoning in large-scale language models.

Figures (5)

• Figure 1: Average accuracy on 2WikiMultiHopQA for comparison task. Despite GPT2-small being a model with 124 million parameters, grokked version achieves almost 100% accuracy, beating the most recent gpt-4o and o1-mini models.
• Figure 2: Exemplary Knowledge Graph with Synthesized Data. Four original nodes (black) and three relations (blue) result in two inferred facts. Two additional synthetic nodes and relations (red) extend the amount of inferred facts by four. Consequently, $\phi_I=\phi_2$ increases from $\frac{2}{3}\approx0.66$ to $\frac{6}{5}=1.2$.
• Figure 3: Conceptual difference between ID and OOD: In-distribution (ID, orange) and Out-of-distribution (OOD, red) inferred facts are shown. All green components are seen during training, including all atomic facts (AF) and some inferred facts (IF).
• Figure 4: (a)Accuracy on the comparison task for original and grokked GPT2-small on OOD-test set. We can observe that the OOD accuracy remains almost the same when it reaches 100%, but the difference becomes evident afterward. The grokked transformer continues to improve in accuracy as training progresses, unlike the original version. (b)Training curves for the structured comparison task. IID and OOD behave similarly. (c)The structured compositional task. We see near-perfect ID accuracy but no late-phase jump in OOD test accuracy. (d)Training curves for the unstructured (full paragraph Wikipedia) comparison setting. Complexity slows convergence and limits OOD gains, although ID accuracy still improves significantly.
• Figure 5: Growth of $\phi_{3,r}$ (y-axis) with $|\mathcal{V}|$ (x-axis) for $b_r=2$. The red line is $\phi_{3,r}=\binom{|\mathcal{V}|}{4}\frac{96}{|\mathcal{V}|(|\mathcal{V}|-1)^3}$ (formula values). The green line are empirical values obtained by randomly generating a graph with the same amount of nodes ($|\mathcal{V}|$) and edges ($|\mathcal{V}| b_r$). Due to randomness, our generated/empirical graph can have locally higher branching factors, resulting in overall more inferred facts. Our formula for $\phi_{3,r}$ therefore effectively underestimates the true value of $\phi_{3,r}$. Nevertheless, the formula correctly approximates the shape and order of magnitude. This holds also for other combinations of $b_r$ and $n$.

Theorems & Definitions (19)

• Definition 2.1: Knowledge Graph
• Definition 2.2: Norm over edges
• Definition 2.3: Average branching factor
• Definition 2.4: Atomic Facts
• Example 1: Running Example: Basic KG
• Definition 2.5: Inference Step
• Definition 2.6: Inference Path
• Definition 2.7: N-th Order Deductions
• Example 2: 2-Hop Deductions
• Definition 2.8: Inferred Facts