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

We thank the reviewer for the thoughtful and detailed feedback, which has helped us refine both the presentation and scope of the paper. Below, we address each of the raised concerns.

A. On theoretical presentation and unused formalisms: We appreciate your observation regarding the mathematical complexity and notation density in the theory section. In response, we have carefully reviewed all definitions and formal constructs. A dependency analysis shows that each definition contributes to Lemma 1 or 2, either directly or through intermediate derivations. While we agree that the presentation could be simplified, we found that deferring key definitions to the appendix would compromise the readability and cohesion of the main theoretical argument. We will, however, streamline redundant notation (e.g., the unused edge norm) and improve clarity in the final version. Thank you also for catching the typo on line 192; we have corrected it.

B. On in-distribution vs. training data distinction:
Thank you for highlighting the confusion here. We will clarify this distinction in the final version. In short: the training data consists of atomic facts and associated inferred facts used in specific combinations. In-distribution (ID) evaluation involves novel combinations of seen atomic facts, not seen together during training. For example, if the model saw the atomic facts for "Avignon Rocher des Doms" and "Paris Louvre Museum" in training, an ID question might involve "Avignon Rocher des Doms" and the "Eiffel Tower", a new combination of familiar elements. This contrasts with out-of-distribution (OOD) examples, which include entirely new atomic facts.

C. On structured vs. unstructured settings and realism:
We agree that structured settings simplify linguistic complexity. Our core claim is not that we have solved the problem for all real-world unstructured data, but that the grokking phenomenon (previously limited to toy datasets) can occur with real-world entities and linguistic structure, provided key conditions are met (notably a sufficiently high $\phi_r$). The structured format allowed us to isolate this effect. In the unstructured setting, generalization circuits are harder to induce due to syntactic variability and ambiguity. We view extending grokking to unstructured data as an important avenue for future work. Our current contribution lies in demonstrating that grokking is not limited to synthetic data, and that it can emerge from real-world facts when structured and augmented appropriately.

D. On relation-specific generalization ratios: The assumption that $\phi_r$ can be modeled as relation-specific follows from prior work [1], which demonstrates that transformer-based generalization circuits rely primarily on relation-specific atomic facts. This behavior has been empirically confirmed in our own replications of their experiments. While entity frequency affects the number of available paths, it is the relation that governs the type and reusability of inference patterns, making $\phi_r$ effectively relation-dependent. We clarify this distinction in Lemma 1 and Appendix A.1, where we show that while the number of unique paths grows with the number of entities $|\mathcal{V}|$, the asymptotic behavior is dominated by the branching factor $b^{n-1}$, not by entity coverage.

E. On model size effects: We are currently conducting experiments with varying GPT-2 sizes (124M–1.5B) to investigate the impact of scale on generalization dynamics. Early results suggest that larger models reach generalization faster, but the critical $\phi_r$ threshold remains stable. For additional details, please see our response to Reviewer hZJa (Section A).

F. On prompt transparency and training details: Thank you for raising this point. We will include full prompt templates, sampling parameters, and training details for both structured and unstructured settings in the final version. Briefly, we trained two separate GPT-2 small models from scratch—one for structured data (triple format) and one for unstructured data (natural language questions). Neither model used pretraining; both were trained solely on task-specific data.

G. On circuit flexibility and multi-hop reasoning: The generalization circuits observed in our setting are relatively rigid and rely on deterministic retrieval and composition of atomic facts. For 2-hop reasoning, this involves a single inference operation; for $n$-hop paths, multiple sequential inference steps are required. Based on our theoretical framework and ongoing experiments, we are confident that $n$-hop grokking is feasible without fundamental changes to the architecture -- provided the augmented dataset supports sufficient multi-hop path coverage.

We again thank the reviewer for the constructive feedback and will incorporate these improvements in the final submission.

We thank the reviewer for the detailed and constructive feedback.

A. On transparency and reproducibility:
We appreciate the emphasis on methodological clarity. In the final version, we will provide full details of the training procedure, including the exact prompts used for structured and unstructured data synthesis, temperature and sampling parameters, and other relevant hyperparameters. We will also include the target responses used for comparison and make code and data available to support reproducibility.

B. On comparisons with large pre-trained models: We agree that comparing against models such as GPT-4o presents challenges due to potential training data leakage, especially with benchmarks like 2WikiMultiHopQA, which overlap with common pretraining corpora such as Wikipedia. We note that this limitation affects most publicly available models and motivates our decision not to report IID/OOD splits for those baselines. Nonetheless, the fact that our grokked GPT-2-small model surpasses these models in structured comparison tasks underscores the potential of $\phi_r$-driven grokking as a viable, lightweight alternative for reasoning over factual knowledge.

C. On model scaling and $\phi_r$ thresholds: We are currently conducting experiments across a range of GPT2 model sizes (124M–1.5B parameters) to evaluate the effect of scaling on generalization dynamics. Preliminary results suggest that while grokking behavior remains consistent, larger models tend to reach generalization more rapidly. Notably, the critical $\phi_r$ threshold required for generalization appears stable across model sizes. Please also refer to our response to Reviewer hZJa (Section A) for more details.

D. On theoretical rigor:
We are in the process of finalizing complete versions of the theoretical results, including clearly stated assumptions and formal proofs to replace the current sketches in Appendix A.2–A.4. These revisions will be included in the final version.

We thank the reviewer for the thoughtful and constructive feedback.

A. On model size and generalization: We appreciate the suggestion regarding scaling effects. We are currently conducting additional experiments across GPT2 model sizes (124M to 1.5B parameters). Preliminary results indicate that the grokking behavior is qualitatively consistent across sizes, especially on structured data. However, larger models tend to reach generalization more quickly. We will include these findings, along with supporting graphs and analysis, in the final version. Notably, the critical ratio $\phi$ remains stable across model sizes.

B. On applicability to other domains: We agree that extending this work to broader domains is an important next step. We are actively exploring applications beyond fact-based reasoning, including logic-intensive tasks such as mathematical problem solving, and real-world scientific domains. These follow-up studies aim to evaluate the generality of $\phi$-driven generalization. We believe this line of inquiry could inform the design of data-efficient training strategies in broader AI applications.

C. On theoretical rigor: Thank you for the feedback on the formal results. We are currently refining the theoretical section to provide complete, rigorous proofs with clearly stated assumptions. The final version will replace the current sketches with detailed derivations, and we will revise the terminology (e.g., replacing “Lemma” if more appropriate) to ensure precision. We will also discuss the limitations and scope of the theoretical claims to better contextualize the results.