在推理时使用读出反馈引导循环推理器

# 使用读取反馈引导循环推理模型的推理过程
来源：https://arxiv.org/html/2608.24136
小山雅典 | 机构：东京大学研究生院 工学研究科
郑成哲 | 机构：东京大学研究生院 工学研究科
内山文哉 | 机构：东京大学研究生院 工学研究科
久保健二 | 机构：东京大学研究生院 工学研究科
林浩平 | 机构：东京大学研究生院 工学研究科
铃木正浩 | 机构：东京大学研究生院 工学研究科
松尾丰 | 机构：东京大学研究生院 工学研究科

###### 摘要

通过共享计算块反复更新潜在状态的循环模型，已成为解决复杂推理任务的强大架构。现有的推理时方法通过运行更多步骤或采样更多轨迹来扩展计算，但忽略了每条轨迹内部揭示的信息。本文表明，循环模型可以在推理时通过使用自身的读取概率来引导潜在动态，从而无需重新训练即可得到改进。我们引入了读取反馈（RoFB），这是一种测试时干预方法，它将中间预测转换为基于token的成对耦合力，并注入到潜在动态中。在数独和迷宫任务上对三种循环模型（AKOrN、ItrSA++、TRM）的评估表明，RoFB在六个模型-任务组合中的四个中带来了显著提升，达到了仅靠运行更多步骤或从多个轨迹中选择无法实现的性能，而计算成本相当或更低。这些结果表明，对潜在动态的闭环引导可以作为循环推理模型的补充推理时控制机制。

## 1引言

人类专家常常通过更长时间的思考来解决难题：数学家完善中间结构，棋手在走出一步棋之前会搜索并修改候选走法。现代推理系统越来越多地利用类似原理——在推理时增加计算可以提升性能。在大语言模型中，这一思想体现在思维链提示（Wei等人，2022 (https://arxiv.org/html/2608.24136#bib.bib1)）、自一致性（Wang等人，2022 (https://arxiv.org/html/2608.24136#bib.bib2)）和基于搜索的推理（Yao等人，2023 (https://arxiv.org/html/2608.24136#bib.bib3)）中，其中推理时计算被用于生成、评估或聚合候选推理路径。这些进展促使推理时计算成为模型规模和训练数据之外的另一个扩展维度。

另一条互补的研究路线研究循环推理模型，这类模型通过共享计算块反复更新潜在状态来解决问题。与具有固定计算深度的前馈模型不同，循环推理器通过循环更新的次数明确地暴露了推理时计算。这一原理既出现在算法设置中，也出现在语言模型设置中：循环网络可以通过迭代计算来推断所学习的算法（Bansal等人，2022 (https://arxiv.org/html/2608.24136#bib.bib4)），而循环或循环深度的Transformer则使用重复的潜在计算来增加有效深度并提升推理性能（Yang等人，2023 (https://arxiv.org/html/2608.24136#bib.bib5)；Saunshi等人，2025 (https://arxiv.org/html/2608.24136#bib.bib6)；Geiping等人，2025 (https://arxiv.org/html/2608.24136#bib.bib7)）。与此同时，为结构化推理设计的紧凑型循环推理器，如HRM（Wang等人，2025 (https://arxiv.org/html/2608.24136#bib.bib8)）、TRM（Jolicoeur-Martineau，2025 (https://arxiv.org/html/2608.24136#bib.bib9)）、URM（Gao等人，2025 (https://arxiv.org/html/2608.24136#bib.bib10)）、AKOrN（Miyato等人，2024 (https://arxiv.org/html/2608.24136#bib.bib11)）以及相关的迭代Transformer和自注意力模型（Kubo等人，2026 (https://arxiv.org/html/2608.24136#bib.bib12)）在数独、迷宫和ARC-AGI 1 & 2（Chollet，2019 (https://arxiv.org/html/2608.24136#bib.bib13)；Chollet等人，2026 (https://arxiv.org/html/2608.24136#bib.bib14)）等谜题式推理基准测试中展现出强劲性能。

尽管取得了这些进展，当前对循环推理器的推理时改进主要被动地利用其迭代结构：运行更多步动态（Wang等人，2025 (https://arxiv.org/html/2608.24136#bib.bib8)；Geiping等人，2025 (https://arxiv.org/html/2608.24136#bib.bib7)），或采样多条轨迹并根据某个标准选择一条，例如选择动态上定义的势能最低的轨迹（Miyato等人，2024 (https://arxiv.org/html/2608.24136#bib.bib11)），或选择最自信的轨迹（Kubo等人，2026 (https://arxiv.org/html/2608.24136#bib.bib12)）。尽管这些策略有效，但它们有固有的局限性：更长的轨迹可能饱和或陷入不成功的动态状态，而多轨迹投票的计算量大致随候选数量线性增长。这引出了一个自然的问题：我们能否通过主动引导每条潜在轨迹在推理时改进一个冻结的循环推理器，而不是仅仅运行更久或采样更多轨迹？

在这项工作中，我们提出了读取反馈（RoFB），一种用于循环推理模型的闭环推理时干预方法（图1 (https://arxiv.org/html/2608.24136#S1.F1)）。RoFB源于一个关于成功循环推理的简单观察：当一条轨迹接近正确解时，token在潜在空间中发展出依赖于类的聚类结构。被预测属于同一类的token会对齐，而被分配到不同类的token会分离。RoFB使用模型自身的读取概率将这一观察转化为反馈信号。在每个选定的推理步骤，它会在潜在动态中注入一个耦合项，该耦合项基于token读取概率之间的距离来鼓励token聚类，同时保持所有模型权重不变。

我们在数独和迷宫任务上评估RoFB，使用了三种具有不同架构设计的互补循环推理器：AKOrN，一个受振荡器启发的单层模型（Miyato等人，2024 (https://arxiv.org/html/2608.24136#bib.bib11)）；ItrSA++（Kubo等人，2026 (https://arxiv.org/html/2608.24136#bib.bib12)），一个在这些基准测试上表现强大的迭代自注意力模型；以及TRM（Jolicoeur-Martineau，2025 (https://arxiv.org/html/2608.24136#bib.bib9)），一个紧凑的层级递归模型。在这些模型和任务上，RoFB无需重新训练即可提升推理性能，并且其增益与基于置信度的投票（当投票适用时）互补。我们进一步分析了性能作为推理步数、轨迹数量和归一化推理计算量的函数，表明RoFB不仅提高了最终准确率，还改善了循环推理的计算量-准确率权衡。这些结果表明，冻结的循环推理器包含未被充分利用的推理时能力，可以通过对其潜在动态的闭环控制来释放。

图1说明文字：传统推理与我们的方法对比。（a）循环模型无反馈的传统推理。（b）使用RoFB的推理。RoFB基于读取概率之间的距离，向潜在动态注入反馈耦合，同时保持所有模型权重不变。
## 2预备知识

### 2.1用于推理的循环模型

我们首先给出本文考虑的循环模型的一般设置（图1 (https://arxiv.org/html/2608.24136#S1.F1)(a)）。简而言之，这些模型由三个主要组成部分构成：(I) 一个输入嵌入映射φ\varphi，(II) 一个循环模块R，和 (III) 一个读取头ψ\psi。

更具体地说，输入嵌入φ\varphi将原始输入token x\text{raw}=\{x^{(i)}\text{raw}\}_{i=1}^N\in\{1,...,V\}^N\mathbf{\{x\}}_{\text{\{raw\}}}=\{\{\mathbf{\{x\}}^{\{\(i\)\}}_{\text{\{raw\}}}\}\}_{i=1}^{\{N\}}\in\{\{1,\ldots,V\}\}^{\{N\}}映射为固定的嵌入x=\{x^{(i)}=φ(x^{(i)}\text{raw})\}_{i=1}^N\in\mathbb{R}^{N\times D}\mathbf{\{x\}}=\{\{\mathbf{\{x\}}^{\{\(i\)\}}=φ(\mathbf{\{x\}}^{\{\(i\)\}}_{\text{\{raw\}}})\}\}_{i=1}^{\{N\}}\in\mathbb{\{R\}}^{\{N\times D\}}，其中V是词表大小，N是token数量，D是嵌入维度。循环模块R在T步上迭代地对潜在状态z\mathbf{\{z\}}_{t}=\{\{\mathbf{\{z\}}^{\{\(i\)\}}_{t}\}\}_{i=1}^{\{N\}}\in(\mathbb{\{R\}}^{\{D\}})^{\{N\}}应用更新：
z^{(i)}_{t+1}=R^{(i)}(z_{t};x), t=0,1,...,T-1.\mathbf{\{z\}}^{\{\(i\)\}}_{t+1}=R^{\{\(i\)\}}(\mathbf{\{z\}}_{t};\mathbf{\{x\}}), \quad t=0,1,\ldots,T\{\-}1\}.\} (1)
这一步进更新可以是连续时间更新的离散化：\dot{z}^{(i)}_{t}=f^{(i)}(z_{t},x), t\in[0,T].\dot{\{\mathbf{\{z\}}^{\{\(i\)\}}_{t}=f^{\{\(i\)\}}(\mathbf{\{z\}}_{t};\mathbf{\{x\}}),~~t\in[0,T]\}. 每个潜在token z^{(i)}_{t}\mathbf{\{z\}}^{\{\(i\)\}}_{t}通常被限制在紧致流形M\subset\mathbb{R}^{D}\mathcal{\{M\}}\subset\mathbb{\{R\}}^{\{D\}}（例如，D维超球面）内。初始值z_{0}\in\mathbf{\{z\}}_{0}\in\mathcal{\{M\}}是从某个分布随机采样，或作为固定参数学习得到。经过T步后，读取头ψ: M\to\psi:\mathcal{\{M\}}\to\mathbb{\{R\}}^{\{C\}}，通常是一个几层的多层感知机（MLP），产生每个token的概率向量p^{(i)}=\text{softmax}(\psi(z^{(i)}_{T}))\mathbf{\{p\}}^{\{\(i\)\}}=\operatorname\{softmax\}(\psi(\mathbf{\{z\}}^{\{\(i\)\}}_{T}))，其中C是输出类别数。

具体而言，我们选择了三种代表性的循环模型用于本研究中的推理任务：AKOrN（Miyato等人，2024 (https://arxiv.org/html/2608.24136#bib.bib11)）、ItrSA++（Kubo等人，2026 (https://arxiv.org/html/2608.24136#bib.bib12)）和TRM（Jolicoeur-Martineau，2025 (https://arxiv.org/html/2608.24136#bib.bib9)）。我们基于它们在推理任务中的强劲表现以及它们多样化的架构设计（包括层级结构和潜在状态几何）做出了这一选择。关键的架构差异总结在表1 (https://arxiv.org/html/2608.24136#S2.T1)中。下文我们将详细描述这三种循环模型的具体架构。

表1：本工作中使用的模型的架构差异。该表总结了本文分析的三个循环模型（AKOrN、ItrSA++和TRM）的关键架构特征，包括它们的潜在状态几何、耦合机制和读取策略。
### 2.2AKOrN

人工Kuramoto振荡神经元（AKOrN）（Miyato等人，2024 (https://arxiv.org/html/2608.24136#bib.bib11)）是一种在推理任务（如数独、迷宫和物体识别）中表现强劲的循环模型。在AKOrN中，每个潜在token z^{(i)}_{t}\mathbf{\{z\}}_{t}^{\{\(i\)\}}位于流形(\mathbb{S}^{m-1})^{D/m}(\mathbb{\{S\}}^{\{m\-1\}})^{\{D/m\}}上，其中m是振荡器维度，并遵循类似Kuramoto振荡器（Kuramoto，1984 (https://arxiv.org/html/2608.24136#bib.bib15)）的动力学：
Δz^{(i)}_{t}=\text{Proj}_{z^{(i)}_{t}}\Big(x^{(i)}+\sum_{j}M^{(ij)}_{t}(z_{t})z^{(j)}_{t}\Big), \quad z^{(i)}_{t+1}=\Pi\left[z^{(i)}_{t}+\gamma\Delta z^{(i)}_{t}\right].\Delta\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}=\operatorname\{Proj\}_{\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}\}\Big(\{\mathbf{\{x\}}^{\{\(i\)\}}+\sum_{\{j\}}\{\mathbf{\{M\}}^{\{\(ij\)\}}_{t}\}(\{\mathbf{\{z\}}_{t}\})\{\mathbf{\{z\}}_{t}^{\{\(j\)\}}\}\Big),\quad\{\mathbf{\{z\}}_{t+1}^{\{\(i\)\}}=\Pi\left[\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}+\gamma\Delta\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}\right]\}.\} (2)
其中，M^{(ij)}_{t}\mathbf{\{M\}}^{\{\(ij\)\}}_{t}是可学习的耦合矩阵，\text{Proj}_{z}(f)=f-\langle f,z\rangle z\operatorname\{Proj\}_{\{\mathbf{\{z\}\}}(\{\mathbf{\{f\}})=\{\mathbf{\{f\}}\-\langle\{\mathbf{\{f\}},\{\mathbf{\{z\}}\rangle\{\mathbf{\{z\}}将f\mathbf{\{f\}}投影到(\mathbb{S}^{m-1})^{D/m}(\mathbb{\{S\}}^{\{m\-1\}})^{\{D/m\}}上z点的切空间，\Pi(\cdot)=\cdot/\|\cdot\| \Pi(\cdot)=\cdot/\{\|\cdot\|\}是逐token的归一化算子，将每个token投影回乘积球面，z_{t}=\{z^{(i)}_{t}\}_{i=1}^{N}\mathbf{\{z\}}_{t}=\{\{\mathbf{\{z\}}^{\{\(i\)\}}_{t}\}\}\_{i=1}^{\{N\}}是所有潜在状态的集合，\gamma\gamma是可学习的步长。耦合矩阵M^{(ij)}_{t}\mathbf{\{M\}}^{\{\(ij\)\}}_{t}被实现为当前潜在状态的自注意力输出，允许token间可学习的全对全耦合。为了简化，我们有意省略了原始公式中存在的固有频率项Ω^{i}\boldsymbol\{\{\Omega\}\_{i}}。

在本工作中，在简要探索了多种设计选择之后，我们使用无归一化的Transformer实现了略微修改形式的AKOrN：
Δz^{(i)}_{t}=\text{Proj}_{z^{(i)}_{t}}\Big(\text{mlp}\Big[\big(z_{t}+x+\text{SelfAttn}(z_{t}+x)\big)^{(i)}\Big]\Big),\Delta\{\mathbf{\{z\}}^{\{\(i\)\}}_{t}=\operatorname\{Proj\}_{\{\mathbf{\{z\}}^{\{\(i\)\}}_{t}\}\left(\{\mathrm{\{mlp\}}\left[\bigl(\{\{\mathbf{\{z\}}_{t}+\{\mathbf{\{x\}}+\{\mathrm{\{SelfAttn\}}(\{\{\mathbf{\{z\}}_{t}+\{\{\mathbf{\{x\}}\}\bigr)^{\{\(i\)\}}\right]\right\},\} (3)
其中\text{mlp}\mathrm{\{mlp\}}是带有GELU激活的两层前馈网络，\text{SelfAttn}\mathrm{\{SelfAttn\}}是带有位置编码的多头自注意力块。

### 2.3ItrSA++

ItrSA++是Kubo等人（2026 (https://arxiv.org/html/2608.24136#bib.bib12)）提出的另一种循环模型，它在数独和迷宫中展现出更强的推理能力。与AKOrN不同，ItrSA++被设计为具有两级层次结构，其中低级过程\mathbf{z}旨在捕获快速变化的细节，而高级过程\mathbf{y}则期望捕获更慢变化的全局结构。在低级过程\mathbf{z}中，动力学更新如下：
z^{(i)}_{t+1}=
\begin{cases}
\text{RN}\Big(z^{(i)}_{t}+\text{SelfAttn}(z_{t})\Big) & \text{if } t\not\equiv 0 \pmod{L},\\
\text{RN}\Big(z^{(i)}_{t}+\text{CrossAttn}\big(\text{RN}(y_{t}),\text{RN}(x)\big)\Big) & \text{if } t\equiv 0 \pmod{L},
\end{cases}\mathbf{\{z\}}_{t+1}^{\{\(i\)\}}=\begin\{\{cases\}\operatorname\{RN\}\left(\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}+\operatorname\{SelfAttn\}(\{\mathbf{\{z\}}_{t}\)\right\}&\text\{if $t\not\equiv 0\pmod\{L\}$,\}\{\operatorname\{RN\}\left(\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}+\operatorname\{CrossAttn\}\left(\operatorname\{RN\}(\{\mathbf{\{y\}}_{t}\),\operatorname\{RN\}(\{\mathbf{\{x\}}\right\)\right\}&\text\{if $t\equiv 0\pmod\{L\}$,\}\end\{\{cases\}\} (4)
其中\text{RN}\operatorname\{RN\}表示RMS归一化。高级过程\mathbf{y}每L步更新一次：
y^{(i)}_{t+1}=
\begin{cases}
\text{RN}\Big(z^{(i)}_{t}+\text{SwiGLU}(z^{(i)}_{t})\Big) & \text{if } t\equiv 0 \pmod{L},\\
y^{(i)}_{t} & \text{otherwise}.
\end{cases}\mathbf{\{y\}}_{t+1}^{\{\(i\)\}}=\begin\{\{cases\}\operatorname\{RN\}\left(\{\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}+\operatorname\{SwiGLU\}(\{\{\mathbf{\{z\}}_{t}^{\{\(i\)\}}\)\right\}&\text\{if $t\equiv 0\pmod\{L\}$,\}\{\mathbf{\{y\}}_{t}^{\{\(i\)\}}&\text\{otherwise.\}\end\{\{cases\}\} (5)
由于RMS归一化，ItrSA++的潜在状态空间可以等同于(\mathbb{S}^{D-1})(\sqrt{D})\mathbb{\{S\}}^{\{D\-1\}\ (\{\sqrt\{D\}\})。

### 2.4微型递归模型（TRM）

TRM（Jolicoeur-Martineau，2025 (https://arxiv.org/html/2608.24136#bib.bib9)）是另一种在推理基准任务（如数独、迷宫和ARC-AGI（Chollet，2019 (https://arxiv.org/html/2608.24136#bib.bib13)；Chollet等人，2026 (https://arxiv.org/html/2608.24136#bib.bib14)）中表现强劲的循环模型。TRM被提议作为开创性层级推理模型（HRM）（Wang等人，2025 (https://arxiv.org/html/2608.24136#bib.bib8)）的简化版本，在减小模型规模和计算复杂性的同时，展现出更强大的推理能力。TRM具有两级隐藏状态层次结构，旨在...