DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning

Implicit in-weight multi-hop reasoning---composing multiple pieces of parametric knowledge within a single forward pass---is a fundamental yet challenging task of language models. Vanilla transformers consistently fail at this task, and even looped transformers, whose recurrent structure naturally fits its iterative nature, generalize imperfectly, particularly out-of-distribution (OOD). We make two contributions toward closing this gap. First, on a symbolic two-hop reasoning task, we mechanistically diagnose why looped transformers fall short: after the first loop, the intermediate answer is already decodable from the residual stream with certain probability, yet the continuous hidden vector carrying it is noisy and geometrically misaligned with the clean discrete embedding of the intermediate answer that the next loop would ideally consume. A training-free intervention that re-aligns this representation nearly closes the OOD gap, indicating that the representational mismatch is a bottleneck. Second, building on 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 two-hop accuracy with substantially less training across symbolic and synthetic-language multi-hop reasoning tasks. When applied to real pretraining, DiscoLoop attains lower training loss and stronger performance on pretraining benchmarks than looped-transformer baselines, suggesting that the mixed-channel design transfers to practical language modeling.