REVIEW 2 cited by
Adding identity supervision on bridge tokens enables out-of-distribution two-hop reasoning in simple transformers, with a nuclear-norm theory explaining the benefit.
Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →
T0 review · deepseek-v4-flash
2026-08-04 13:52 UTC pith:KUWBA5H3
Unveiling the Mechanisms of Multi-Hop Reasoning in Transformers via Identity Bridge
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Why would a B-to-B lesson help? The model stores facts as directions in a high-dimensional space. The identity bridge forces the representation of B to point both to itself and toward its C. Because the one-hop task A-to-B already points to B, the two pieces line up, so the network can read off C directly. The paper proves this with a simplified model where training is described by a nuclear-norm minimization. With identity bridge, the optimal solution has a positive margin on every held-out composition; without it, the margin is negative, matching the observed failure.
The effect weakens when the dataset is more complex. The authors show that small initialization or weight decay, which strengthen the implicit regularization, keeps the benefit. Tests on real pretrained language models are more tentative: correct two-hop answers correlate with a stronger subject-to-object signal, but explicitly fine-tuning with identity data gives no significant accuracy gain. The theory also leans on several unproved assumptions, so the mechanism is not fully closed.
Core claim
The paper's central assertion is that adding a zero-hop identity supervision on bridge tokens unlocks OOD two-hop generalization: 'Under this supervision, even a one-layer transformer with uniform attention (Emb-MLP) can achieve out-of-distribution (OOD) two-hop generalization.' Theorem 1 formalizes this for the nuclear-norm program: for every OOD query X=(a_i,r1,r2) with label y=c_i, the multiclass margin q(X,y)>0, so the composed map g2∘g1 is recovered on held-out compositions.
Load-bearing premise
The proof of Theorem 1 relies on unproved structural assumptions: Assumption 2 (the optimal solution satisfies 1^T W^T = 0) is stated 'on the basis of observations from numerical experiments', Assumption 4 asserts c1>0 and Θ(1) parameters as non-degeneracy, and condition (12) (a1=1 with constraint g2 tight) is assumed then justified only through a sketchy KKT case analysis. Additionally, the theorem is about the convex nuclear-norm solution, not the actual gradient-trained network; the bridge between Eq. 1 and Eq. 2 is empirical, and the paper admits a KKT point of Eq. 1 does not in general certify optimality for Eq. 2. If any of these fail, the positive OOD margin conclusion may not hold.
Editorial analysis
A structured set of objections, weighed in public.
Axiom & Free-Parameter Ledger
free parameters (2)
- initialization scaling exponent gamma =
gamma > 0.5 (exact value not reported)
- weight decay coefficient =
not reported
axioms (8)
- domain assumption A KKT point of the homogeneous margin program (Eq. 1) corresponds to the solution of the nuclear-norm program (Eq. 2) for Emb-MLP.
- standard math Huang et al. Lemma 1: an optimal solution of Eq. 2 with identity task exists in restricted form with parameters a1...h.
- standard math Assumption 1 (Huang et al.): if the solution is not of restricted form, then 1^T W^T != 0.
- ad hoc to paper Assumption 2: 1^T_{2n} W^T = 0 for the optimal solution.
- standard math Assumption 3: any optimal solution satisfies KKT conditions.
- ad hoc to paper Assumption 4: c1 > 0 and parameters are Theta(1).
- domain assumption Emb-MLP with uniform attention is a faithful proxy for GPT-2's two-hop mechanism.
- standard math Permutation symmetry of the dataset implies symmetry of the optimal solution after averaging.
invented entities (1)
-
Identity bridge (zero-hop identity supervision on bridge tokens)
independent evidence
read the original abstract
Large Language Models (LLMs) excel at multi-hop reasoning in distribution, yet fail on unseen compositions, a phenomenon known as the curse of two-hop reasoning. In this work, we argue that this phenomenon can be attributed to a missing supervision on the bridge entity. We formalize this gap by introducing identity bridge, a minimal supervision that enforces a identity mapping on bridge tokens. Under this supervision, even a one-layer transformer with uniform attention (Emb-MLP) can achieve out-of-distribution (OOD) two-hop generalization. We provide a theoretical analysis demonstrating that identity bridge induces an implicit regularization effect, leading the model to establish a direct subject-to-answer association. From an empirical perspective, the performance of standard GPT-2 models aligns closely with simple Emb--MLP models across varying levels of problem complexity. Finally, analyses of fine-tuned mainstream LLMs indicate that correct two-hop predictions consistently coincide with the establishment of a subject-to-answer relationship, extending our findings to realistic settings.
Figures
Forward citations
Cited by 2 Pith papers
-
Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge
Identity-bridge regularization, rephrased into an out-of-context reasoning form, yields ~40% reversal accuracy in a 1B LLM and provably fixes reversal in an idealized one-layer transformer.
-
The Latent Space: Foundation, Evolution, Mechanism, Ability, and Outlook
A large survey organizes latent-space work in language-based models by foundation, evolution, four mechanisms, seven abilities, and open challenges.
Reference graph
Works this paper leans on
-
[1]
Physics of language models: part 3.1, knowledge storage and extraction
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: part 3.1, knowledge storage and extraction. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024
2024
-
[2]
Physics of language models: Part 3.2, knowledge manipulation
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=oDbiL9CLoS
2025
-
[3]
The two-hop curse: Llms trained on a to b, b to c fail to learn a to c
Mikita Balesni, Tomek Korbak, and Owain Evans. The two-hop curse: Llms trained on a to b, b to c fail to learn a to c. arXiv preprint arXiv:2411.16353, 2024
arXiv 2024
-
[4]
The reversal curse: LLM s trained on a is b fail to learn b is a
Lukas Berglund, Meg Tong, Maximilian Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. The reversal curse: LLM s trained on a is b fail to learn b is a . In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=GPKTIktA0k
2024
-
[5]
Birth of a transformer: A memory viewpoint
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=3X2EbBLNsk
2023
-
[6]
Hopping too late: Exploring the limitations of large language models on multi-hop queries
Eden Biran, Daniela Gottesman, Sohee Yang, Mor Geva, and Amir Globerson. Hopping too late: Exploring the limitations of large language models on multi-hop queries. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 14113--14130, Miami, Florida, USA, N...
-
[7]
Learning associative memories with gradient descent, 2024
Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent, 2024. URL https://arxiv.org/abs/2402.18724
Pith/arXiv arXiv 2024
-
[8]
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442, 2024
Pith/arXiv arXiv 2024
-
[9]
Universal transformers
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=HyzdRiR9Y7
2019
-
[10]
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, and Jena D. Hwang. Faith and fate: Limits of transformers on compositionality. ArXiv, abs/2305.18654, 2023
Pith/arXiv arXiv 2023
-
[11]
Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024. URL https://arxiv.org/abs/2402.11004
Pith/arXiv arXiv 2024
-
[12]
Reverse training to nurse the reversal curse
Olga Golovneva, Zeyuan Allen-Zhu, Jason Weston, and Sainbayar Sukhbaatar. Reverse training to nurse the reversal curse. arXiv preprint arXiv:2403.13799, 2024
Pith/arXiv arXiv 2024
-
[13]
Mitigating reversal curse in large language models via semantic-aware permutation training
Qingyan Guo, Rui Wang, Junliang Guo, Xu Tan, Jiang Bian, and Yujiu Yang. Mitigating reversal curse in large language models via semantic-aware permutation training. arXiv preprint arXiv:2403.00758, 2024
Pith/arXiv arXiv 2024
-
[14]
How do llms perform two-hop reasoning in context? arXiv preprint arXiv:2502.13913, 2025
Tianyu Guo, Hanlin Zhu, Ruiqi Zhang, Jiantao Jiao, Song Mei, Michael I Jordan, and Stuart Russell. How do llms perform two-hop reasoning in context? arXiv preprint arXiv:2502.13913, 2025
Pith/arXiv arXiv 2025
-
[15]
Scalable complexity control facilitates reasoning ability of llms
Liangkai Hang, Junjie Yao, Zhiwei Bai, Tianyi Chen, Yang Chen, Rongjie Diao, Hezhou Li, Pengxiao Lin, Zhiwei Wang, Cheng Xu, et al. Scalable complexity control facilitates reasoning ability of llms. arXiv preprint arXiv:2505.23013, 2025
Pith/arXiv arXiv 2025
-
[16]
Generalization or hallucination? understanding out-of-context reasoning in transformers
Yixiao Huang, Hanlin Zhu, Tianyu Guo, Jiantao Jiao, Somayeh Sojoudi, Michael I Jordan, Stuart Russell, and Song Mei. Generalization or hallucination? understanding out-of-context reasoning in transformers. arXiv preprint arXiv:2506.10887, 2025
arXiv 2025
-
[17]
Block-recurrent transformers
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block-recurrent transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=uloenYmLCAo
2022
-
[18]
Understanding Finetuning for Factual Knowledge Extraction from Language Models
Mehran Kazemi, Sid Mittal, and Deepak Ramachandran. Understanding finetuning for factual knowledge extraction from language models. CoRR, abs/2301.11293, 2023. URL https://doi.org/10.48550/arXiv.2301.11293
work page internal anchor Pith review Pith/arXiv arXiv doi:10.48550/arxiv.2301.11293 2023
-
[19]
Chain of history: Learning and forecasting with llms for temporal knowledge graph completion
Ruilin Luo, Tianle Gu, Haoling Li, Junzhe Li, Zicheng Lin, Jiayi Li, and Yujiu Yang. Chain of history: Learning and forecasting with llms for temporal knowledge graph completion. arXiv preprint arXiv:2401.06072, 2024
Pith/arXiv arXiv 2024
-
[20]
Are we falling in a middle-intelligence trap? an analysis and mitigation of the reversal curse
Ang Lv, Kaiyi Zhang, Shufang Xie, Quan Tu, Yuhan Chen, Ji-Rong Wen, and Rui Yan. Are we falling in a middle-intelligence trap? an analysis and mitigation of the reversal curse. arXiv preprint arXiv:2311.07468, 2023
Pith/arXiv arXiv 2023
-
[21]
Gradient descent maximizes the margin of homogeneous neural networks
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890, 2019
Pith/arXiv arXiv 1906
-
[22]
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=jNM4imlHZv
2024
-
[23]
In-context learning and induction heads
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022
Pith/arXiv arXiv 2022
-
[24]
Measuring and narrowing the compositionality gap in language models
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. arXiv preprint arXiv:2210.03350, 2022
Pith/arXiv arXiv 2022
-
[25]
An investigation of llms' inefficacy in understanding converse relations
Chengwen Qi, Bowen Li, Binyuan Hui, Bailin Wang, Jinyang Li, Jinwang Wu, and Yuanjun Laili. An investigation of llms' inefficacy in understanding converse relations. arXiv preprint arXiv:2310.05163, 2023
Pith/arXiv arXiv 2023
-
[26]
The mechanistic basis of data dependence and abrupt learning in an in-context classification task
Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=aN4Jf6Cx69
2024
-
[27]
Jo MA : Demystifying multilayer transformers via joint dynamics of MLP and attention
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Jo MA : Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=LbJqRGNYCf
2024
-
[28]
Miles Turpin, Julian Michael, Ethan Perez, and Samuel R. Bowman. Language models don't always say what they think: Unfaithful explanations in chain-of-thought prompting. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=bzs4uPLXvi
2023
-
[29]
Grokking of implicit reasoning in transformers: A mechanistic journey to the edge of generalization
Boshi Wang, Xiang Yue, Yu Su, and Huan Sun. Grokking of implicit reasoning in transformers: A mechanistic journey to the edge of generalization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id=D4QgSWxiOb
2024
-
[30]
Chi, Quoc V
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In NeurIPS, 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/9d5609613524ecf4f15af0f7b31abca4-Abstract-Conference.html
2022
-
[31]
An explanation of in-context learning as implicit bayesian inference
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021
Pith/arXiv arXiv 2021
-
[32]
Do large language models have compositional ability? an investigation into limitations and scalability
Zhuoyan Xu, Zhenmei Shi, and Yingyu Liang. Do large language models have compositional ability? an investigation into limitations and scalability. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. URL https://openreview.net/forum?id=4XPeF0SbJs
2024
-
[33]
Do large language models latently perform multi-hop reasoning? arXiv preprint arXiv:2402.16837, 2024
Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. Do large language models latently perform multi-hop reasoning? arXiv preprint arXiv:2402.16837, 2024
Pith/arXiv arXiv 2024
-
[34]
An analysis for reasoning bias of language models with small initialization, 2025
Junjie Yao, Zhongwang Zhang, and Zhi-Qin John Xu. An analysis for reasoning bias of language models with small initialization, 2025. URL https://arxiv.org/abs/2502.04375
Pith/arXiv arXiv 2025
-
[35]
How does transformer learn implicit reasoning? 2025
Jiaran Ye, Zijun Yao, Zhidian Huang, Liangming Pan, Jinxin Liu, Yushi Bai, Amy Xin, Liu Weichuan, Xiaoyin Che, and Lei Hou. How does transformer learn implicit reasoning? 2025
2025
-
[36]
Do llms really think step-by-step in implicit reasoning? 2024
Yijiong Yu. Do llms really think step-by-step in implicit reasoning? 2024
2024
-
[37]
Yedi Zhang, Aaditya K. Singh, Peter E. Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention, 2025 a . URL https://arxiv.org/abs/2501.16265
Pith/arXiv arXiv 2025
-
[38]
Initialization is critical to whether transformers fit composite functions by reasoning or memorizing
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Initialization is critical to whether transformers fit composite functions by reasoning or memorizing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id=YOBGdVaYTS
2024
-
[39]
Complexity control facilitates reasoning-based compositional generalization in transformers
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Complexity control facilitates reasoning-based compositional generalization in transformers. arXiv preprint arXiv:2501.08537, 2025 b
Pith/arXiv arXiv 2025
-
[40]
Towards a theoretical understanding of the'reversal curse'via training dynamics
Hanlin Zhu, Baihe Huang, Shaolun Zhang, Michael Jordan, Jiantao Jiao, Yuandong Tian, and Stuart J Russell. Towards a theoretical understanding of the'reversal curse'via training dynamics. Advances in Neural Information Processing Systems, 37: 0 90473--90513, 2024
2024
-
[41]
write newline
" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION format.date year duplicate empty "emp...
-
[42]
@esa (Ref
\@ifxundefined[1] #1\@undefined \@firstoftwo \@secondoftwo \@ifnum[1] #1 \@firstoftwo \@secondoftwo \@ifx[1] #1 \@firstoftwo \@secondoftwo [2] @ #1 \@temptokena #2 #1 @ \@temptokena \@ifclassloaded agu2001 natbib The agu2001 class already includes natbib coding, so you should not add it explicitly Type <Return> for now, but then later remove the command n...
-
[43]
\@lbibitem[] @bibitem@first@sw\@secondoftwo \@lbibitem[#1]#2 \@extra@b@citeb \@ifundefined br@#2\@extra@b@citeb \@namedef br@#2 \@nameuse br@#2\@extra@b@citeb \@ifundefined b@#2\@extra@b@citeb @num @parse #2 @tmp #1 NAT@b@open@#2 NAT@b@shut@#2 \@ifnum @merge>\@ne @bibitem@first@sw \@firstoftwo \@ifundefined NAT@b*@#2 \@firstoftwo @num @NAT@ctr \@secondoft...
-
[44]
c6lؐ >Fcr>]F e+TqIJ t 8:ﻸaÆ gÆ 5jw 5 oa С BȢߚE QQQoĔ k/ԥuQ5h@UT' DǏ7i+VM^yw|
@open @close @open @close and [1] URL: #1 \@ifundefined chapter * \@mkboth \@ifxundefined @sectionbib * \@mkboth * \@mkboth\@gobbletwo \@ifclassloaded amsart * \@ifclassloaded amsbook * \@ifxundefined @heading @heading NAT@ctr thebibliography [1] @ \@biblabel @NAT@ctr \@bibsetup #1 @NAT@ctr @ @openbib .11em \@plus.33em \@minus.07em 4000 4000 `\.\@m @bibit...
arXiv 2046
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.