DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning
作者: Hengyu Fu, Tianyu Guo, Zixuan Wang, Hanlin Zhu, Jason D. Lee, Jiantao Jiao, Stuart Russell, Song Mei
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-07-01
备注: 16 pages, 7 figures
💡 一句话要点
提出DiscoLoop以解决多步推理中的表示瓶颈问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 多步推理 循环Transformer 离散嵌入 隐状态 自然语言处理 知识图谱 智能问答
📋 核心要点
- 现有的Transformer模型在多步推理任务中存在深度局部存储问题，导致早期层学习的事实在后续层无法使用。
- 本文提出DiscoLoop架构，通过结合离散嵌入和连续隐状态，解决了表示不足的问题，提升了模型的推理能力。
- 实验结果表明，DiscoLoop在符号和合成语言的多步推理任务中实现了接近完美的准确率，并且训练步骤显著减少。
📝 摘要（中文）
大型语言模型在许多推理任务中表现出色，尤其是在允许外部化中间步骤的情况下。然而，某些问题要求模型在单次前向传递中内化多步推理。本文研究了两步推理任务，发现标准的非递归Transformer在深度局部存储方面存在问题。尽管循环Transformer通过重用相同的内存缓解了这一问题，但仍存在表示不足的瓶颈。我们提出DiscoLoop架构，结合离散嵌入通道和连续隐状态通道，显著提高了多步推理任务的准确性和训练效率。
🔬 方法详解
问题定义：本文旨在解决多步推理任务中，标准Transformer模型在深度局部存储方面的不足，导致早期学习的知识无法在后续推理中使用。
核心思路：提出DiscoLoop架构，通过引入离散嵌入通道和连续隐状态通道的循环设计，增强模型在推理过程中的信息传递和表示能力。
技术框架：DiscoLoop的整体架构包括两个主要通道：一个用于存储离散嵌入，另一个用于维护连续的隐状态。这种设计允许模型在推理过程中有效地重用信息。
关键创新：DiscoLoop的核心创新在于其混合通道设计，使得模型在进行多步推理时能够更好地对齐隐状态与嵌入，从而显著提升了推理的准确性。
关键设计：在模型训练中，采用了特定的损失函数以优化离散嵌入和隐状态的对齐，同时调整了网络结构以支持循环机制，确保信息在推理过程中的有效传递。
🖼️ 关键图片
📊 实验亮点
实验结果显示，DiscoLoop在多步推理任务中实现了接近完美的准确率，训练步骤显著减少。与循环Transformer基线相比，DiscoLoop在真实世界预训练中表现出更低的训练损失和更强的基准性能，验证了其设计的有效性。
🎯 应用场景
DiscoLoop架构在多步推理任务中展现出强大的性能，具有广泛的应用潜力，尤其是在需要复杂推理的自然语言处理任务中。未来，该模型可能在智能问答、对话系统和知识图谱推理等领域发挥重要作用。
📄 摘要（原文）
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.