Compositional generalization remains a persistent limitation of transformer language models: even when a model learns to combine known facts correctly in-distribution (ID), the same compositional mechanism may fail when applied to out-of-distribution (OOD) knowledge. We investigate this compositionality gap in a controlled two-hop reasoning setting using a GPT-2-like transformer model trained from scratch. Building on prior work identifying a circuit associated with compositional generalization, we test whether the failure of OOD transfer is related to how atomic knowledge is structurally exposed during training. We modify the training distribution by presenting both ID and OOD atomic facts in a format positionally aligned with the second-hop retrieval step of the learned compositional circuit, without modifying the model architecture or adding OOD compositional supervision. Under this intervention, the model successfully transfers the learned ID compositional behaviour to held-out OOD atomic facts, closing the observed compositionality gap in the controlled setting. Our results provide evidence that the observed compositionality gap is not a fundamental architecture limitation and that it can be overcome through mechanistically informed training distribution design.
Figures & tables
Figure 1: OOD Compositional Generalization under the Full Position-Aligned Intervention. ID and OOD test accuracy on the two-hop compositional task where both ID and OOD atomic facts are additionally presented in the second-hop-compatible format during training. OOD accuracy rises with delayed generalization in the grokking training regime, eventually reaching complete compositional generalization and closing the compositionality gap in the controlled experimental setting. The dashed gray curve shows OOD test accuracy under the original training distribution, without position-aligned atomic exposure, which remains near zero.
Figure 2: Generalization Performance Across Ablation Experiments. Comparison of ID and OOD test accuracy on the two-hop compositional task across three controlled ablations. (a) Removing ID atomic facts from the second-hop-compatible format prevents OOD compositional rule transfer. (b) Restricting the second-hop-compatible format to only a subset of OOD atomic facts results in partial OOD compositional generalization. (c) Removing two-hop compositional training examples prevents the model from learning the compositional task for both ID and OOD atomic facts.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Identity-Bridge Control without Compositional Supervision. Accuracy under identity-bridge supervision after removing all two-hop compositional training examples. While training accuracy remains near perfect, both ID and OOD compositional test accuracy remain at zero, showing that identity-bridge supervision alone does not induce compositional generalization in our setting.
Figure 4: Impact of Dataset Size-Density on OOD Compositional Generalization . ( 4(a) ) When trained on the full but less dense dataset, the model demonstrates delayed convergence in both training and testing metrics. As a result, the generalization effect driven by grokking does not achieve full accuracy within the performed training steps. ( 4(b) ) In contrast, training on a smaller but denser dataset leads to significantly improved generalization speed. The model reaches full accuracy on the OOD test set within the given optimization steps, highlighting the benefits of dataset density in accelerating generalization.
Figure 5: Ablation Analysis Results on the Full-Model for the experiments described in section 3.3 . ( 5(a) ): ablation experiment (a) where no in-distribution (ID) atomic facts are provided in the new 2°hop format. Under this condition, the Full Model similarly fails to generalize to OOD examples, resembling the behaviour previously observed in the Smaller Model under the same setting. ( 5(b) ): ablation experiment (b) where out-of-distribution (OOD) facts are partially placed in the second-hop position. In the Full Model experiment, this partial OOD placement appears to completely hinder OOD generalization. This contrasts with the behaviour observed in the Smaller Model, where the same setting resulted in only a partial reduction in final accuracy ( fig. 2(b) ).
Large language models are able to compose skills in order to perform complex tasks, many of which might not have been seen during training. The details of how exactly this composition occurs remain elusive. In this paper, we study a mechanism for compositional generalization in transformers by considering a simple controlled setting involving variable assignment and modular addition. By partitioning our training data into disjoint sets, we observe that small transformers are able to generalize to previously unseen combinations of variables and numbers. Our mechanistic analysis shows that the same ``modular addition'' MLP module is used whether the inputs are given directly or indirectly through a separate variable assignment mechanism. We also analyze the training dynamics from an empirical lens, which reveals three phases of learning: first, modular addition is learned, then the structure required for variable assignment, and finally a refinement phase where the model generalizes to some hard sequences not seen in training. Finally, we provide a theoretical framework to explain how compositionality emerges from training dynamics. These results suggest that compositional generalization can be a natural consequence of the compositionality of internal mechanisms in~transformers.
Large language models (LLMs) can solve complex multi-hop problems yet exhibit puzzling failures on simple two-hop queries: although a model may correctly store each individual hop, it often fails to combine them. To understand the internal mechanisms of this phenomenon, we train transformers from scratch in a controlled symbolic environment. Our experiments reveal a pattern in two-hop generalization: models generalize reliably when the second hop follows the training distribution, but always fail when it deviates. Through mechanistic analysis, we provide a complete explanation for these distinct generalization behaviors: in settings where models generalize successfully, performance is driven by the emergence of consistent intermediate representations for the same entities across contexts, whereas failures on settings where the second hop is out-of-distribution arise from a mismatch across layers: lower layers correctly construct these intermediate representations, but upper layers, while trained on corresponding atomic facts, primarily learn to map them to outputs rather than to reason over them. Driven by this insight, we propose a recurrent-style training strategy, which enables transformers to reuse their reasoning circuitry across input forms and substantially improves generalization on out-of-distribution two-hop queries.
Zili Zhang, Yilin Wang, Heng Wang +2
1Xi’an Jiaotong University · University of Illinois Urbana-Champaign
Large Language Models (LLMs) excel at multi-hop reasoning in distribution, yet fail on unseen compositions, a phenomenon known as the curse of two-hop reasoning. In this work, we argue that this phenomenon can be attributed to a missing supervision on the bridge entity. We formalize this gap by introducing identity bridge, a minimal supervision that enforces a identity mapping on bridge tokens. Under this supervision, even a one-layer transformer with uniform attention (Emb-MLP) can achieve out-of-distribution (OOD) two-hop generalization. We provide a theoretical analysis demonstrating that identity bridge induces an implicit regularization effect, leading the model to establish a direct subject-to-answer association. From an empirical perspective, the performance of standard GPT-2 models aligns closely with simple Emb--MLP models across varying levels of problem complexity. Finally, analyses of fine-tuned mainstream LLMs indicate that correct two-hop predictions consistently coincide with the establishment of a subject-to-answer relationship, extending our findings to realistic settings.
Pengxiao Lin, Zheng-An Chen, Zhi-Qin John Xu
School of Mathematical Sciences, Shanghai Jiao Tong University · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University · Shanghai Seres Information Technology Co., Ltd, Shanghai 200040, China.