Why Knowing Both Hops Is Not Enough: Understanding Two-Hop Generalization in Language Models

# Why Knowing Both Hops Is Not Enough: Understanding Two-Hop Generalization in Language Models
Source: [https://arxiv.org/html/2608.07261](https://arxiv.org/html/2608.07261)
Zili Zhang1Yilin Wang111footnotemark:1Heng Wang2Herun Wan1Minnan Luo1 1Xi’an Jiaotong University2University of Illinois Urbana\-Champaign \{zhangzl666, 13148035071xjtu\}@stu\.xjtu\.edu\.cn

###### Abstract

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\.

Why Knowing Both Hops Is Not Enough: Understanding Two\-Hop Generalization in Language Models

Zili Zhang1††thanks:Equal contribution\.Yilin Wang111footnotemark:1Heng Wang2Herun Wan1Minnan Luo1††thanks:Corresponding author\.1Xi’an Jiaotong University2University of Illinois Urbana\-Champaign\{zhangzl666, 13148035071xjtu\}@stu\.xjtu\.edu\.cn

## 1Introduction

Large language models \(LLMs\) have achieved impressive performance across multiple tasks\(Zhaoet al\.,[2025a](https://arxiv.org/html/2608.07261#bib.bib39)\), yet their ability to perform multi\-hop reasoning remains fragile and poorly understood\(Huang and Chang,[2023](https://arxiv.org/html/2608.07261#bib.bib41)\)\. Even in the simplest case—two\-hop reasoning—models exhibit a persistent failure: although they can correctly store each atomic fact \(e\.g\., A→\\toB and B→\\toC\), they often struggle to reliably combine them into a compositional inference \(A→\\toC\)\(Biranet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib9); Balesniet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib4)\)\.

Therefore, enhancing the ability of multi\-hop reasoning has recently become a research focus\(Liet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib6); OpenAI,[2024](https://arxiv.org/html/2608.07261#bib.bib5); Pettyet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib8); Luoet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib7)\)\. Beyond explicit prompting strategies such as Chain\-of\-Thought \(CoT\)\(Kojimaet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib12); Weiet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib13)\), a growing line of work suggests that LLMs can answer compositional queries without generating intermediate steps, indicating the presence of implicit multi\-hop reasoning\(Denget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib14); Yanget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib15)\)\. However, the underlying mechanism governing how transformers internally complete two\-hop queries using latent reasoning remains not well understood\.

In this study, we aim to dissect the internal mechanisms of this phenomenon without the confounding factors of pretraining data opacity—such as simple memorization or spurious correlations\(Xuet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib3); Wanget al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib1); Juet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib2)\)\. To this end, in §[2](https://arxiv.org/html/2608.07261#S2), we use a symbolic environment to train a transformer from scratch and test our trained model on different splits of data\.

Using this approach, in §[3](https://arxiv.org/html/2608.07261#S3), we observe an asymmetric generalization pattern: models generalize reliably when the second hop follows the same distributional patterns observed during training \(Test\-II, Test\-OI\), but fail when the second hop is out\-of\-distribution \(Test\-IO, Test\-OO\)\.

To give a complete explanation, we conduct a mechanistic analysis in §[4](https://arxiv.org/html/2608.07261#S4)\. First, we observe a progressive reasoning process across layers using logit lens\. Based on this experiment, we select layer 5 as the anchor for further analysis and define its output at positionr1\{r\_\{1\}\}as the representation of the bridge entity\. Then, through causal patching, we find that the alignment of the same entity’s representations is positively correlated with the model’s generalization on Tset\-II/OI\. To explain the failure on Test\-OO and Test\-IO, we hypothesize that the upper layers of the model only learn the mapping ability during atomic facts training, and validate it by linear transformation and attention knockout experiments\.

Dirven by the above insights, in §[5](https://arxiv.org/html/2608.07261#S5), we heuristically propose that the looped architecture is a way to solve multi\-hop reasoning failure: the model eliminates the difference between training and inference by aligning representations across different layers\. Extending experiments verify the scalability of our method and its effectiveness on real\-world corpora\.

In summary, this paper makes the following contributions:

- •Identify a systematic asymmetry in two\-hop generalization: models generalize reliably when the second hop stays within the training distribution, and, conversely, fail when the second hop is out\-of\-distribution, despite correct atomic knowledge\.
- •Provide a complete mechanistic explanation for the two\-hop phenomenon: we attribute the generalization success on Test\-II/OI to the same entity’s representations alignment, and its failure on Test\-IO/OO to a functional mismatch where upper layers specialize in representation mapping rather than reasoning\.
- •Propose a mechanism\-grounded solution: we adopt a looped architecture and show that aligning representation formats across layers enables robust two\-hop generalization\. Extending experiments show that our method is also applicable to modern architectures and real\-world language data\.

## 2Experimental Settings

We construct the data in a symbolic environment for controlled analysis, followingWanget al\.\([2024](https://arxiv.org/html/2608.07261#bib.bib10)\)\.

Data Construction:The overall construction pipeline is shown in Figure[1](https://arxiv.org/html/2608.07261#S2.F1): We define an entity setℰ=\{e1,…,e\|ℰ\|\}\\mathcal\{E\}=\\\{e\_\{1\},\\ldots,e\_\{\|\\mathcal\{E\}\|\}\\\}and a relation setℛ=\{r1,…,r\|ℛ\|\}\\mathcal\{R\}=\\\{r\_\{1\},\\ldots,r\_\{\|\\mathcal\{R\}\|\}\\\}\. For each head entitye1∈ℰe\_\{1\}\\in\\mathcal\{E\}, we sample 20 relations fromℛ\\mathcal\{R\}and pair each with a randomly chosen tail entitye2∈ℰe\_\{2\}\\in\\mathcal\{E\}, yielding atomic facts\(e1,r,e2\)\(e\_\{1\},r,e\_\{2\}\)\. We partition atomic facts into in\-distribution \(ID\) and out\-of\-distribution \(OOD\) subsets with proportionϕ\\phi\. We then form two\-hop facts\(e1,r1,r2,e3\)\(e\_\{1\},r\_\{1\},r\_\{2\},e\_\{3\}\)by chaining\(e1,r1,e2\)\(e\_\{1\},r\_\{1\},e\_\{2\}\)and\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\)that share bridge entitye2e\_\{2\}\. These facts fall into four types: II, OO, IO, and OI\. II is further split into Train\-II and Test\-II, while OO/IO/OI are used as Test\-OO/Test\-IO/Test\-OI\. The cross\-distribution splits followYeet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib11)\)\. Details are in Appendix[A](https://arxiv.org/html/2608.07261#A1)\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x1.png)Figure 1:The pipeline of data construction\.Training Configuration:Our training split consists of all atomic facts \(ID and OOD\) and Train\-II\. All other two\-hop splits \(Test\-II, Test\-IO, Test\-OI, and Test\-OO\) are used for evaluation to comprehensively assess the generalization ability of the model\. We train a GPT2\-style decoder\-only transformer\([Radfordet al\.,](https://arxiv.org/html/2608.07261#bib.bib40)\)with 8 layers and a hidden dimension of 768\. The model is optimized using AdamW with a learning rate of1×10−41\\times 10^\{\-4\}\. More detailed training configuration is provided in appendix[B](https://arxiv.org/html/2608.07261#A2)\.

## 3Training Results

![Refer to caption](https://arxiv.org/html/2608.07261v1/x2.png)Figure 2:Training dynamics of a transformer trained on all atomic facts and in\-distribution two\-hop compositions\. Accuracy is shown for atomic facts, the two\-hop training set \(Train\-II\), and multiple evaluation splits\. The model rapidly fits the training data, followed by delayed generalization to the in\-distribution test set \(Test\-II\) and partial improvement on Test\-OI\. In contrast, performance on Test\-IO and Test\-OO remains near chance throughout training, indicating a failure to generalize when the second hop is out\-of\-distribution\.Figure[2](https://arxiv.org/html/2608.07261#S3.F2)shows the training dynamics across different data splits\. We observe that:\(A\)The model first achieves near\-perfect accuracy on Train\-II after a relatively small number of optimization steps\.\(B\)After continued training, the accuracy on Test\-II begins to increase steadily and eventually reaches perfect accuracy, suggesting that the model transitions from memorization to generalization\.\(C\)At a later stage of training, the model also exhibits noticeable gains on the Test\-OI split, indicating the emergence of an additional generalization phase\.\(D\)However, even with longer training, the model cannot generalize on Test\-IO and Test\-OO splits\.

Overall, the model exhibits a highly asymmetric generalization pattern: it generalizes reliably, only when the second hop stays within the training distribution \(Test\-II and Test\-OI\), and consistently fails when the second hop is out\-of\-distribution \(Test\-IO and Test\-OO\)\. In the following section, we investigate the underlying reasons for this disparity by analyzing how intermediate representations are formed and processed during two\-hop reasoning\.

## 4Mechanism of Two\-Hop Reasoning

### 4\.1Identifying the Bridge Entity Representation in Hidden Space

To explain the asymmetric training dynamics observed above, we begin by identifying where and how the bridge entitye2e\_\{2\}is represented in the model\. To this end, we conduct a logit lens\(nostalgebraist,[2020](https://arxiv.org/html/2608.07261#bib.bib23)\)analysis on the trained model at both ther1r\_\{1\}andr2r\_\{2\}positions\. At ther1r\_\{1\}position, we track how the probability ofe2e\_\{2\}evolves across layers, probing when the model acquires sufficient information about the intermediate entity\. At ther2r\_\{2\}position, we examine the probabilities of the relation tokenr2r\_\{2\}itself and the target entitye3e\_\{3\}, which together reflect the beginning of second\-hop reasoning\.

As shown in Figure[3](https://arxiv.org/html/2608.07261#S4.F3), layer 5 emerges as a critical transition layer\. By this layer, the model has already accumulated strong evidence for the bridge entitye2e\_\{2\}, as indicated by the high probability ofe2e\_\{2\}atr1r\_\{1\}position\. More importantly, after layer 5, we observe a qualitative shift in behavior at ther2r\_\{2\}position: the probability of the surface tokenr2r\_\{2\}begins to decrease, while the probability of the final answere3e\_\{3\}rises, suggesting upper layers may start integrating contextual information and performing the second hop of reasoning\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x3.png)Figure 3:Logit lens reveals a progressive reasoning process, and that layer 5 is a transition layer between two hops\.This progressive reasoning process leads us to conjecture that the output at positionr1r\_\{1\}of layer 5 may contain information about the bridge entitye2e\_\{2\}that is necessary for the model’s reasoning\. Therefore, we define it as the hidden representation of the bridge entity, denoted byℛ\\mathscr\{R\}, which serves as an anchor for subsequent analysis111In Appendix[D](https://arxiv.org/html/2608.07261#A4), we repeat this experiment at other training stages and find that the key concept layer only changes in the early memorization stage and is stable thereafter, so using this layer as the analysis anchor is reasonable\.\.

### 4\.2Probing Bridge Entity Representation Consistency via Entity Patching

![Refer to caption](https://arxiv.org/html/2608.07261v1/x4.png)Figure 4:\(Left\)A two\-hop query composed of two single\-hop facts, where the bridge entitye2e\_\{2\}connects\(e4,r4\)\(e\_\{4\},r\_\{4\}\)and\(e2,r3\)\(e\_\{2\},r\_\{3\}\)to yield the target entitye3e\_\{3\}\.\(Right\)The entity patching procedure: the hidden representation of the bridge entitye2e\_\{2\}at layer 5 is replaced with an alternative representation of the same entity extracted from a different single\-hop context \(e\.g\.,\(e1,r1\)\(e\_\{1\},r\_\{1\}\)\)\. The model’s ability to recover the original answer after patching tests whether representations of the same entity are aligned across contexts\.To perform robust multi\-hop reasoning, models must move beyond surface\-level pattern matching and internalize the underlying graph structure of knowledge\. A fundamental requirement for such capability is concept invariance: the internal representation of a bridge entitye2e\_\{2\}should remain consistent regardless of the specific single\-hop context that elicits it\. Based on this intuition, we hypothesize that the model’s ability to generalize on Test\-II and Test\-OI is driven by the emergence of such consistent representations at the critical bridge layer\.

To test this hypothesis, we introduce anentity representation consistency patching experiment, illustrated in Figure[4](https://arxiv.org/html/2608.07261#S4.F4)\. We consider a two\-hop query from the in\-distribution training set \(Train\-II\), where the bridge entitye2e\_\{2\}appears at positionr2r\_\{2\}and is represented by the hidden state at layer 5, following the definition established in Section[4\.1](https://arxiv.org/html/2608.07261#S4.SS1)\. In this instance, the single\-hop context that gives rise toe2e\_\{2\}is\(e4,r4\)\(e\_\{4\},r\_\{4\}\)\.

We then select another single\-hop atomic fact with the same tail entitye2e\_\{2\}but a different context, e\.g\.,\(e1,r1\)\(e\_\{1\},r\_\{1\}\)\. The representationℛe2\\mathscr\{R\}\_\{e\_\{2\}\}extracted from this context is used to replace the original bridge entity representation in the two\-hop query, while keeping all other components unchanged222To rule oute3e\_\{3\}being pre\-encoded atr3r\_\{3\}in early training, we replacer3r\_\{3\}with that from another two\-hop query whose answer is note3e\_\{3\}, so predictinge3e\_\{3\}must use the patched representation rather than residual earlier\-layer information\.\. After patching, we evaluate whether the model still produces the correct final answere3e\_\{3\}\.

We perform this replacement using bridge entity representations drawn from two sources: \(i\) in\-distribution \(ID\) atomic facts and \(ii\) out\-of\-distribution \(OOD\) atomic facts, and record the patching success rate throughout training\. Intuitively, if the model has learned a context\-invariant representation fore2e\_\{2\}, substituting its representation from another context should preserve the correctness of the second\-hop inference\. As a control, we also try to inject Gaussian noise, destroying the original representation to verify that this representation is causally necessary for the final prediction; the resulting performance drop \(The grey line in Figure[5](https://arxiv.org/html/2608.07261#S4.F5)\) ensures that patching success stems from representation alignment rather than the model ignoring the position\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x5.png)Figure 5:Patching success rate using ID and OOD bridge entity representations during training, compared with Test\-II and Test\-OI accuracy\.We show the success rate of the patching experiment and the model’s performance on Test\-II and Test\-OI in Figure[5](https://arxiv.org/html/2608.07261#S4.F5)as they change during the training process, which reveals a clear training\-dependent pattern\. In the early stages of training, patching success rates are low for both ID and OOD replacements, indicating that representations of the same entity remain highly context\-dependent\. During this phase, the model operates in a rote memorization regime, encoding entity information in a fragmented and inefficient manner\.

As training progresses, the success rate of ID\-based patching increases first\. Notably, this rise coincides with the improvement of performance on Test\-II, suggesting generalization within the in\-distribution setting emerges once representations of the same entity become aligned across different ID contexts\. With further training, the success rate of OOD\-based patching also begins to increase, accompanied by a corresponding improvement on Test\-OI\.

These results support the assumption that generalization in two\-hop reasoning is driven by the gradual alignment of entity representations across contexts\. Moreover, representations derived from in\-distribution atomic facts converge more rapidly than those from out\-of\-distribution facts, explaining why generalization to Test\-II precedes that to Test\-OI\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x6.png)Figure 6:High cosine similarity and low MSE from layer 5 onward in single\-hop queries indicate that upper layers act as linear mappers, whereas the lower linearity in two\-hop queries suggests nonlinear information aggregation at the bridge layer\.
### 4\.3Why Second\-Hop OOD Fails: Reasoning vs\. Mapping in Upper Layers

![Refer to caption](https://arxiv.org/html/2608.07261v1/x7.png)Figure 7:\(Left\)In a two\-hop query\(e1,r1,r2\)→e3\(e\_\{1\},r\_\{1\},r\_\{2\}\)\\\!\\rightarrow\\\!e\_\{3\}, lower layers first form the bridge entity representatione2e\_\{2\}atr1r\_\{1\}\. Upper layers must then combineℛe2\\mathscr\{R\}\_\{e\_\{2\}\}withr2r\_\{2\}to infere3e\_\{3\}\.\(Right\)For a single\-hop fact\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\), the representation ofe3e\_\{3\}is already formed atr2r\_\{2\}in lower layers, so upper layers only map it to the outputhe3Lh^\{L\}\_\{e\_\{3\}\}without relational reasoning\.We next investigate why generalization emerges on Test\-II and Test\-OI, but consistently fails on Test\-IO and Test\-OO\. An observation is that the essential difference between these test sets lies in whether the second hop corresponds to in\-distribution or out\-of\-distribution knowledge\. Given that the first\-hop representation has already been shown to be well\-formed and robust \(Section[4\.2](https://arxiv.org/html/2608.07261#S4.SS2)\), the remaining bottleneck must reside in the second\-hop computation, namely, how the model operates on intermediate representations\.

In a two\-hop query of the form\(e1,r1,r2\)→e3\(e\_\{1\},r\_\{1\},r\_\{2\}\)\\rightarrow e\_\{3\}, during the second\-hop inference, the model is required to use the representation of the bridge entitye2e\_\{2\}formed at positionr1r\_\{1\}, and then reason through relationr2r\_\{2\}to infer the target entitye3e\_\{3\}\. As shown in the left side of Figure[7](https://arxiv.org/html/2608.07261#S4.F7), this requires the upper layers to implement a computation that implementsℛe2→r2e3\\mathscr\{R\}\_\{e\_\{2\}\}\\xrightarrow\{r\_\{2\}\}\{e\_\{3\}\}, where the input is an intermediate entity representation rather than a surface token embedding\.

In contrast, during single\-hop training on atomic facts\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\), the situation is fundamentally different\. At positionr2r\_\{2\}, the upper layers already have direct access to the representation ofe3e\_\{3\}\. Consequently, the upper layersare not requiredto perform relational reasoning; instead, they can learn a much simpler transformation:ℛe3⟶e3\\mathscr\{R\}\_\{e\_\{3\}\}\\longrightarrow\{e\_\{3\}\}, namely a mapping from an already\-formed representation to the final output space, as shown in the right side of Figure[7](https://arxiv.org/html/2608.07261#S4.F7)\.

This distinction leads to atraining\-inference mismatchhypothesis: single\-hop training primarily teaches the upper layers to perform representation mapping rather than representation\-based reasoning\. Under this view, the observed generalization behavior follows naturally\. For Test\-II and Test\-OI, all their second hops in the form of intermediate representation have appeared during training \(i\.e\., within Train\-II\)\. Consequently, the model learns to consume theℛe2\\mathscr\{R\}\_\{e\_\{2\}\}and reason throughr2r\_\{2\}, enabling it to transfer this representation\-based reasoning capability to these test subsets\. In contrast, for Test\-IO and Test\-OO, the second hop has never been trained in a representational form—it has only been encountered as a single\-hop atomic fact\. The upper layers, therefore, attempt to apply a learned mapping where reasoning is required, leading to systematic failure\.

To validate this hypothesis, we design a linear probing experiment333Besides, we also conduct an attention blocking experiment\(Gevaet al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib38)\)in Appendix[E](https://arxiv.org/html/2608.07261#A5), which yields the same conclusion\.\. Specifically, for atomic training examples, we collect: the hidden state at the⟨r⟩\\langle r\\rangleposition from an intermediate layerℓ\\ellas input, denoted𝐡rℓ\\mathbf\{h\}^\{\\ell\}\_\{r\}; and the hidden state at the same position from the final layerLLas output, denoted𝐡rL\\mathbf\{h\}^\{L\}\_\{r\}\. We then fit a linear model𝐡rL≈𝐖𝐡rℓ\\mathbf\{h\}^\{L\}\_\{r\}\\approx\\mathbf\{W\}\\mathbf\{h\}^\{\\ell\}\_\{r\}on a subset of these\(𝐡ℓ,𝐡L\)\(\\mathbf\{h\}^\{\\ell\},\\mathbf\{h\}^\{L\}\)pairs, and evaluate its performance on held\-out data using both mean squared error \(MSE\) and cosine similarity as comprehensive evaluation metrics, similar to the idea inKhandelwal and Pavlick \([2025](https://arxiv.org/html/2608.07261#bib.bib22)\)\. We compare the linearity of this mapping in two scenarios: therrposition for single\-hop queries and ther2r\_\{2\}position for two\-hop queries\.

Figure[6](https://arxiv.org/html/2608.07261#S4.F6)reveals a clear divergence in linearity, particularly in layer 5, the bridge representation layer\.For single\-hop queries, the transformation from layer 5 to the final output exhibits near\-perfect linearity \(high cosine similarity, low MSE\)\. This confirms our hypothesis that for atomic facts, once the entity is resolved, the upper layers function primarily as linear mappers\. In contrast,for two\-hop queries, the linearity at layer 5 is significantly lower\. This disparity indicates that at this critical stage, a simple linear mapping is insufficient; instead, the model requires non\-linear operations \(e\.g\., attention mechanisms\) to aggregate the bridge entity information fromr1r\_\{1\}intor2r\_\{2\}\. The gradual improvement in linearity across subsequent layers reflects the progressive completion of this information aggregation\.

Consequently, the failure of OOD generalization can be attributed to a functional mismatch: the upper layers are trained on atomic facts to perform linear mapping, but multi\-hop reasoning requires them to execute non\-linear information aggregation at the bridge layer\.

## 5Bridging the Representation\-Reasoning Gap for Robust Two\-Hop Generalization

![Refer to caption](https://arxiv.org/html/2608.07261v1/x8.png)\(a\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x9.png)\(b\)

Figure 8:\(a\)Continual training of upper layers on atomic facts presented in a representational form\. We only depict the process of continual training; the standard training process is the same as in Figure[2](https://arxiv.org/html/2608.07261#S3.F2)\.\(b\)Looped training with shared parameters yields consistent generalization across all test splits by aligning intermediate representations with lower\-layer inputs\.Our analysis so far indicates that the failure of two\-hop generalization does not stem from missing knowledge, but is most likely caused by a structural mismatch between the roles played by different layers of the model\. This diagnosis suggests a natural direction for achieving robust generalization: bridging the gap between lower\-layer representations and upper\-layer reasoning\. In the following, we explore two strategies to achieve this goal\.

### 5\.1Explicitly training upper layers for representation\-based reasoning

A direct approach is to equip the upper layers with the ability to reason over representations\. Concretely, after standard training, we introduce a continuous training process applied only to the upper layers \(layer 5 and above\)\. In this auxiliary phase, the model is trained on OOD atomic facts, but presented purely in a representational form: the inputs to the upper layers are pairs of hidden states\(he,hr\)\(h\_\{e\},h\_\{r\}\), rather than surface tokens \(It is necessary to note that this continual training only exposes the model to atomic facts and does not leak any two\-hop supervision\)\. We provide more details about the continual training process in Appendix[F](https://arxiv.org/html/2608.07261#A6)

The evaluation results on Test\-OO and Test\-IO as training processes are shown in Figure[8\(a\)](https://arxiv.org/html/2608.07261#S5.F8.sf1)\. We observe that this additional process substantially improves generalization on previously failing settings, namely Test\-IO and Test\-OO\. This improvement confirms our earlier diagnosis that the upper layer struggles to reason over out\-of\-distribution atomic facts in representational form\. However, this approach comes with notable drawbacks: it introduces additional training stages and increases computational cost\.

### 5\.2Eliminating layer\-wise mismatch via looped architecture

An alternative and more principled solution is to remove the structural mismatch altogether\. Instead of explicitly training the upper layers on representational inputs, we enforce alignment between representations and lower\-layer inputs by sharing parameters across layers \(Here, we share the parameters of the top four layers and the bottom four layers\), forming alooped architecture\(Yanget al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib24); Fanet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib25)\)\. In this design, the hidden states produced by the model’s first forward pass are fed back as input to a second forward pass using the same set of parameters\. Because both passes share weights, the model is forced to process surface tokens and intermediate representations in a unified format\.

As a result, even when the model has never encountered OOD atomic knowledge in a representational form, the alignment between representations and lower\-layer inputs allows it to reuse the same reasoning circuit, thus bridging the training\-inference gap\. Figure[8\(b\)](https://arxiv.org/html/2608.07261#S5.F8.sf2)demonstrates that this looped training strategy yields substantial improvement in generalization performance, without requiring additional supervision on OOD representations\.

### 5\.3Evidence of Representation\-Input Alignment under Looped Training

The intuition behind looped training is that robust two\-hop generalization emerges from an explicit alignment between intermediate representations and lower\-layer inputs\. In this subsection, we provide direct empirical evidence for this claim by measuring representational similarity across layers\.

To quantify the alignment between the model’s internal representations and its input forms, we calculate the cosine similarity for both the bridge entitye2e\_\{2\}and the relationr2r\_\{2\}at critical transition layers\. We define two alignment metrics:

1. 1\.Loop\-End Alignment \(SendS\_\{\\text\{end\}\}\): Measures how close the representation at the end of the loop block \(layer 4\) is to the input embeddings\.Send​\(e2\)=cos⁡\(𝐡r14,𝐄e2\),Send​\(r2\)=cos⁡\(𝐡r24,𝐄r2\)S\_\{\\text\{end\}\}\(e\_\{2\}\)=\\cos\(\\mathbf\{h\}^\{4\}\_\{r\_\{1\}\},\\mathbf\{E\}\_\{e\_\{2\}\}\),\\\>S\_\{\\text\{end\}\}\(r\_\{2\}\)=\\cos\(\\mathbf\{h\}^\{4\}\_\{r\_\{2\}\},\\mathbf\{E\}\_\{r\_\{2\}\}\)
2. 2\.Bridge\-Layer Alignment \(SbridgeS\_\{\\text\{bridge\}\}\): Measures the similarity between the bridge layer \(layer 5\) and layer 1\.Sbridge​\(e2\)=cos⁡\(𝐡r15,𝐡e21\),Sbridge​\(r2\)=cos⁡\(𝐡r25,𝐡r21\)S\_\{\\text\{bridge\}\}\(e\_\{2\}\)=\\cos\(\\mathbf\{h\}^\{5\}\_\{r\_\{1\}\},\\mathbf\{h\}^\{1\}\_\{e\_\{2\}\}\),\\\>S\_\{\\text\{bridge\}\}\(r\_\{2\}\)=\\cos\(\\mathbf\{h\}^\{5\}\_\{r\_\{2\}\},\\mathbf\{h\}^\{1\}\_\{r\_\{2\}\}\)

Table 1:Cosine similarity analysis of representation alignment\. Looped training forces high compatibility between intermediate representations and lower\-layer inputs\.The results in Table[1](https://arxiv.org/html/2608.07261#S5.T1)reveal a striking contrast between the two training regimes\. Under looped training, intermediate representations exhibit strong alignment with lower\-layer inputs\. In particular, the representation of the bridge entitye2e\_\{2\}shows consistently high cosine similarity across layers, and the representation of relationr2r\_\{2\}becomes highly aligned when comparing higher\-layer outputs to lower\-layer inputs\. In contrast, under standard training, this alignment is weak or absent\.

Driven by the need for parameter sharing \(e\.g\., layer 5 and layer 1 have the same parameters\), the model is compelled to converge their hidden states into a unified format\. This alignment ensures the reasoning machinery learned on surface forms can be seamlessly re\-applied to intermediate states, explaining the generalization observed in Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3)\.

### 5\.4Scalability of Looped Architecture

We finally examine whether the benefits of looped training extend to modern architectures and real\-world two\-hop data\. We benchmark performance on two model backbones: the GPT\-style model used in previous sections and a Llama\-based architecture\. We also move beyond symbolic tasks to a real dataset derived from Wikipedia, consisting of natural\-language multi\-hop reasoning samples444The construction details are shown in Appendix[K](https://arxiv.org/html/2608.07261#A11)\.\. Results are reported in Table[2](https://arxiv.org/html/2608.07261#S5.T2)\. Loops prove robust on modern architecture\. On the real\-world dataset, the model successfully generalizes to OOD splits, confirming that the convergence of intermediate states is not an artifact of symbolic data but a viable strategy for realistic reasoning tasks\. We also study how changing the number of trainable layers under the looped architecture affects generalization; results in Appendix[I](https://arxiv.org/html/2608.07261#A9)show that allowing more trainable layers substantially improves the model’s generalization ability and breaks the bottleneck in Figure[8](https://arxiv.org/html/2608.07261#S5.F8)\.

Table 2:Generalization performance across different architectures and datasets\.Realdenotes natural language 2\-hop data extracted from Wikipedia\.

## 6Related Works

### 6\.1Mechanistic Interpretability

Mechanistic interpretability seeks to reverse engineer how LLMs implement specific computations\(Olah,[2022](https://arxiv.org/html/2608.07261#bib.bib26)\)\. A line of work studies hidden states across layers, including logit lens methods that project intermediate states through the unembedding matrix to examine token\-level predictions\(nostalgebraist,[2020](https://arxiv.org/html/2608.07261#bib.bib23); Daret al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib27); Yu and Ananiadou,[2024](https://arxiv.org/html/2608.07261#bib.bib28); Katz and Belinkov,[2023](https://arxiv.org/html/2608.07261#bib.bib29)\)\. Other approaches use causal interventions on hidden states to measure their effect on outputs and localize functionally relevant components\(Stolfoet al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib32); Menget al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib31); Viget al\.,[2020](https://arxiv.org/html/2608.07261#bib.bib30)\)\. Motivated by the superposition hypothesis, sparse autoencoders extract more interpretable features from activations\(Scherliset al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib34); Elhageet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib33); Balcellset al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib36); Zhaoet al\.,[2025b](https://arxiv.org/html/2608.07261#bib.bib35)\)\. Together, these studies provide a toolkit for dissecting how complex behaviors emerge from layered computations\.

### 6\.2Implicit Reasoning

LLMs can perform multi\-step or compositional reasoning without explicitly generating intermediate chains, a phenomenon termed implicit reasoning\(Kojimaet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib12); Weiet al\.,[2022](https://arxiv.org/html/2608.07261#bib.bib13)\)\. Prior work shows that transformers encode latent reasoning states internally\(Denget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib14)\), and pretrained models can answer multi\-hop queries without exposing intermediate steps\(Yanget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib15); Balesniet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib4)\)\. Mechanistic analyses further link the emergence of implicit reasoning to representation structure and grokking\-like dynamics\(Wanget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib10); Yeet al\.,[2025](https://arxiv.org/html/2608.07261#bib.bib11)\)\. Based on this line of work, we attribute second\-hop OOD failure to a layer\-wise mismatch between representation and reasoning\.

### 6\.3Two\-Hop Reasoning

Two\-hop reasoning is a canonical multi\-hop QA setting that requires retrieving and composing multiple atomic facts into a single inference\(Yanget al\.,[2018](https://arxiv.org/html/2608.07261#bib.bib16)\)\. Prior work identifies a persistent gap between factual recall and compositional reasoning in LLMs:Presset al\.\([2023](https://arxiv.org/html/2608.07261#bib.bib17)\)formalizes this as the compositionality gap, where models fail on two\-hop queries despite succeeding on the corresponding single\-hop facts\. Subsequent studies show that such failures are unstable across domains and settings, questioning whether LLMs robustly perform genuine multi\-hop reasoning\(Berglundet al\.,[2023](https://arxiv.org/html/2608.07261#bib.bib18); Yanget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib15)\)\. Mechanistic analyses further reveal that even when the bridge entity is correctly retrieved at the first hop, it is often not utilized by higher layers, indicating a disconnect between representation formation and reasoning\(Biranet al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib9); Yuet al\.,[2025](https://arxiv.org/html/2608.07261#bib.bib19)\)\. Related work also studies two\-hop reasoning circuits from theoretical, controlled finetuning, and in\-context learning perspectives\(Fenget al\.,[2024](https://arxiv.org/html/2608.07261#bib.bib20); Guoet al\.,[2025](https://arxiv.org/html/2608.07261#bib.bib21)\)\.

## 7Conclusion

This work provides a mechanistic explanation of the two\-hop problem in Transformers\. In a controlled symbolic setting, we identify an asymmetric pattern: models generalize when the second hop is in\-distribution but fail when it is out\-of\-distribution\. We trace this behavior to representation dynamics: successful cases exhibit consistent bridge\-entity representations across contexts, whereas failures arise because upper layers mainly map resolved representations to outputs rather than reason over intermediate states\. Based on these findings, we propose a looped architecture that aligns intermediate representations with lower\-layer inputs, enabling reuse of the same reasoning circuitry across input forms\. Experiments on symbolic and natural\-language tasks improve OOD generalization\.

## Limitations

While our looped architecture effectively enables representation reuse and substantially improves generalization, our training procedure remains relatively simple: we adopt a standard pre\-training strategy and apply only a single loop\. It is possible that alternative training schedules, multiple looping iterations, or more adaptive mechanisms could further enhance efficiency or performance\(Zhuet al\.,[2025](https://arxiv.org/html/2608.07261#bib.bib37)\)\. Exploring these directions may lead to more effective architectures, but it falls outside the scope of the present study\.

## References

- R\. Abramov, F\. Steinbauer, and G\. Kasneci \(2025\)Grokking in the wild: data augmentation for real\-world multi\-hop reasoning with transformers\.InProceedings of the 42nd International Conference on Machine Learning,A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\),Proceedings of Machine Learning Research, Vol\.267,pp\. 231–245\.External Links:[Link](https://proceedings.mlr.press/v267/abramov25a.html)Cited by:[Appendix K](https://arxiv.org/html/2608.07261#A11.p2.5)\.
- D\. Balcells, B\. Lerner, M\. Oesterle, E\. Ucar, and S\. Heimersheim \(2024\)Evolution of sae features across layers in llms\.arXiv preprint arXiv:2410\.08869\.Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- M\. Balesni, T\. Korbak, and O\. Evans \(2024\)Lessons from studying two\-hop latent reasoning\.External Links:[Link](https://api.semanticscholar.org/CorpusID:274234948)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p1.3),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- L\. Berglund, A\. C\. Stickland, M\. Balesni, M\. Kaufmann, M\. Tong, T\. Korbak, D\. Kokotajlo, and O\. Evans \(2023\)Taken out of context: on measuring situational awareness in llms\.arXiv preprint arXiv:2309\.00667\.Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- E\. Biran, D\. Gottesman, S\. Yang, M\. Geva, and A\. Globerson \(2024\)Hopping too late: exploring the limitations of large language models on multi\-hop queries\.InConference on Empirical Methods in Natural Language Processing,External Links:[Link](https://api.semanticscholar.org/CorpusID:270562421)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p1.3),[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- G\. Dar, M\. Geva, A\. Gupta, and J\. Berant \(2023\)Analyzing transformers in embedding space\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),A\. Rogers, J\. Boyd\-Graber, and N\. Okazaki \(Eds\.\),Toronto, Canada,pp\. 16124–16170\.External Links:[Link](https://aclanthology.org/2023.acl-long.893/),[Document](https://dx.doi.org/10.18653/v1/2023.acl-long.893)Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- Y\. Deng, Y\. Choi, and S\. Shieber \(2024\)From explicit cot to implicit cot: learning to internalize cot step by step\.arXiv preprint arXiv:2405\.14838\.Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- N\. Elhage, T\. Hume, C\. Olsson, N\. Schiefer, T\. Henighan, S\. Kravec, Z\. Hatfield\-Dodds, R\. Lasenby, D\. Drain, C\. Chen,et al\.\(2022\)Toy models of superposition\.arXiv preprint arXiv:2209\.10652\.Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- Y\. Fan, Y\. Du, K\. Ramchandran, and K\. Lee \(2024\)Looped transformers for length generalization\.arXiv preprint arXiv:2409\.15647\.Cited by:[§5\.2](https://arxiv.org/html/2608.07261#S5.SS2.p1.1)\.
- J\. Feng, S\. Russell, and J\. Steinhardt \(2024\)Extractive structures learned in pretraining enable generalization on finetuned facts\.ArXivabs/2412\.04614\.External Links:[Link](https://api.semanticscholar.org/CorpusID:274581557)Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- M\. Geva, J\. Bastings, K\. Filippova, and A\. Globerson \(2023\)Dissecting recall of factual associations in auto\-regressive language models\.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing,H\. Bouamor, J\. Pino, and K\. Bali \(Eds\.\),Singapore,pp\. 12216–12235\.External Links:[Link](https://aclanthology.org/2023.emnlp-main.751/),[Document](https://dx.doi.org/10.18653/v1/2023.emnlp-main.751)Cited by:[footnote 3](https://arxiv.org/html/2608.07261#footnote3)\.
- T\. Guo, H\. Zhu, R\. Zhang, J\. Jiao, S\. Mei, M\. I\. Jordan, and S\. Russell \(2025\)How do llms perform two\-hop reasoning in context?\.arXiv preprint arXiv:2502\.13913\.Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- J\. Huang and K\. C\. Chang \(2023\)Towards Reasoning in Large Language Models: A Survey\.arXiv\.External Links:2212\.10403,[Document](https://dx.doi.org/10.48550/arXiv.2212.10403)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p1.3)\.
- T\. Ju, Y\. Chen, X\. Yuan, Z\. Zhang, W\. Du, Y\. Zheng, and G\. Liu \(2024\)Investigating multi\-hop factual shortcuts in knowledge editing of large language models\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),L\. Ku, A\. Martins, and V\. Srikumar \(Eds\.\),Bangkok, Thailand,pp\. 8987–9001\.External Links:[Link](https://aclanthology.org/2024.acl-long.486/),[Document](https://dx.doi.org/10.18653/v1/2024.acl-long.486)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p3.1)\.
- S\. Katz and Y\. Belinkov \(2023\)VISIT: visualizing and interpreting the semantic information flow of transformers\.arXiv preprint arXiv:2305\.13417\.Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- A\. Khandelwal and E\. Pavlick \(2025\)How do language models compose functions?\.arXiv preprint arXiv:2510\.01685\.Cited by:[§4\.3](https://arxiv.org/html/2608.07261#S4.SS3.p5.9)\.
- T\. Kojima, S\. S\. Gu, M\. Reid, Y\. Matsuo, and Y\. Iwasawa \(2022\)Large language models are zero\-shot reasoners\.Advances in neural information processing systems35,pp\. 22199–22213\.Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- Z\. Li, G\. Jiang, H\. Xie, L\. Song, D\. Lian, and Y\. Wei \(2024\)Understanding and patching compositional reasoning in LLMs\.InFindings of the Association for Computational Linguistics: ACL 2024,L\. Ku, A\. Martins, and V\. Srikumar \(Eds\.\),Bangkok, Thailand,pp\. 9668–9688\.External Links:[Link](https://aclanthology.org/2024.findings-acl.576/),[Document](https://dx.doi.org/10.18653/v1/2024.findings-acl.576)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1)\.
- L\. Luo, Y\. Liu, R\. Liu, S\. Phatale, H\. Lara, Y\. Li, L\. Shu, Y\. Zhu, L\. Meng, J\. Sun, and A\. Rastogi \(2024\)Improve mathematical reasoning in language models by automated process supervision\.ArXivabs/2406\.06592\.External Links:[Link](https://api.semanticscholar.org/CorpusID:270379625)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1)\.
- K\. Meng, D\. Bau, A\. Andonian, and Y\. Belinkov \(2022\)Locating and editing factual associations in gpt\.Advances in neural information processing systems35,pp\. 17359–17372\.Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- nostalgebraist \(2020\)Note:LessWrong blog postExternal Links:[Link](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)Cited by:[§4\.1](https://arxiv.org/html/2608.07261#S4.SS1.p1.8),[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- C\. Olah \(2022\)Note:Transformer Circuits essayExternal Links:[Link](https://www.transformer-circuits.pub/2022/mech-interp-essay)Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- OpenAI \(2024\)Learning to reason with llms\.Note:[https://openai\.com/zh\-Hans\-CN/index/learning\-to\-reason\-with\-llms/](https://openai.com/zh-Hans-CN/index/learning-to-reason-with-llms/)Accessed: 2025\-02\-14Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1)\.
- J\. Petty, S\. Steenkiste, I\. Dasgupta, F\. Sha, D\. Garrette, and T\. Linzen \(2024\)The impact of depth on compositional generalization in transformer language models\.InProceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(Volume 1: Long Papers\),K\. Duh, H\. Gomez, and S\. Bethard \(Eds\.\),Mexico City, Mexico,pp\. 7239–7252\.External Links:[Link](https://aclanthology.org/2024.naacl-long.402/),[Document](https://dx.doi.org/10.18653/v1/2024.naacl-long.402)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1)\.
- O\. Press, M\. Zhang, S\. Min, L\. Schmidt, N\. Smith, and M\. Lewis \(2023\)Measuring and narrowing the compositionality gap in language models\.InFindings of the Association for Computational Linguistics: EMNLP 2023,H\. Bouamor, J\. Pino, and K\. Bali \(Eds\.\),Singapore,pp\. 5687–5711\.External Links:[Link](https://aclanthology.org/2023.findings-emnlp.378/),[Document](https://dx.doi.org/10.18653/v1/2023.findings-emnlp.378)Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- \[26\]A\. Radford, J\. Wu, R\. Child, D\. Luan, D\. Amodei, and I\. SutskeverLanguage Models are Unsupervised Multitask Learners\.Cited by:[§2](https://arxiv.org/html/2608.07261#S2.p3.1)\.
- A\. Scherlis, K\. Sachan, A\. S\. Jermyn, J\. Benton, and B\. Shlegeris \(2022\)Polysemanticity and capacity in neural networks\.arXiv preprint arXiv:2210\.01892\.Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- A\. Stolfo, Y\. Belinkov, and M\. Sachan \(2023\)A mechanistic interpretation of arithmetic reasoning in language models using causal mediation analysis\.InConference on Empirical Methods in Natural Language Processing,External Links:[Link](https://api.semanticscholar.org/CorpusID:258865170)Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- J\. Vig, S\. Gehrmann, Y\. Belinkov, S\. Qian, D\. Nevo, Y\. Singer, and S\. M\. Shieber \(2020\)Investigating gender bias in language models using causal mediation analysis\.InNeural Information Processing Systems,External Links:[Link](https://api.semanticscholar.org/CorpusID:227275068)Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- B\. Wang, X\. Yue, Y\. Su, and H\. Sun \(2024\)Grokking of implicit reasoning in transformers: a mechanistic journey to the edge of generalization\.Advances in Neural Information Processing Systems37,pp\. 95238–95265\.Cited by:[Appendix A](https://arxiv.org/html/2608.07261#A1.p1.1),[Appendix A](https://arxiv.org/html/2608.07261#A1.p5.2),[Appendix K](https://arxiv.org/html/2608.07261#A11.p2.5),[Appendix L](https://arxiv.org/html/2608.07261#A12.p1.1),[§2](https://arxiv.org/html/2608.07261#S2.p1.1),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- F\. Wang, W\. Mo, Y\. Wang, W\. Zhou, and M\. Chen \(2023\)A causal view of entity bias in \(large\) language models\.InFindings of the Association for Computational Linguistics: EMNLP 2023,H\. Bouamor, J\. Pino, and K\. Bali \(Eds\.\),Singapore,pp\. 15173–15184\.External Links:[Link](https://aclanthology.org/2023.findings-emnlp.1013/),[Document](https://dx.doi.org/10.18653/v1/2023.findings-emnlp.1013)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p3.1)\.
- J\. Wei, X\. Wang, D\. Schuurmans, M\. Bosma, F\. Xia, E\. Chi, Q\. V\. Le, D\. Zhou,et al\.\(2022\)Chain\-of\-thought prompting elicits reasoning in large language models\.Advances in neural information processing systems35,pp\. 24824–24837\.Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- N\. Xu, F\. Wang, B\. Li, M\. Dong, and M\. Chen \(2022\)Does your model classify entities reasonably? diagnosing and mitigating spurious correlations in entity typing\.InProceedings of the 2022 Conference on Empirical Methods in Natural Language Processing,Y\. Goldberg, Z\. Kozareva, and Y\. Zhang \(Eds\.\),Abu Dhabi, United Arab Emirates,pp\. 8642–8658\.External Links:[Link](https://aclanthology.org/2022.emnlp-main.592/),[Document](https://dx.doi.org/10.18653/v1/2022.emnlp-main.592)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p3.1)\.
- L\. Yang, K\. Lee, R\. Nowak, and D\. Papailiopoulos \(2023\)Looped transformers are better at learning learning algorithms\.arXiv preprint arXiv:2311\.12424\.Cited by:[§5\.2](https://arxiv.org/html/2608.07261#S5.SS2.p1.1)\.
- S\. Yang, E\. Gribovskaya, N\. Kassner, M\. Geva, and S\. Riedel \(2024\)Do large language models latently perform multi\-hop reasoning?\.arXiv preprint arXiv:2402\.16837\.Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p2.1),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1),[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- Z\. Yang, P\. Qi, S\. Zhang, Y\. Bengio, W\. Cohen, R\. Salakhutdinov, and C\. D\. Manning \(2018\)HotpotQA: a dataset for diverse, explainable multi\-hop question answering\.InProceedings of the 2018 Conference on Empirical Methods in Natural Language Processing,E\. Riloff, D\. Chiang, J\. Hockenmaier, and J\. Tsujii \(Eds\.\),Brussels, Belgium,pp\. 2369–2380\.External Links:[Link](https://aclanthology.org/D18-1259/),[Document](https://dx.doi.org/10.18653/v1/D18-1259)Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- X\. Yao, R\. Ren, Y\. Liao, and Y\. Liu \(2025\)Unveiling the Mechanisms of Explicit CoT Training: How CoT Enhances Reasoning Generalization\.arXiv\.External Links:2502\.04667,[Document](https://dx.doi.org/10.48550/arXiv.2502.04667)Cited by:[Appendix L](https://arxiv.org/html/2608.07261#A12.p1.1)\.
- J\. Ye, Z\. Yao, Z\. Huang, L\. Pan, J\. Liu, Y\. Bai, A\. Xin, L\. Weichuan, X\. Che, L\. Hou,et al\.\(2025\)How does transformer learn implicit reasoning?\.arXiv preprint arXiv:2505\.23653\.Cited by:[Appendix A](https://arxiv.org/html/2608.07261#A1.p1.1),[§H\.2](https://arxiv.org/html/2608.07261#A8.SS2.p1.1),[§2](https://arxiv.org/html/2608.07261#S2.p2.11),[§6\.2](https://arxiv.org/html/2608.07261#S6.SS2.p1.1)\.
- Z\. Yu and S\. Ananiadou \(2024\)Interpreting arithmetic mechanism in large language models through comparative neuron analysis\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,Y\. Al\-Onaizan, M\. Bansal, and Y\. Chen \(Eds\.\),Miami, Florida, USA,pp\. 3293–3306\.External Links:[Link](https://aclanthology.org/2024.emnlp-main.193/),[Document](https://dx.doi.org/10.18653/v1/2024.emnlp-main.193)Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- Z\. Yu, Y\. Belinkov, and S\. Ananiadou \(2025\)Back attention: understanding and enhancing multi\-hop reasoning in large language models\.ArXivabs/2502\.10835\.External Links:[Link](https://api.semanticscholar.org/CorpusID:276409219)Cited by:[§6\.3](https://arxiv.org/html/2608.07261#S6.SS3.p1.1)\.
- W\. X\. Zhao, K\. Zhou, J\. Li, T\. Tang, X\. Wang, Y\. Hou, Y\. Min, B\. Zhang, J\. Zhang, Z\. Dong, Y\. Du, C\. Yang, Y\. Chen, Z\. Chen, J\. Jiang, R\. Ren, Y\. Li, X\. Tang, Z\. Liu, P\. Liu, J\. Nie, and J\. Wen \(2025a\)A Survey of Large Language Models\.arXiv\.External Links:2303\.18223,[Document](https://dx.doi.org/10.48550/arXiv.2303.18223)Cited by:[§1](https://arxiv.org/html/2608.07261#S1.p1.3)\.
- Y\. Zhao, A\. Devoto, G\. Hong, X\. Du, A\. P\. Gema, H\. Wang, X\. He, K\. Wong, and P\. Minervini \(2025b\)Steering knowledge selection behaviours in LLMs via SAE\-based representation engineering\.InProceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies \(Volume 1: Long Papers\),L\. Chiruzzo, A\. Ritter, and L\. Wang \(Eds\.\),Albuquerque, New Mexico,pp\. 5117–5136\.External Links:[Link](https://aclanthology.org/2025.naacl-long.264/),[Document](https://dx.doi.org/10.18653/v1/2025.naacl-long.264),ISBN 979\-8\-89176\-189\-6Cited by:[§6\.1](https://arxiv.org/html/2608.07261#S6.SS1.p1.1)\.
- R\. Zhu, Z\. Wang, K\. Hua, T\. Zhang, Z\. Li, H\. Que, B\. Wei, Z\. Wen, F\. Yin, H\. Xing,et al\.\(2025\)Scaling latent reasoning via looped language models\.arXiv preprint arXiv:2510\.25741\.Cited by:[Limitations](https://arxiv.org/html/2608.07261#Sx1.p1.1)\.

## Appendix ADataset Details

Our training set is constructed following the approach ofWanget al\.\([2024](https://arxiv.org/html/2608.07261#bib.bib10)\), while the test set is built following the methodology ofYeet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib11)\)\. Specifically, we first construct 2,000 distinct entities \(<e\_id\>\) and 200 distinct relations \(<r\_id\>\)\.

For atomic knowledge construction, each entity is randomly assigned 20 distinct outgoing relations, where each relation points to another entity—for example,\(⟨e0⟩,⟨r0⟩\)→⟨e1⟩\(\\langle e\_\{0\}\\rangle,\\langle r\_\{0\}\\rangle\)\\rightarrow\\langle e\_\{1\}\\rangle\. This results in a total of 40,000 atomic facts\. To create in\-distribution \(ID\) and out\-of\-distribution \(OOD\) facts, 95% of the atomic facts are assigned to the in\-distribution set \(ID\_atomic\), and the remaining 5% are assigned to the out\-of\-distribution set \(OOD\_atomic\)\.

Two atomic facts can be combined into a two\-hop fact if the tail entity of one fact matches the head entity of the other\. In a two\-hop fact, the first\-hop and second\-hop can each come from either the ID or the OOD atomic set\. We first consider two\-hop facts where both the first and second hops are drawn from the ID atomic set\. These two\-hop facts are further divided: one portion is included in the test set as Test\-II, and from the remaining portion, a number of two\-hop facts corresponding toϕ\\phitimes the number of ID atomic facts are randomly selected and included in the training set\.

Additionally, we construct Test\-IO, Test\-OI, and Test\-OO to more comprehensively evaluate the model’s generalization ability\. Specifically, Test\-IO contains two\-hop facts where the first hop comes from ID atomic facts and the second hop comes from OOD atomic facts; Test\-OI and Test\-OO are defined analogously\.

In this paper, we setϕ=7\.2\\phi=7\.2because, according toWanget al\.\([2024](https://arxiv.org/html/2608.07261#bib.bib10)\), 7\.2 represents a relatively intermediate value\. This corresponds to a training set size that allows the model to generalize on Test\-II while avoiding an excessively large dataset that would significantly increase training cost\. Choosingϕ=7\.2\\phi=7\.2is appropriate for our study of why the model fails to generalize on Test\-IO and Test\-OO\.

More detailed information about the dataset is provided in Table[3](https://arxiv.org/html/2608.07261#A1.T3)and Table[4](https://arxiv.org/html/2608.07261#A1.T4)\. Table[3](https://arxiv.org/html/2608.07261#A1.T3)shows the formats of atomic and two\-hop facts\. Table[4](https://arxiv.org/html/2608.07261#A1.T4)presents the number of facts in each data split\.

Table 3:Example formats of facts in the dataset\.Table 4:Detailed statistics of the dataset\.
## Appendix BTraining Details

All experiments are conducted on Tesla V100\-PCIE\-32GB GPUs\. Training is performed using four Tesla V100\-PCIE\-32GB GPUs, with a batch size of 1024 per GPU\. We use a learning rate of1×10−41\\times 10^\{\-4\}and a weight decay of 0\.1 for optimization\. Inference on the test set is conducted using greedy decoding\. The number of optimization steps and corresponding training time for the main training experiments reported in this paper are summarized in Table[5](https://arxiv.org/html/2608.07261#A2.T5)\.

Table 5:Optimization steps and training time for the main training experiments in this paper\. Standard Training corresponds to Section[3](https://arxiv.org/html/2608.07261#S3), Rep\.\-Based Training \(denoting representation\-based training\) corresponds to Section[5\.1](https://arxiv.org/html/2608.07261#S5.SS1), and Looped Training corresponds to Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3)\.
## Appendix CIntermediate Representations of Bridge Entities Reflect Graph Structure

In this section, we explore the formation of entity representations in relation to the graph structure within the dataset\. We first analyze the structure of the graph composed of all atomic facts, and then examine how this graph structure influences the formation of entity representations\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x10.png)\(a\)Source entity sete1e\_\{1\}
![Refer to caption](https://arxiv.org/html/2608.07261v1/x11.png)\(b\)Target entity sete3e\_\{3\}
![Refer to caption](https://arxiv.org/html/2608.07261v1/x12.png)\(c\)Incoming relation setr1r\_\{1\}
![Refer to caption](https://arxiv.org/html/2608.07261v1/x13.png)\(d\)Outgoing relation setr2r\_\{2\}

Figure 9:Histograms of the set size distributions for thee1e\_\{1\},e3e\_\{3\},r1r\_\{1\}, andr2r\_\{2\}sets across all entities\.![Refer to caption](https://arxiv.org/html/2608.07261v1/x14.png)Figure 10:Illustration of the local graph structure centered at entitye2e\_\{2\}, wherer1r\_\{1\}andr2r\_\{2\}denote the sets of incoming and outgoing relations, respectively, connecting source entitiese1e\_\{1\}and target entitiese3e\_\{3\}\.### C\.1Graph Structure in the Dataset

The 40,000 atomic facts constructed in the dataset jointly form a graph structure\. The resulting graph contains 2,000 nodes, where each node corresponds to an entity\. The graph includes 200 edge types, with each edge type representing a distinct relation\. Each node emits 20 outgoing edges connecting to other nodes, and every node is also the target of incoming edges from other nodes\.

For ease of presentation, we consider an entitye2e\_\{2\}\. The edges emitted bye2e\_\{2\}define a set of outgoing relationsr2r\_\{2\}, and the target entities ofr2r\_\{2\}form a set denoted ase3e\_\{3\}\. The edges pointing toe2e\_\{2\}define a set of incoming relationsr1r\_\{1\}, and the source entities ofr1r\_\{1\}form a set denoted ase1e\_\{1\}\. Consequently, each entity is associated with four sets: a set of target entitiese3e\_\{3\}, a set of outgoing relationsr2r\_\{2\}, a set of source entitiese1e\_\{1\}, and a set of incoming relationsr1r\_\{1\}, as illustrated in Figure[10](https://arxiv.org/html/2608.07261#A3.F10)\.

For all entities, we collect statistics on the sizes of thee1e\_\{1\},r1r\_\{1\},e3e\_\{3\}, andr2r\_\{2\}sets, and plot histograms of the size distributions for each of the four sets, as shown in Figure[9](https://arxiv.org/html/2608.07261#A3.F9)\. As can be clearly observed, the sizes of thee1e\_\{1\}andr1r\_\{1\}sets approximately follow a Gaussian\-like distribution centered around 20\. In contrast, the size of ther2r\_\{2\}set is exactly 20 for all entities, which is consistent with our atomic fact construction procedure\. The size of thee3e\_\{3\}set is 20 for the majority of entities, with a small fraction having size 19\. This deviation arises because, among the 20 outgoing edges of an entity, multiple edges may point to the same target entity, resulting in a reduced number of unique entities ine3e\_\{3\}\.

### C\.2How Does Graph Structure Shape the Formation of Entity Representations?

![Refer to caption](https://arxiv.org/html/2608.07261v1/x15.png)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x16.png)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x17.png)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x18.png)

Figure 11:Relationship between pairwise set overlap and representation similarity for the four attribute setse1e\_\{1\},e3e\_\{3\},r1r\_\{1\}, andr2r\_\{2\}\. The x\-axis denotes the intersection\-over\-union \(IoU\) between the corresponding sets of two entities\. The left y\-axis \(blue curves\) shows the kernel density estimation \(KDE\) of entity pairs at each IoU value\. The right y\-axis \(orange curves\) reports the average cosine similarity of the intermediate representations for entity pairs with IoU greater than or equal to the corresponding threshold\.In this paper, we identify and define the output hidden state atr1r\_\{1\}in layer 5 as the intermediate representation of the bridge entity\. We further observe that, as training progresses, the intermediate representations of the same bridge entity in different contexts gradually converge\. A natural question arises: what does this aggregated entity representation reflect, or what determines its position in the representation space?

An intuitive hypothesis is that the formation of entity representations is related to the graph structure within the dataset\. Specifically, if two entities have similar surrounding attributes—i\.e\., high overlap in theire1e\_\{1\},r1r\_\{1\},e3e\_\{3\}, andr2r\_\{2\}sets—then their representations in the representation space should also be close\.

Based on this hypothesis, we design the following experiment to empirically examine the relationship between local graph structure and entity representation similarity:

1. 1\.For each entity, extract its surrounding attribute set, consisting ofe1e\_\{1\},r1r\_\{1\},e3e\_\{3\}, andr2r\_\{2\}\.
2. 2\.Collect the intermediate hidden state representations of all entities using ID atomic facts, and compute the mean representation for each entity\.
3. 3\.Pair all entities and compute the overlap between their corresponding attribute sets, and plot kernel density estimation \(KDE\) curves for the resulting overlap distributions\.
4. 4\.For multiple overlap thresholds, calculate the average cosine similarity between the representations of all entity pairs exceeding each threshold\.
5. 5\.Plot the average cosine similarity as a function of the overlap threshold\.

This experiment enables us to investigate how local graph structure influences the formation of entity representations\. The results are shown in Figure[11](https://arxiv.org/html/2608.07261#A3.F11)\. The following provides a detailed analysis of these results:

![Refer to caption](https://arxiv.org/html/2608.07261v1/x19.png)

\(a\)

![Refer to caption](https://arxiv.org/html/2608.07261v1/x20.png)

\(b\)

Figure 12:Logit lens results at different stages of the standard training process\. \(a\) Memorization stage \(checkpoint 3k\)\. \(b\) First generalization stage \(Test\-II generalization, checkpoint 100k\)\.KDE distributions\.We first examine the KDE curves\. Most entity pairs concentrate at low set overlap values, indicating that only a small fraction of entity pairs share highly similar surrounding attribute sets\.

Cosine similarity trends\.Next, we consider the cosine similarity curves\. For thee1e\_\{1\},e3e\_\{3\}, andr1r\_\{1\}sets, the cosine similarity consistently increases as the overlap threshold grows, consistent with the expectation that greater structural overlap corresponds to more similar representations\. In contrast, ther2r\_\{2\}set shows decreasing cosine similarity as the overlap threshold increases, appearing to contradict this expectation\.

Magnitude and stability of effects\.A closer inspection reveals an important difference\. The cosine similarity curves fore1e\_\{1\},e3e\_\{3\}, andr1r\_\{1\}remain above zero and exhibit relatively large variations, whereas ther2r\_\{2\}curve is roughly centered around zero, with both positive and negative effects and a much smaller range of variation\. This suggests that the overlap of ther2r\_\{2\}set has only a negligible influence on the cosine similarity between entity representations\.

Overall, higher overlap in thee1e\_\{1\},e3e\_\{3\}, andr1r\_\{1\}attribute sets corresponds to greater similarity between the representations of two entities\. This indicates that the intermediate entity representations capture implicit graph structures within the dataset, and that, as training progresses, the model tends to organize its hidden representations in a more structured and efficient manner\.

## Appendix DMore Detailed Logit Lens Results

Figure[3](https://arxiv.org/html/2608.07261#S4.F3)in Section[4\.1](https://arxiv.org/html/2608.07261#S4.SS1)shows the logit lens results of the final model, where the bridge entity representation is identified ashr15h\_\{r\_\{1\}\}^\{5\}by locating the starting point of the second hop\. As a supplement to that analysis, Figure[12](https://arxiv.org/html/2608.07261#A3.F12)presents logit lens results from earlier training stages\.

During the memorization stage \(Figure[12](https://arxiv.org/html/2608.07261#A3.F12)a\), the probability ofe3e\_\{3\}becomes significant even when the representation of the bridge entitye2e\_\{2\}is still not clearly formed\. This indicates that, at this stage, the model does not rely on the bridge entity representation when answering two\-hop queries, but rather memorizes the final two\-hop answer directly, without performing genuine multi\-hop reasoning\.

In the first generalization stage \(Figure[12](https://arxiv.org/html/2608.07261#A3.F12)b\), the logit lens result is consistent with that of the final model\. At ther2r\_\{2\}position, the probability ofr2r\_\{2\}begins to decrease significantly at Layer 6, while the probability ofe3e\_\{3\}starts to increase, marking the beginning of the second hop\. This indicates that the bridge entity representation emerges athr15h\_\{r\_\{1\}\}^\{5\}, providing empirical support for the choice made in Section[4\.2](https://arxiv.org/html/2608.07261#S4.SS2), where this representation is patched at different stages of training\.

## Appendix ELayer\-wise Attention Masking with a Sliding Window

As discussed in Section[4\.3](https://arxiv.org/html/2608.07261#S4.SS3), when answering two\-hop queries\(e1,r1,r2,e3\)\(e\_\{1\},r\_\{1\},r\_\{2\},e\_\{3\}\), the upper layers \(Layers 6, 7, and 8\) of the model must perform a nonlinear transformation at the position ofr2r\_\{2\}, mappinghe2→r2e3h\_\{e\_\{2\}\}\\xrightarrow\{r\_\{2\}\}e\_\{3\}\. In contrast, when learning single\-hop facts\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\), the upper layers only learn a linear mappinghe3→Ee3h\_\{e\_\{3\}\}\\rightarrow E\_\{e\_\{3\}\}\. To further support the findings in Section[4\.3](https://arxiv.org/html/2608.07261#S4.SS3), we conduct a layer\-wise attention masking experiment using a sliding window\. The masking window spans three consecutive layers\.

For atomic facts\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\), we mask the attention fromr2r\_\{2\}toe2e\_\{2\}using a three\-layer window with different window centers, and evaluate whether the model can still correctly predicte3e\_\{3\}after masking\. For two\-hop queries\(e1,r1,r2,e3\)\(e\_\{1\},r\_\{1\},r\_\{2\},e\_\{3\}\), we instead mask the attention fromr2r\_\{2\}tor1r\_\{1\}using the same window with different centers and test whether the model can still correctly predicte3e\_\{3\}\. The purpose of this experiment is to verify whether the upper layers truly rely on information from thee2e\_\{2\}position when answering single\-hop facts\. The results are shown in Figure[13](https://arxiv.org/html/2608.07261#A5.F13)\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x21.png)Figure 13:Accuracy of predicting the target entity under attention masking with a three\-layer window\. For atomic facts\(e2,r2,e3\)\(e\_\{2\},r\_\{2\},e\_\{3\}\), we mask attention fromr2r\_\{2\}toe2e\_\{2\}; for two\-hop queries\(e1,r1,r2,e3\)\(e\_\{1\},r\_\{1\},r\_\{2\},e\_\{3\}\), we mask attention fromr2r\_\{2\}tor1r\_\{1\}\. The x\-axis denotes the center layer of the masking window\.As shown in Figure[13](https://arxiv.org/html/2608.07261#A5.F13), for atomic facts, masking attention in the lower layers significantly reduces the model’s accuracy, whereas masking attention in the upper layers \(i\.e\., when the masking window lies entirely within Layers 6–8\) has only a minor effect\.

In contrast, for two\-hop queries, masking attention in the lower layers has almost no effect on the model’s accuracy\. However, when the masking window covers Layers 6–8, the model’s accuracy drops to nearly zero\.

These results confirm that when answering single\-hop facts, the upper layers indeed do not require information from thee2e\_\{2\}position\. Consequently, during training on atomic facts, the ability to perform single\-hop reasoning is not learned in the upper layers\. However, answering two\-hop queries relies on this single\-hop reasoning capability in the upper layers\. This creates a mismatch between the capability required for two\-hop reasoning and the capability actually learned from atomic facts during training\. As a result, the model fails to generalize on Test\-IO and Test\-OO, since learning OOD atomic facts alone does not equip the upper layers with the ability to perform OOD single\-hop reasoning\.

## Appendix FDetails of Representation\-based Training Experiments

This section provides additional details on the training procedure described in Section[5\.1](https://arxiv.org/html/2608.07261#S5.SS1)\. The representation\-based training procedure is illustrated in Figure[14](https://arxiv.org/html/2608.07261#A6.F14)\. In this experiment, besides the standard training procedure, we additionally train Layers 6–8 to learn OOD atomic facts\(e2,r2\)\(e\_\{2\},r\_\{2\}\)in a representation\-based form\(he2,hr2\)\(h\_\{e\_\{2\}\},h\_\{r\_\{2\}\}\)\. The procedure is as follows:

1. 1\.Standard training:We first train the model normally until it achieves high generalization performance on Test\-II and Test\-OI, as only when the model generalizes well on these benchmarks do the intermediate representations of the same entity become consistent and stable, which in turn provides a reliable foundation for subsequent representation\-based training\. For this experiment, we directly use the model after 3,500,000 optimization steps of standard training\.
2. 2\.Multi\-round training with representation\-based reasoning:The following procedure is repeated for multiple rounds\. In each round, we first perform standard training to maintain the model’s generalization performance and prevent the additional representation\-based reasoning training from degrading it\. After this, we collect the intermediate representations from Train\-II:he=h5​\(r1\)h\_\{e\}=h\_\{5\}\(r\_\{1\}\)for all entities andhr=h5​\(r2\)h\_\{r\}=h\_\{5\}\(r\_\{2\}\)for all relations\. These representations are then used to construct representation\-based forms\(he2,hr2\)\(h\_\{e\_\{2\}\},h\_\{r\_\{2\}\}\)for all OOD atomic facts\(e2,r2\)\(e\_\{2\},r\_\{2\}\)\. The pair\(he2,hr2\)\(h\_\{e\_\{2\}\},h\_\{r\_\{2\}\}\)is fed into Layer 6 of the model and propagated through the subsequent layers, with supervision applied to guide the model to output the correct answers\.

Test accuracy is evaluated throughout this process\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x22.png)Figure 14:Illustration of the representation\-based training process, where intermediate representations of entities and relations for OOD atomic facts are extracted from Train\-II \(layer 5\) and used to train Layers 6–8\.
## Appendix GLogit Lens Analysis of Looped Training Models

Figure[15](https://arxiv.org/html/2608.07261#A7.F15)shows the logit lens results of the final model trained under the recurrent architecture when answering two\-hop questions\. For a token at positionttin layerll, we extract its hidden statehl​\(t\)h\_\{l\}\(t\), apply a LayerNorm, and project it through the model’s embedding transposeE⊤E^\{\\top\}to obtain the probability of each target entity or relation:

pl​\(x∣t\)\\displaystyle p\_\{l\}\(x\\mid t\)=softmax​\(E⊤​LayerNorm​\(hl​\(t\)\)\),\\displaystyle=\\text\{softmax\}\\Big\(E^\{\\top\}\\,\\text\{LayerNorm\}\(h\_\{l\}\(t\)\)\\Big\),\(1\)x∈\{e2,e3,r2\}\.\\displaystyle\\quad x\\in\\\{e\_\{2\},e\_\{3\},r\_\{2\}\\\}\.This allows us to inspect how the model’s predictions fore2,e3e\_\{2\},e\_\{3\}andr2r\_\{2\}evolve across layers\.

From Figure[15](https://arxiv.org/html/2608.07261#A7.F15), we observe that Layer 6 marks the onset of the second\-hop reasoning: at the position ofr2r\_\{2\}, the probability ofr2r\_\{2\}begins to decrease while the probability ofe3e\_\{3\}starts to increase\. This indicates that Layer 6 is the first layer where the model initiates second\-hop inference\. Following the approach in Section[4\.1](https://arxiv.org/html/2608.07261#S4.SS1), we define the hidden state at layer 5 corresponding to ther1r\_\{1\}position, denoted ash5​\(r1\)h\_\{5\}\(r\_\{1\}\), as the intermediate representation of the bridge entitye2e\_\{2\}, and the hidden state at layer 5 corresponding to ther2r\_\{2\}position, denoted ash5​\(r2\)h\_\{5\}\(r\_\{2\}\), as the intermediate representation of relationr2r\_\{2\}\. During the second\-hop reasoning, the model performs inference based on these intermediate representationsh5​\(r1\)h\_\{5\}\(r\_\{1\}\)andh5​\(r2\)h\_\{5\}\(r\_\{2\}\)\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x23.png)Figure 15:Logit lens of the final model trained under the looped architecture, showing hidden state predictions fore2e\_\{2\},e3e\_\{3\}, andr2r\_\{2\}\. Layer 6 marks the start of second\-hop reasoning\.
## Appendix HExploring Equivalent Alternatives to the Shared\-Parameter Model

![Refer to caption](https://arxiv.org/html/2608.07261v1/x24.png)Figure 16:Comparison of two training strategies described in Appendix[H](https://arxiv.org/html/2608.07261#A8)\.Left:the training strategy described in Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1), where in each training round, embeddings of atomic facts are additionally used to train the upper four layers\.Right:the training strategy described in Appendix[H\.2](https://arxiv.org/html/2608.07261#A8.SS2), where atomic facts are removed from the original training set, and in each training round, embeddings of atomic facts are additionally used to train both the upper four layers and the lower four layers\.In Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3), we empirically demonstrate that the shared\-parameter model generalizes well on both Test\-IO and Test\-OO\. We further explain its effectiveness: sharing the parameters of the lower and upper four layers forces the model to align the embedding\-level representations and hidden\-state representations of atomic facts, while naturally transferring the reasoning capability learned in the lower layers to the upper layers\. As a result, even when the model is only exposed to embedding\-based reasoning patterns of out\-of\-distribution atomic facts during training, it can still perform reasoning over the hidden representations of out\-of\-distribution atomic facts at test time\.

As a complement to Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3), in this section we explore alternative mechanisms that can serve as equivalent substitutes for the shared\-parameter model, with the aim of gaining a deeper understanding of the fundamental sources of its effectiveness\. Specifically, we investigate two alternative training strategies that aim to replicate the key effects of parameter sharing without explicitly tying parameters across layers, as illustrated in Figure[16](https://arxiv.org/html/2608.07261#A8.F16)\. The left panel of the figure corresponds to a strategy that trains only the upper layers using embedded atomic facts, while the right panel illustrates a strategy that trains both the upper and lower layers separately under the same embedded\-atomic\-fact supervision\.

### H\.1Training Upper Layers with Embedded Atomic Facts

![Refer to caption](https://arxiv.org/html/2608.07261v1/x25.png)\(a\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x26.png)\(b\)

Figure 17:Training results for the two training strategies described in Appendix[H](https://arxiv.org/html/2608.07261#A8): \(a\) the strategy corresponding to the left panel in Figure[16](https://arxiv.org/html/2608.07261#A8.F16), where embeddings of atomic facts are additionally used to train the upper layers; \(b\) the strategy corresponding to the right panel in Figure[16](https://arxiv.org/html/2608.07261#A8.F16), where embeddings of atomic facts are used to train both the upper and lower layers\. Note that the accuracy on atomic facts is not shown in \(b\) because this training set does not include supervision that uses all layers to answer atomic facts\.In this section, we augment the standard training pipeline with an additional procedure\. After each training epoch, we further train the upper four layers of the model using all atomic facts, including both in\-distribution \(ID\) and out\-of\-distribution \(OOD\) atomic facts\. Specifically, we directly feed the embeddings of atomic facts\(Ee,Er\)\(E\_\{e\},E\_\{r\}\)into the fifth layer of the model, and supervise the model outputs using the target entities of these atomic facts\. This procedure is designed to simulate a key property of the shared\-parameter model, namely, enabling the upper layers to perform embedding\-based reasoning\.

The training results are shown in Figure[17\(a\)](https://arxiv.org/html/2608.07261#A8.F17.sf1)\. Compared with the standard training setting, the results differ in the following aspects:

- •Test\-IO and Test\-OO\.The augmented training procedure enables strong generalization on Test\-IO, achieving nearly 100% accuracy, and weak but non\-zero generalization on Test\-OO, with around 4% accuracy\. In contrast, under the standard training setting, the model achieves almost zero accuracy on both Test\-IO and Test\-OO\.
- •Test\-OI\.Despite the improvements on Test\-IO and Test\-OO, the model only achieves limited generalization on Test\-OI, with around 6% accuracy\. By comparison, the standard training procedure yields over 80% accuracy on Test\-OI\.

We interpret these results as follows\. Training the upper layers with embedding\-based atomic facts reshapes their parameters and forces the intermediate representations of entities and relations to align with their corresponding embeddings\. Moreover, since the upper layers are explicitly trained on OOD atomic facts, the model is able to generalize to Test\-IO\. However, introducing this additional training procedure also alters the parameters of the upper layers, which interferes with the model’s ability to generalize on Test\-OI\.

Despite this limitation, the experiment remains informative\. The fact that this training strategy enables generalization on Test\-IO supports our hypothesis that the effectiveness of the shared\-parameter model stems from endowing the upper layers with embedding\-based reasoning capability\.

### H\.2Training Upper and Lower Layers Separately with Embedded Atomic Facts

In this section, we attempt to address the issue of disrupted generalization on Test\-OI observed in Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1)\. According to our analysis in Section[4\.2](https://arxiv.org/html/2608.07261#S4.SS2), the model’s ability to generalize on Test\-OI arises from the consistency of representations for the same entity across different contexts\. Furthermore, as analyzed byYeet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib11)\), the aggregation and alignment of representations generated from out\-of\-distribution atomic facts are enabled by the presence of in\-distribution atomic facts that share the same target entity\. In the training strategy adopted in Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1), the upper layers of the model are trained on atomic facts in two different phases: once through the standard dataset training and once through the additional embedding\-based training\. We hypothesize that this repeated training of atomic facts in different forms interferes with the model’s ability to generalize on Test\-OI\.

To test this hypothesis, we decouple the standard atomic\-fact training from the additional embedding\-based training\. Specifically, in each training epoch, we train the lower four layers and the upper four layers separately using all atomic facts, while removing all atomic facts from the original training set and retaining only Train\-II\. The training results are shown in Figure[17\(b\)](https://arxiv.org/html/2608.07261#A8.F17.sf2)\. Compared with Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1), we observe that, while maintaining high accuracy on Test\-IO, the model achieves 25% accuracy on Test\-OI and 15% accuracy on Test\-OO\. Both results are approximately four times higher than those reported in Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1)\. These results support our hypothesis that the degradation in Test\-OI generalization is indeed caused by interference introduced by the additional training procedure\.

Although the proposed training strategy does not reach the accuracy achieved by the shared\-parameter model on Test\-OI and Test\-OO \(80% and 60%, respectively\), it nonetheless attains a substantial level of performance and provides a closer functional approximation to shared\-parameter training\. Compared with the setting in Appendix[H\.1](https://arxiv.org/html/2608.07261#A8.SS1), this experiment more strongly suggests that the success of shared\-parameter training stems from enabling both the upper and lower layers of the model to perform embedding\-based reasoning over atomic facts\.

## Appendix IEffect of Model Scale on the Training Performance of Looped Transformers

![Refer to caption](https://arxiv.org/html/2608.07261v1/x27.png)Figure 18:Test accuracy during training on looped transformers with varying model scale, controlled by the number of layers\.In this section, we investigate how model scale affects the training dynamics of looped transformers by varying the number of layers\. All models follow the GPT2 architecture with a hidden size of 768\. We adopt the looped transformer design, where the parameters of the upper and lower halves of the network are shared across layers\.

Specifically, we experiment with models of 4, 8, 16, and 24 layers\. In the 4\-layer setting, the parameters of the bottom 2 layers are shared with the top 2 layers\. Similarly, the 8\-layer model shares parameters between the bottom 4 and top 4 layers, and so on\. The training results are shown in Figure[18](https://arxiv.org/html/2608.07261#A9.F18)\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x28.png)Figure 19:Visualization of how a model trained with a looped architecture transfers embedding\-based atomic fact reasoning learned in lower layers to higher layers\. The bottom part illustrates the reasoning process for a two\-hop instance, while the top part shows the embedding\-based reasoning process corresponding to the second hop of the same instance\.From Figure[18](https://arxiv.org/html/2608.07261#A9.F18), we observe the following trends:

- •Generalization accuracy\.Looped transformers with different layer counts can all achieve nearly 100% accuracy on Test\-II and Test\-IO\. However, on Test\-IO and Test\-OO, models with more layers demonstrate better generalization performance\. Models with fewer layers exhibit weaker generalization ability\. In particular, the 4\-layer model shows unstable training behavior and achieves only around 10% accuracy on Test\-OO\. This phenomenon is likely due to the limited capacity of small\-scale models, which is insufficient to support robust generalization\.
- •Generalization speed\.Models with more layers tend to generalize faster during training\. In addition, deeper models show more stable generalization dynamics, with smaller fluctuations in accuracy throughout the training process\.

## Appendix JDetailed Evidence on Alignment under Looped Training

This section supplements Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3)\. We provide a detailed description of how the results in Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3)are computed and present additional evidence for*Representation–Input Alignment under Looped Training*\. In Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3), we only report the alignment results for \(i\)\(h4​\(r1\),h4​\(r2\)\)\(h\_\{4\}\(r\_\{1\}\),h\_\{4\}\(r\_\{2\}\)\)vs\.\(Ee2,Er2\)\(E\_\{e\_\{2\}\},E\_\{r\_\{2\}\}\), and \(ii\)\(h5​\(r1\),h5​\(r2\)\)\(h\_\{5\}\(r\_\{1\}\),h\_\{5\}\(r\_\{2\}\)\)vs\.\(h1​\(e2\),h1​\(r2\)\)\(h\_\{1\}\(e\_\{2\}\),h\_\{1\}\(r\_\{2\}\)\)\. Here, we extend this analysis by reporting cosine similarities between hidden states across a broader range of layer pairs\.

We first describe how the results in Section[5\.3](https://arxiv.org/html/2608.07261#S5.SS3)are computed\. For each two\-hop instance from the training set \(Train\-II\), we identify the atomic fact corresponding to its second hop, which is in\-distribution \(ID\)\. For each two\-hop instance from the test set \(Test\-IO\), we identify the atomic fact corresponding to its second hop, which is out\-of\-distribution \(OOD\)\. This matching strategy allows us to examine whether the model reuses embedding\-level representations of atomic facts when performing the second hop of reasoning, as illustrated in Figure[19](https://arxiv.org/html/2608.07261#A9.F19)\.

Specifically, for a two\-hop query⟨e1,r1,r2⟩\\langle e\_\{1\},r\_\{1\},r\_\{2\}\\rangle, we extract the hidden states at the position ofr1r\_\{1\}from all Transformer layers, and compare them with the hidden states at the position ofe2e\_\{2\}from the corresponding second\-hop atomic fact⟨e2,r2⟩\\langle e\_\{2\},r\_\{2\}\\rangle\. We compute a layer\-wise cosine similarity matrix of size\(L\+1\)×\(L\+1\)\(L\+1\)\\times\(L\+1\), whereLLdenotes the number of Transformer layers555Layer 0 corresponds to the input embedding without positional encoding\.\. Similarly, we extract the hidden states at the position ofr2r\_\{2\}in⟨e1,r1,r2⟩\\langle e\_\{1\},r\_\{1\},r\_\{2\}\\rangleand compare them with the hidden states at the position ofr2r\_\{2\}in the atomic fact⟨e2,r2⟩\\langle e\_\{2\},r\_\{2\}\\rangle, again computing a\(L\+1\)×\(L\+1\)\(L\+1\)\\times\(L\+1\)cosine similarity matrix\.

Finally, we average the cosine similarity matrices over all matched two\-hop instances to obtain the final mean cosine similarity matrices\. We present the averaged similarity matrices separately for Train\-II and Test\-IO, and compare models trained with the looped architecture against normally trained models\. Figure[21](https://arxiv.org/html/2608.07261#A10.F21)visualizes the resulting mean cosine similarity matrices\.

From Figure[21](https://arxiv.org/html/2608.07261#A10.F21), it is clear that for models trained under the looped architecture \(left column\), the cosine similarity matrices at positionsr1r\_\{1\}andr2r\_\{2\}for both Train\-II and Test\-IO \(\(a\), \(c\), \(e\), and \(g\)\) show significantly higher values along the diagonal directions within the red\-square\-highlighted regions\. Notably, the entries at position \(1,5\) \(row, column\) exhibit particularly strong similarity across all subfigures\. In contrast, models trained in the standard manner do not display this pattern\. This indicates that for models trained under the looped architecture, the implicit second\-hop reasoning in the two\-hop queries is aligned with the reasoning over the corresponding single\-hop atomic facts \(based on embeddings\), achieving strong alignment at the fifth layer\. Under this input alignment condition, the model can naturally transfer the reasoning ability over out\-of\-distribution atomic facts learned in lower layers to higher layers, enabling effective generalization to both Test\-IO and Test\-OO\.

![Refer to caption](https://arxiv.org/html/2608.07261v1/x29.png)\(a\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x30.png)\(b\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x31.png)\(c\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x32.png)\(d\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x33.png)\(e\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x34.png)\(f\)

![Refer to caption](https://arxiv.org/html/2608.07261v1/x35.png)\(a\)
![Refer to caption](https://arxiv.org/html/2608.07261v1/x36.png)\(b\)

Figure 21:Layerwise cosine similarity heatmaps comparing models trained under a recurrent architecture \(left column\) and models trained in the standard manner \(right column\)\. For each two\-hop query⟨e1,r1,r2⟩\\langle e\_\{1\},r\_\{1\},r\_\{2\}\\rangle, hidden states at the position ofr1r\_\{1\}are compared with the corresponding hidden states ate2e\_\{2\}in the atomic fact⟨e2,r2⟩\\langle e\_\{2\},r\_\{2\}\\rangle, and hidden states at the position ofr2r\_\{2\}are compared with the hidden states atr2r\_\{2\}in the atomic fact\. Subfigures \(a\) and \(b\) show Train\-II at positionr1r\_\{1\}, \(c\) and \(d\) show Train\-II at positionr2r\_\{2\}, \(e\) and \(f\) show Test\-IO at positionr1r\_\{1\}, and \(g\) and \(h\) show Test\-IO at positionr2r\_\{2\}\.
## Appendix KExtraction of a Real\-World Two\-Hop Dataset

To evaluate the effectiveness of looped transformers on real\-world two\-hop reasoning, we construct a dataset from Wikidata5M by extracting a dense subgraph\.

Prior work has shown that successful generalization on two\-hop reasoning critically depends on the ratio between compositional training samples \(Train\-II\) and atomic facts\.Wanget al\.\([2024](https://arxiv.org/html/2608.07261#bib.bib10)\)identify a threshold phenomenon: models exhibit grokking behavior on Test\-II only when this ratio, denoted asϕ\\phi, exceeds approximately 5\.4\. However, real\-world multi\-hop datasets are typically too sparse to satisfy this condition\. For example,Abramovet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib43)\)report that the graph induced by 2WikiMultiHopQA hasϕ≈0\.5\\phi\\approx 0\.5, which is insufficient for grokking to emerge\. To address this issue,Abramovet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib43)\)increaseϕ\\phiby augmenting the knowledge graph with additional edges, reachingϕ=6\.25\\phi=6\.25and enabling generalization on Test\-II, although the model still fails to generalize to Test\-OO\. In our work, instead of augmenting the dataset, we directly extract a dense subgraph from Wikidata5M, which provides a simple and controlled way to increaseϕ\\phiwhile focusing on evaluating the effectiveness of looped transformers on real\-world data\.

### K\.1Subgraph Construction

We construct the subgraph from Wikidata5M by focusing on entities of typehuman\. We first collect all such entities and compute their in\-degrees and out\-degrees within the original graph\. Each entity is then assigned a score:

score​\(e\)=in\-degree​\(e\)×out\-degree​\(e\)in\-degree​\(e\)\+out\-degree​\(e\)\.\\text\{score\}\(e\)=\\frac\{\\text\{in\-degree\}\(e\)\\times\\text\{out\-degree\}\(e\)\}\{\\text\{in\-degree\}\(e\)\+\\text\{out\-degree\}\(e\)\}\.This scoring function favors entities that simultaneously have high in\-degree and out\-degree, i\.e\., entities that are well\-connected both as sources and as targets\. Intuitively, such entities are more likely to serve as intermediate nodes in two\-hop paths, thereby contributing to a denser set of compositional relations\.

We rank entities by this score and select the top 100 as central nodes\. The subgraph is then constructed by including these entities and all their neighboring entities as nodes, along with all edges among them\. This yields a set of single\-hop facts and all candidate two\-hop paths within the extracted subgraph\.

### K\.2Two\-Hop Path Filtering and Balancing

We further refine the two\-hop samples to mitigate shortcut learning and encourage genuine compositional reasoning\.

First, we remove cyclic paths where the head and tail entities are identical\. Second, we control the target distribution for each relation pair\(r1,r2\)\(r\_\{1\},r\_\{2\}\)\. Without such control, certain targets may dominate a given relation pair, allowing the model to exploit spurious correlations between\(r1,r2\)\(r\_\{1\},r\_\{2\}\)ande3e\_\{3\}rather than performing two\-hop reasoning\. Concretely, for each\(r1,r2\)\(r\_\{1\},r\_\{2\}\), we group paths by their target entitye3e\_\{3\}and enforce a maximum proportion constraint\. LetNNdenote the total number of retained paths for this pair\. We ensure that no single target entity accounts for more than0\.2​N0\.2N\. In addition, we discard relation pairs with fewer than 10 distinct target entities to ensure sufficient diversity\.

After filtering, the ratio of Train\-II paths to single\-hop samples \(ϕ\\phi\) is 9\.60\. For the actual training set, we sample Train\-II paths to achieveϕ=7\.2\\phi=7\.2, providing a dense yet controlled setting for evaluating the effectiveness of looped transformers on real\-world two\-hop reasoning\. The statistics of the resulting dataset are summarized in Table[6](https://arxiv.org/html/2608.07261#A11.T6)\.

CategoryData TypeQuantityGraph StatisticsEntities14,995Relations176Training SetID\_atomic15,375OOD\_atomic809Train\-II110,700Test SetID\_atomic1,000OOD\_atomic1,000Train\-II1,000Test\-II1,000Test\-IO1,000Test\-OI1,000Test\-OO480Table 6:Statistics of the extracted two\-hop dataset and subgraph\. The graph statistics include the total number of entities and relations\. Training and test splits are further broken down by sample type\.

## Appendix LComparison with Explicit CoT Training

Although our work focuses on implicit reasoning, it is necessary to compare our approach with explicit Chain\-of\-Thought \(CoT\) supervision\. Prior work byYaoet al\.\([2025](https://arxiv.org/html/2608.07261#bib.bib44)\), using the same dataset setting introduced byWanget al\.\([2024](https://arxiv.org/html/2608.07261#bib.bib10)\), demonstrated that explicit CoT supervision can substantially improve out\-of\-distribution two\-hop generalization, achieving nearly 100% accuracy on the Test\-OO split\.

In contrast, our proposed solution, the looped transformer, does not achieve perfect Test\-OO accuracy\. However, it addresses a fundamental architectural limitation of implicit reasoning models: the upper layers of standard transformers lack robust out\-of\-distribution reasoning capabilities\. Moreover, explicit CoT supervision requires constructing additional supervision signals for intermediate entities, which introduces extra annotation procedures and training overhead\. Our approach, by comparison, only requires modifying the model architecture without introducing additional supervision\.