π€ AI Summary
Large language models often fail to generalize on simple two-hop reasoning tasks due to an inability to effectively compose their learned single-hop knowledge. This work trains Transformers from scratch in a controlled symbolic environment and, through mechanistic interpretability analysis, reveals for the first time that consistency in intermediate representations is crucial for successful two-hop generalization, while mismatches between upper- and lower-layer representations underlie out-of-distribution failures. Building on this insight, the authors propose a recursive training strategy that encourages reuse of reasoning circuits, substantially improving out-of-distribution generalization on two-hop queries. The study further provides a systematic account of the internal mechanisms driving both success and failure in compositional reasoning.
π 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.