用循环变压器让隐式推理追上显式推理,速度提升6.9倍。
Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers
- 用循环结构并行处理隐藏推理块,提升效率。
- 30亿参数下实现与显式推理相当的准确率。
- 适合需要快速推理且关注中间步骤解释性的研究者。
语言模型通常通过显式思维链(CoT)逐标记生成中间步骤进行推理。隐式CoT则在模型隐藏状态中完成多步推理,以连续表示替代解码标记,提升效率。然而,现有隐式CoT方法在超过10亿参数后性能落后于显式CoT,且差距随规模扩大。循环(或递归深度)Transformer通过复用权重增加计算深度而不增加参数,天然适配隐式推理。本文提出一种简单方案:使用循环填充Transformer并行处理K个隐式块,经R次迭代,对每个隐式位置施加交叉熵损失,监督其对应的真实推理步骤,类似显式CoT训练。该方法实现为LOTUS(Looped Transformers with parallel supervision on latents)。据我们所知,LOTUS是首个在30亿参数规模上缩小隐式与显式推理差距的方法,同时将从紧凑数学表达到自然语言的推理阶段延迟降低2.5至6.9倍。将LOTUS的后循环隐藏表示通过基础语言模型头投影,可恢复真实推理步骤,并揭示其他有效中间步骤,表明其隐藏空间具有可解释性且与思维链对齐。消融实验确认循环主干和对真实推理步骤的并行监督均至关重要。代码已公开于 https://github.com/yingfan-bot/lotus。
原文摘要 · Abstract (English)
Language models typically reason via explicit chain-of-thought (CoT), generating intermediate steps token-by-token. Latent CoT offers an alternative: it performs multi-step reasoning in the model's hidden states, replacing decoded tokens with continuous representations for greater efficiency. However, existing latent CoT methods underperform explicit CoT beyond 1B parameters, and the gap widens with scale. Looped, or recurrent-depth, Transformers, which reuse their weights to increase computation depth without adding parameters, are a natural fit for latent reasoning. We therefore ask whether looped Transformers can bridge this gap. We answer affirmatively with a simple recipe: a looped padded Transformer that processes K latent blocks in parallel for R iterations, with a cross-entropy loss on each latent position's gold CoT-step token, similar to explicit CoT supervision. We instantiate it as LOTUS (Looped Transformers with parallel supervision on latents). LOTUS is, to our knowledge, the first latent-CoT method to bridge the gap to explicit CoT at the 3B scale, while cutting thought-phase latency by 2.5x-6.9x from compact math expressions to natural language. Projecting LOTUS's post-loop latents through the base LM head recovers the gold reasoning steps and even surfaces alternative valid intermediate steps, evidence that its latent space is interpretable and CoT-aligned. Ablations confirm that both the looped backbone and the parallel supervision on gold CoT tokens are essential. Code is available at https://github.com/yingfan-bot/lotus.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。