一句话总结

来自微软研究院、苏黎世联邦理工学院、KRAFTON和Ludo Robotics的研究人员提出了LOTUS，一种循环填充Transformer，它并行处理K个潜在块，进行R次迭代，并在黄金思维链步骤token上使用交叉熵监督，在$3B$3B规模上弥合了潜在推理与显式推理之间的差距，同时将思考阶段延迟降低了$2.5 \times -6.9 \times$2.5×−6.9×，并产生了可解释的、与CoT对齐的潜在表示。

核心贡献

• LOTUS是首个在3B规模上缩小与显式思维链性能差距的潜在思维链方法，在Llama-3.2-3B-Instruct上针对GSM8K实现了这一点。
• 与显式思维链相比，它将思考阶段延迟降低了2.5倍至6.9倍，并在数学基准测试上超越了域外平均分。
• 循环后的潜在状态是可解释的：通过基础语言模型头对其进行投影，可以恢复黄金推理步骤和替代有效链，并且消融实验证实了循环骨干网络和并行黄金token监督是必不可少的。

引言

通过思维链推理来扩展推理计算已成为提高语言模型性能的主要方法，但顺序生成每个推理token会带来高延迟。潜在推理提供了一种更高效的替代方案，它通过在连续隐藏状态而非离散token中执行中间计算，将多个步骤浓缩为更少的模型评估。小规模模型受益于这种方法，然而，一旦超过大约1B参数，现有的潜在方法在数学任务上就无法匹配显式CoT的准确性，并且这种差距随着模型规模的增大而扩大。作者指出了两种常见的失败原因：先前的方法以自回归方式生成潜在思维，这保留了CoT昂贵的顺序瓶颈；并且它们缺乏直接的、位置对齐的CoT监督，导致潜在表示漂移并破坏训练稳定性。为了克服这两个问题，作者提出了LOTUS，一种循环填充Transformer，它在少量并行传递中细化固定数量的潜在token块，并通过基础模型自身的语言建模头，直接根据黄金CoT token监督这些潜在表示。这种并行设计消除了顺序依赖，而真实CoT对齐则防止了潜在漂移。LOTUS是首个在3B规模上，在GSM8K上缩小与显式CoT域内准确性差距的潜在推理方法，在域外数学基准测试上超越了CoT，并在冗长的自然语言推理上将思考阶段延迟降低了高达6.9倍。

方法

为了将标准语言模型转变为高效的潜在推理器，作者提出了LOTUS（对潜在表示进行并行监督的循环Transformer）。其核心设计通过两个主要组成部分，将潜在计算直接建立在黄金思维链token之上：一个由循环Transformer处理的填充潜在前缀，以及对精确黄金CoT token的并行交叉熵监督。

输入构建始于一个填充的潜在前缀。对于一个问题$Q$Q，$S$S个CoT步骤$T_1, \ldots, T_S$T1​,…,TS​，以及一个答案$A$A，模型构建一个输入序列，包含$Q$Q、一个可学习的思维开始token $\langle \mathrm{BoT} \rangle$⟨BoT⟩、$K$K个块，每个块包含$c$c个可学习的潜在token $\langle \mathrm{lat} \rangle$⟨lat⟩、一个思维结束token $\langle \mathrm{EoT} \rangle$⟨EoT⟩，以及答案$A$A。块预算$K$K和每块宽度$c$c是固定的超参数，确保潜在区域能够跨样本与CoT步骤对齐。

在循环潜在计算期间，基础因果语言模型$f_{\theta}$fθ​首先处理前缀$[Q, \langle \mathrm{BoT} \rangle]$[Q,⟨BoT⟩]以填充键值缓存$\mathcal{C}_{\mathrm{pre}}$Cpre​。然后，模型在潜在嵌入$\boldsymbol{E} \in \mathbb{R}^{Kc \times d}$E∈RKc×d上迭代$R$R次。迭代$t$t后潜在位置的隐藏状态计算如下：

$$\begin{array}{l} \boldsymbol{h}^{(0)} = f_{\boldsymbol{\theta}}(\boldsymbol{E} \mid \mathcal{C}_{\mathrm{pre}}), \\ \boldsymbol{h}^{(t)} = f_{\boldsymbol{\theta}}\left(\boldsymbol{E} + \boldsymbol{h}^{(t-1)} \mid \mathcal{C}_{\text{pre}}\right), \quad t = 1, \dots, R. \end{array}$$h(0)=fθ​(E∣Cpre​),h(t)=fθ​(E+h(t−1)∣Cpre​),t=1,…,R.​

这种有限展开的循环允许模型在填充的工作空间上联合细化潜在嵌入，同时关注已缓存的问题上下文。

如下图所示：

训练目标结合了步骤对齐的CoT监督和答案监督。在$R$R次迭代之后，步骤CoT监督损失$\mathcal{L}_{\text{step}}$Lstep​通过基础模型的语言模型头$f_{\text{head}}$fhead​，使用单次批量交叉熵，将网格中的每个潜在位置$(i,j)$(i,j)与对应的CoT步骤token $T_{i,j}$Ti,j​对齐：

$$\mathcal{L}_{\text{step}} = \frac{1}{N_{\text{step}}} \sum_{i=1}^{K} \sum_{j=1}^{c} \mathrm{CE}\left(f_{\text{head}}\left(\boldsymbol{h}_{i,j}^{(R)}\right), T_{i,j}\right),$$Lstep​=Nstep​1​i=1∑K​j=1∑c​CE(fhead​(hi,j(R)​),Ti,j​),

其中$N_{\text{step}}$Nstep​是被监督的CoT token总数。答案监督损失$\mathcal{L}_{\text{ans}}$Lans​在一个单独的最终前向传递中计算，该传递重用前缀缓存并插入循环后的潜在隐藏状态$\boldsymbol{h}^{(R)}$h(R)。模型自回归地预测答案后缀$A$A：

$$\mathcal{L}_{\text{ans}} = \frac{1}{|A|} \sum_{m=0}^{|A|-1} \mathrm{CE}(f_{\text{head}}(\boldsymbol{z}_m), A_{m+1}),$$Lans​=∣A∣1​m=0∑∣A∣−1​CE(fhead​(zm​),Am+1​),

其中$\boldsymbol{z}_m$zm​表示答案后缀位置的隐藏状态。完整目标是$\mathcal{L} = \mathcal{L}_{\text{ans}} + \lambda_{\text{step}} \mathcal{L}_{\text{step}}$L=Lans​+λstep​Lstep​。这种并行链似然方法确保了每个位置上黄金token的覆盖，而答案损失则为联合计算的隐藏状态提供了全局选择压力。

为了探索自回归链似然，作者还引入了LOTUS-aux，它通过一个辅助解码器$g_{\phi}$gϕ​而不是基础语言模型头来传递CoT监督。

辅助解码器变体如下图所示：

在LOTUS-aux中，在每次循环迭代$t$t时，辅助解码器读取潜在块$\boldsymbol{h}_{t}^{(t)}$ht(t)​作为前缀，并在教师强制下预测黄金CoT步骤$T_t$Tt​。输出隐藏状态$\widetilde{\boldsymbol{z}}^{(t)}$z(t)由下式生成：

$$\widetilde{\boldsymbol{z}}^{(t)} = g_{\phi}\left([\boldsymbol{h}_{t,c}^{(t)}, T_{t,1}, \dots, T_{t,|T_t|-1}] \mid \boldsymbol{h}_{t,1}^{(t)}, \dots, \boldsymbol{h}_{t,c-1}^{(t)}\right).$$z(t)=gϕ​([ht,c(t)​,Tt,1​,…,Tt,∣Tt​∣−1​]∣ht,1(t)​,…,ht,c−1(t)​).

此变体的步骤损失计算如下：

$$\mathcal{L}_{\text{step}}^{\text{aux}} = \frac{1}{N_{\text{step}}} \sum_{t=1}^{K} \sum_{m=0}^{|T_t|-1} \mathrm{CE}\left(g_{\text{head}}(\tilde{\boldsymbol{z}}_m^{(t)}), T_{t,m+1}\right).$$Lstepaux​=Nstep​1​t=1∑K​m=0∑∣Tt​∣−1​CE(ghead​(z~m(t)​),Tt,m+1​).

辅助解码器仅在训练期间用于自回归地建模思维链，而LOTUS和LOTUS-aux的推理仅依赖于循环后的潜在表示来解码最终答案，从而保留了并行潜在推理的效率优势。

实验

该评估使用GPT-2和LLAMA骨干网络，在GSM8K和域外数学基准测试上，从准确性、效率、设计选择和可解释性方面研究了潜在推理模型LOTUS。LOTUS在大规模上匹配了显式思维链的准确性，而先前的潜在方法则落后，并且它通过将推理压缩到并行潜在迭代中，在思考阶段实现了2.5倍到6.9倍的加速。消融实验表明，循环骨干网络和潜在监督都是有益的，准确性在适中的潜在宽度和循环深度下达到饱和，并且推理时调整宽度或深度无需重新训练即可保持稳健。对潜在表示的分析证实，它们携带了可读的思维链信号，在未见过的有效推理链上分配了非平凡的质量，并依赖于步骤监督和答案监督的互补作用来形成连贯的推理状态。

LOTUS在所有测试规模上几乎匹配了显式思维链的域内准确性，并在最大规模上领先域外平均分，而像CODI+SIM-CoT和KaVa这样的先前潜在方法，其准确性差距随着模型规模的增大而扩大。它实现了比显式CoT快2.5倍的思考阶段，将推理压缩到并行潜在迭代中，并且在自然语言轨迹上加速比增长到6.9倍。LOTUS在每个规模上的GSM8K得分都保持在显式CoT的约1.5分以内，并在3B规模上超越了显式CoT的域外平均分，而CODI+SIM-CoT在3B规模上落后9.2分。LOTUS的思考阶段比显式CoT快2.5倍，在自然语言CoT轨迹上加速比增加到6.9倍，同时保持与显式CoT相当的准确性。

LOTUS通过将推理压缩到并行潜在迭代中，实现了比显式思维链显著更低的思考阶段延迟，同时保持了有竞争力的准确性。思考阶段主导了整体延迟差异，LOTUS在此阶段比显式CoT快2.5倍，比SIM-CoT快1.2倍。CODI由于每步仅解码单个潜在表示，在思考阶段更快，但这是以降低准确性为代价的。LOTUS的思考阶段延迟为133.0毫秒，而显式CoT为338.8毫秒，SIM-CoT为162.7毫秒。CODI实现了最低的思考延迟，为88.2毫秒，但相对于LOTUS和显式CoT牺牲了准确性。每个样本的总推理延迟，LOTUS为181.2毫秒，显式CoT为384.2毫秒，加速超过2倍。查询预填充和答案阶段在不同方法中贡献的延迟相对较小且相似，使得思考阶段成为主要的区分因素。

LOTUS在使用3B参数模型的自然语言GSM8K压力测试中，匹配了显式思维链的准确性，同时将思考阶段延迟降低了6.9倍。它远超其他潜在推理方法，比次优基线高出超过8个百分点。LOTUS将思考阶段延迟从963.6毫秒降至140.8毫秒（6.9倍加速），同时保持在显式思维链准确性的0.3分以内。它优于所有其他潜在推理基线，最佳先前方法KaVa落后8.1个百分点（60.0% vs. 68.13%）。

通过主语言模型头对最终潜在状态进行循环后监督，实现了最高的测试准确性，达到70.0%，优于逐次迭代监督和辅助解码器变体。任何形式的LOTUS潜在监督都比仅使用答案损失的基线（63.3%）和仅使用CODI的模型（64.4%）带来了巨大的提升。通过主语言模型头对h(R)进行循环后监督达到了70.0%的准确性，高于逐次迭代监督的68.2%。通过辅助解码器进行逐次迭代监督的LOTUS-aux达到了69.9%，几乎匹配了最佳配置。仅使用答案损失（63.3%）和仅使用CODI（64.4%）的基线都远低于任何LOTUS变体。

增加循环细化迭代次数R持续提高了GSM8K上的准确性，从R=2到R=5有显著提升，并在R=6时接近饱和。模型受益于在读出之前对潜在前缀进行更深层次的顺序细化。准确性从R=2时的14.6%急剧上升到R=5时的68.1%，然后在R=6时稳定在70.0%附近。提升持续到R=5，然后接近饱和，表明超过五次后，额外的细化迭代带来的回报递减。

LOTUS是一种将思维链压缩为并行迭代细化的潜在推理方法，在GSM8K和域外任务上，针对显式思维链和先前的潜在方法进行了评估。它匹配了显式思维链的准确性，同时实现了高达6.9倍的更快思考阶段延迟，并且它大幅优于早期的潜在方法，这些方法的准确性差距随着规模扩大而扩大。消融实验表明，对最终潜在状态进行循环后监督以及增加细化迭代次数（超过五次后回报递减）对其性能至关重要。