提出 DiscoLoop 架构,用双通道循环提升多跳推理准确率。
DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning

- 采用离散嵌入与连续隐藏状态双通道循环机制
- 在多跳推理任务中接近完美准确率,训练步数更少
- 无需额外训练即可显著提升模型泛化能力,适合实际语言建模
大语言模型在允许外部化中间推理步骤时表现优异,但许多问题需在单次前向传播中完成多步推理。本文研究两跳推理任务,发现标准非循环 Transformer 存在深度局部存储问题:早期层学习的知识在第二跳检索时不可用。环形 Transformer 虽缓解此问题但仍存在泛化不足。我们发现瓶颈在于表征不匹配:第一轮循环已近乎完美解码桥接实体,但对应隐藏状态与桥接标记嵌入对齐不佳。令人惊讶的是,仅通过一次无训练的对齐干预几乎完全弥补泛化差距。基于此,我们提出 DiscoLoop,其循环结构同时携带离散嵌入通道与连续隐藏状态通道。在符号与合成语言多跳推理任务中,DiscoLoop 实现近似完美准确率且训练步数显著减少。应用于真实世界预训练时,DiscoLoop 达到更低训练损失和更强基准表现,表明混合通道设计可有效迁移至实际语言建模。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。