使用循环 Transformer 弥合潜在推理和显式推理之间的差距

语言模型通常通过显式思维链 (CoT) 进行推理，逐个生成中间步骤。潜在 CoT 提供了一种替代方案：它在模型的隐藏状态中执行多步推理，用连续表示替换解码的标记，以提高效率。然而，现有的潜在 CoT 方法在 1B 参数之外的表现不如显式 CoT，并且差距随着规模的扩大而扩大。循环深度或循环深度 Transformer 重用权重来增加计算深度而不添加参数，非常适合潜在推理。因此，我们询问循环 Transformer 是否可以弥补这一差距。我们用一个简单的方法来肯定地回答：一个循环填充的 Transformer，它并行处理 K 个潜在块进行 R 次迭代，每个潜在位置的标准 CoT 步词元上都有交叉熵损失，类似于显式 CoT 监督。我们将其实例化为 LOTUS（循环 Transformer，对潜在对象进行并行监督）。据我们所知，LOTUS 是第一个在 3B 规模上弥合与显式 CoT 差距的潜在 CoT 方法，同时将紧凑数学表达式到自然语言的思维阶段延迟缩短了 2.5 倍至 6.9 倍。通过基础 LM 头投射 LOTUS 的后环潜伏可以恢复标准推理步骤，甚至表面替代有效的中间步骤，证明其潜伏空间是可解释的且 CoT 对齐的。消融证实，循环主干网和对标准 CoT 词元的并行监管都是必不可少的。代码可在 https://github.com/yingfan-bot/lotus 获取。