Multiplex Thinking: Reasoning via Token-wise Branch-and-Merge
作者: Yao Tang, Li Dong, Yaru Hao, Qingxiu Dong, Furu Wei, Jiatao Gu
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-01-13
备注: 21 pages. Code available at https://github.com/GMLR-Penn/Multiplex-Thinking
🔗 代码/项目: GITHUB
💡 一句话要点
提出Multiplex Thinking,通过token级分支合并进行高效推理。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 语言模型 推理 思维链 强化学习 token嵌入 软推理 数学推理
📋 核心要点
- 现有Chain-of-Thought方法推理能力强,但生成序列过长,效率较低,且难以优化。
- Multiplex Thinking通过采样和聚合token嵌入,在每一步保留多个候选推理路径,实现软推理。
- 实验表明,该方法在数学推理任务上优于CoT和强化学习基线,且生成序列更短。
📝 摘要(中文)
大型语言模型通常通过思维链(CoT)更有效地解决复杂的推理任务,但代价是产生长且低带宽的token序列。相比之下,人类通常通过保持对合理后续步骤的分布来进行软推理。受此启发,我们提出了一种随机软推理机制Multiplex Thinking,它在每个思考步骤中采样K个候选token,并将它们的嵌入聚合为单个连续的多路复用token。这保留了标准离散生成的词汇嵌入先验和采样动态,同时诱导了多路复用展开上的易处理概率分布。因此,多路复用轨迹可以直接使用在线强化学习(RL)进行优化。重要的是,Multiplex Thinking是自适应的:当模型有信心时,多路复用token几乎是离散的,并且表现得像标准的CoT;当模型不确定时,它紧凑地表示多个合理的后续步骤,而不会增加序列长度。在具有挑战性的数学推理基准测试中,从Pass@1到Pass@1024,Multiplex Thinking始终优于强大的离散CoT和RL基线,同时产生更短的序列。
🔬 方法详解
问题定义:大型语言模型在复杂推理任务中,使用Chain-of-Thought (CoT) 方法虽然能提升性能,但会生成很长的token序列,导致推理效率降低。此外,这种离散的token序列也使得直接优化变得困难。现有方法难以在推理性能和效率之间取得平衡。
核心思路:模仿人类的软推理方式,即在每一步考虑多种可能的推理路径,而不是只选择一个。通过维护一个token分布,模型可以更灵活地探索不同的推理方向,并在不确定时保留多种可能性。这种方式可以避免过早地陷入错误的推理路径,同时也能更有效地利用token序列的长度。
技术框架:Multiplex Thinking 的核心在于将多个候选token的嵌入表示聚合为一个“multiplex token”。具体来说,在每个推理步骤,模型首先采样 K 个候选token。然后,将这些token的嵌入表示进行聚合,例如通过加权平均,得到一个单一的、连续的 multiplex token。这个 multiplex token 随后被输入到模型中进行下一步的推理。整个过程可以看作是在token级别进行“分支”和“合并”。
关键创新:Multiplex Thinking 的关键创新在于将离散的token选择过程转化为连续的嵌入空间操作。这使得模型能够以一种更紧凑的方式表示多种可能的推理路径,而无需增加序列长度。此外,由于 multiplex token 是连续的,因此可以使用基于梯度的优化方法,例如强化学习,直接优化推理过程。
关键设计:在具体实现上,需要考虑如何选择候选token的数量 K,以及如何聚合这些token的嵌入表示。论文中可能使用了不同的聚合方法,例如加权平均,并可能引入了可学习的权重参数。此外,损失函数的设计也至关重要,需要鼓励模型探索不同的推理路径,并在最终结果正确时给予奖励。强化学习的奖励函数也需要仔细设计,以引导模型学习更有效的推理策略。
📊 实验亮点
实验结果表明,Multiplex Thinking 在数学推理基准测试中,从Pass@1到Pass@1024,均优于传统的Chain-of-Thought和强化学习基线。更重要的是,该方法在取得更好性能的同时,还生成了更短的序列,验证了其高效性。具体性能提升数据未知,请参考论文原文。
🎯 应用场景
Multiplex Thinking 具有广泛的应用前景,可以应用于各种需要复杂推理的场景,例如数学问题求解、代码生成、知识图谱推理等。该方法可以提高推理效率,并允许模型在不确定性下进行更鲁棒的推理。未来,该方法可以进一步扩展到其他模态,例如图像和语音,以实现更通用的推理能力。
📄 摘要(原文)
Large language models often solve complex reasoning tasks more effectively with Chain-of-Thought (CoT), but at the cost of long, low-bandwidth token sequences. Humans, by contrast, often reason softly by maintaining a distribution over plausible next steps. Motivated by this, we propose Multiplex Thinking, a stochastic soft reasoning mechanism that, at each thinking step, samples K candidate tokens and aggregates their embeddings into a single continuous multiplex token. This preserves the vocabulary embedding prior and the sampling dynamics of standard discrete generation, while inducing a tractable probability distribution over multiplex rollouts. Consequently, multiplex trajectories can be directly optimized with on-policy reinforcement learning (RL). Importantly, Multiplex Thinking is self-adaptive: when the model is confident, the multiplex token is nearly discrete and behaves like standard CoT; when it is uncertain, it compactly represents multiple plausible next steps without increasing sequence length. Across challenging math reasoning benchmarks, Multiplex Thinking consistently outperforms strong discrete CoT and RL baselines from Pass@1 through Pass@1024, while producing shorter sequences. The code and checkpoints are available at https://github.com/GMLR-Penn/Multiplex-Thinking.