Research map / Large language models / Retrieval and knowledge
Grokking in the Wild: Data Augmentation for Real-World Multi-Hop Reasoning with Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends grokking to real-world multi-hop QA by synthetically augmenting 2WikiMultiHopQA (Wiki2Hop) with LLM-generated atomic and inferred facts to raise the inferred-to-atomic ratio, then training a GPT-2-small-style Transformer from scratch long enough to grok.
Problem
Transformers struggle with implicit multi-step factual reasoning on real-world sparse knowledge; grokking has only been shown on small synthetic tasks, and real KGs lack enough inferred facts relative to atomic facts to trigger it.
Approach
Augment sparse real-world KG-style facts from 2WikiMultiHopQA with LLM-generated atomic and inferred facts so the per-relation ratio phi_r exceeds the grokking threshold, then train a small GPT-2-style model long enough to exhibit a late-phase OOD accuracy jump.
Key results
- Grokked GPT2-small reaches 0.97 average IID and 0.52 average OOD accuracy on structured 2WikiMultiHopQA, with comparison OOD at 0.96 and IID at 1.00.
- Baselines GPT-4o (0.56 avg) and o1-mini (0.60 avg) are outperformed on average, though o1-mini is better on composition OOD (0.32 vs 0.07 for grokked model).
- Grokked model reaches only 0.07 OOD accuracy on the composition task, with no late-phase OOD jump there.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add synthetic data augmentation and train until grokking | GPT2-Small · 2WikiMultiHopQA Comparison OOD | accuracy: 0.59 -> 0.96 | Helped |
| Add synthetic data augmentation and train until grokking | GPT2-Small · 2WikiMultiHopQA Composition OOD | accuracy: 0.03 -> 0.07 | Mixed |
| Add synthetic data augmentation and train until grokking | GPT2-Small · 2WikiMultiHopQA Average OOD | accuracy: 0.31 -> 0.52 | Helped |
| Train on original structured comparison data without augmentation (ablation) | GPT2-Small · 2WikiMultiHopQA Comparison OOD | accuracy | No clear effect |
| Use structured compositional tasks instead of structured comparison tasks | Grokked GPT2-Small · 2WikiMultiHopQA OOD | accuracy: 0.96 -> 0.07 | Hurt |
| Use unstructured full-paragraph text instead of structured triplets | Grokked GPT2-Small · 2WikiMultiHopQA Comparison OOD | accuracy | Hurt |
| Compare Grokked GPT2-Small with GPT-4o | 2WikiMultiHopQA Comparison · 2WikiMultiHopQA Comparison | accuracy: 0.87 -> 1.00 | Helped |
| Compare Grokked GPT2-Small with o1-mini | 2WikiMultiHopQA Composition · 2WikiMultiHopQA Composition | accuracy: 0.32 -> 0.93; 0.07 | Mixed |
Surprising findings
- Factually incorrect synthetic data was found to strengthen emergent reasoning circuits rather than degrade accuracy.
- Composition OOD accuracy stays low (0.07) despite augmentation, showing grokking does not transfer uniformly across task types.
Limitations
- Composition OOD accuracy remains very low (0.07), so grokking gains are limited to comparison queries.
- Unstructured text slows convergence and yields only modest OOD gains; data synthesis relies on hallucination-prone LLM generation with factuality risks.
- Training for hundreds of thousands of steps is computationally expensive; evaluation is limited to 2WikiMultiHopQA.
Experimental setup
Datasets: 2WikiMultiHopQA Wiki2Hop
Benchmarks: 2WikiMultiHopQA
Metrics: accuracy (in-distribution and out-of-distribution)
Baselines: GPT2-Small GPT-4o o1-mini
Builds on: grokking 2WikiMultiHopQA AdamW
Closest papers by meaning
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization (NeurIPS 2024 · similarity 0.91)
- Learning from Synthetic Data Improves Multi-hop Reasoning (ICLR 2026 · similarity 0.89)
- How do Transformers Learn Implicit Reasoning? (NeurIPS 2025 · similarity 0.88)
- Finding the Minimal Parameter Budget for Implicit Reasoning: A Data Complexity Driven Scaling Law for Language Models (ICML 2026 · similarity 0.88)
- Scaling Reasoning Hop Exposes Weaknesses: Demystifying and Improving Hop Generalization in Large Language Models (ICLR 2026 · similarity 0.88)
- Boule or Baguette? A Study on Task Topology, Length Generalization, and the Benefit of Reasoning Trace (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.