DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning
2026-07-01 • Computation and Language
Computation and LanguageArtificial IntelligenceMachine Learning
AI summaryⓘ
The authors study how large language models handle multi-step reasoning in one go, focusing on a task called two-hop reasoning. They find that standard Transformers forget important facts learned earlier when trying to combine information in later steps. Using Looped Transformers helps by reusing memory, but they still don't perform perfectly. The authors identify a problem with how information is represented and introduce DiscoLoop, a model that carries two types of information during its repeated passes, achieving near-perfect accuracy faster and working well on real-world data.
Large Language ModelsChain-of-ThoughtTwo-hop ReasoningTransformersLooped TransformersHidden StateToken EmbeddingRepresentation AlignmentRecurrenceLanguage Modeling
Authors
Hengyu Fu, Tianyu Guo, Zixuan Wang, Hanlin Zhu, Jason D. Lee, Jiantao Jiao, Stuart Russell, Song Mei
Abstract
Large language models achieve strong performance on many reasoning tasks when allowed to externalize intermediate steps as Chain-of-Thought (CoT). However, many questions require the model to internalize the multi-step reasoning within a single forward pass before generating the answer. We study this challenge through two-hop reasoning, a representative task where the model must compose multiple pieces of parametric knowledge within a single forward pass. Standard non-recurrent Transformers suffer from a depth-local storage problem: facts learned in earlier layers are unavailable where second-hop retrieval happens. We found that Looped Transformers mitigate this issue by reusing the same memory, but still generalize imperfectly. We show that the remaining bottleneck is representational. In the two-hop reasoning task, the first loop often makes the correct bridge entity nearly perfectly decodable, yet the corresponding hidden state remains poorly aligned with the bridge token embedding. Surprisingly, an easy training-free realignment intervention nearly closes the generalization gap. Building upon this insight, we propose DiscoLoop, a looping architecture whose recurrence carries both a discrete embedding channel and a continuous hidden-state channel. DiscoLoop achieves near-perfect accuracy with substantially fewer training steps across symbolic and synthetic-language multi-hop reasoning tasks. When applied to real-world pretraining, DiscoLoop attains lower training loss and stronger benchmark performance than looped-transformer baselines, suggesting that the mixed-channel design transfers to practical language modeling.