论文

MUX:经复用token的连续推理

MUX: Continuous Reasoning via Multiplexed Tokens

模型架构新型架构

摘要

语言模型用自然语言表述中间推理步骤来解决复杂问题。这有效但受计算瓶颈制约:每个推理步只传达一个子词,且大量token用于表达想法而非执行计算。我们提出MUX,一个高带宽、紧凑推理的简单方法:把离散推理蒸馏到潜空间中的连续复用token。每个潜token经训练表示一段离散推理子词的加权线性叠加(复用),该叠加按构造无损、且该段可被完整恢复(解复用)。我们证明简单的位置相关加权——如合适的几何衰减——即支持无损复用,进而防止潜坍缩引起的捷径行为。我们进一步证明复用推理能在需要搜索的问题中执行并行探索。跨四个语言模型的32个评估设置中,MUX超过强潜推理基线。消融与探针分析进一步显示,学到的潜token编码了忠实且可解释的推理。结果表明,无损叠加作为局部学习目标构成强而高效的潜连续推理的充分条件。

MUX:经复用token的连续推理:论文配图
图 1:MUX 概述。给定一个问题,语言模型会预测一系列连续的潜在推理标记。每个标记都线性投影到词汇空间,并经过训练以表示离散推理步骤的局部范围。我们从每个离散跨度构建学习目标作为 one-hot 编码的加权平均值,并使用局部 KL 散度损失训练每个潜在标记。答案是使用标准交叉熵损失进行训练的。