← 返回 2026-05-21

生成式递归推理:把递归潜状态推理改造为概率多轨迹生成框架 Generative Recursive Reasoning

Junyeob Baek, Mingyu Jo, Minsu Kim, Mengye Ren, Yoshua Bengio, Sungjin Ahn 📅 2026-05-20 👍 31 2026-07-13 08:36
变分推断 推理时计算扩展 无条件生成 概率潜状态空间模型 潜变量生成模型 递归推理模型

把递归推理改成概率多轨迹生成模型,支持深度与宽度双重扩展。

前置知识

递归推理模型 (Recursive Reasoning Model, RRM)

通过反复施加共享转移函数来精炼持久潜在状态的模型,把推理深度与参数规模、输出长度解耦,使小模型也能完成多步内部计算,代表有 HRM、TRM、Looped Transformer、Universal Transformer。

本文的 GRAM 正是建立在 RRM 之上,把它从确定性改造为概率生成式,必须先理解 RRM 的潜状态更新机制和深度监督策略,才能读懂 GRAM 的随机性注入点。

变分推断与 ELBO

当条件似然 $p(y|x)$ 因需对潜在推理轨迹积分而不可计算时,引入变分后验 $q_\phi(\tau|x,y)$,最大化证据下界 ELBO,即重建项减去 KL 散度项,来联合训练生成参数与变分参数。

GRAM 的训练目标就是轨迹级 ELBO 的截断近似,理解变分推断才能看懂为何训练时用后验采样、推理时用先验采样,以及 KL 项如何对齐二者。

深度监督与截断梯度传播

在 $N_{sup}$ 个连续监督步上施加密集监督信号,而非只监督最终状态;同时用截断梯度传播,每步只在最后一个转移回传梯度,以换取训练长计算链的内存效率。

GRAM 沿用了 HRM/TRM 的深度监督与截断梯度策略,这是其能训练长潜推理链的关键工程技巧,读懂训练流程必须先掌握。

自适应计算时间 ACT 与宽度扩展

自适应计算时间 ACT 让每条轨迹在学到的停止深度处终止,实现深度方向自适应推理;宽度扩展则并行采样多条潜在轨迹并用价值头择优,二者共同构成推理时扩展。

GRAM 的核心卖点之一是宽度扩展,理解 ACT 才能区分深度与宽度两条推理时扩展轴,并看懂 LPRM 价值头如何从 N 条候选中择优。

研究动机

现有递归推理模型(RRM)如 HRM、TRM 和 Looped Transformer 虽然在结构化推理上展现了潜力——通过迭代精炼潜状态、深度监督和分层潜动力学把推理深度与参数规模解耦——但它们在本质上都是确定性的:给定相同输入和初始化,模型只沿单一潜轨迹演化并收敛到单一预测。这种确定性递归把所有合理推理路径坍缩到单个吸引子,无法维持不确定性、考虑替代假设或探索多种解法。论文图4右侧的实验直观印证了这一点:在 N-Queens 任务中,随着有效解数量增加,三种确定性基线的准确率急剧下降;表1显示它们的覆盖率最高仅 36.1%。在 N-Queens、图着色这类本质存在多解、或单条精炼路径可能陷入次优推理轨迹的问题上,确定性 RRM 存在结构性失效。

本文的目标是本文的核心目标是把递归潜状态推理从确定性过程改造为概率化的多轨迹计算:让推理过程本身成为一条随机潜轨迹,每次递归都从前一状态条件采样转移,从而定义出推理轨迹的分布。这样模型既能维持多个假设、探索替代解策略,又能在推理时不仅靠增加递归深度(深度扩展),还能并行采样多条轨迹(宽度扩展)来扩展计算。更进一步,作者希望同一个递归过程既能建模条件分布 $p_\theta(y|x)$ 用于推理求解,也能在输入缺失或固定时建模无条件分布 $p_\theta(x)$ 充当生成模型。目标是既保留紧凑循环模型的高效性,又获得概率生成模型的多假设推理与生成能力。

与已有工作不同的是,已有的概率循环潜状态空间模型(如用于视频预测、模型预测控制的随机循环网络)虽用随机潜转移和变分推断,但都把潜动力学解释为对时序观测的建模,而非推理计算本身;已有 RRM 则把递归视为确定性精炼。两者之间存在空白:能否把随机潜动力学重新解释为'计算',让潜转移定义可能的推理轨迹,从而在 RRM 范式内实现多假设探索与生成。GRAM 的独特切入角度正是把概率潜状态空间模型的变分视角嫁接到递归推理架构上,用随机残差扰动 $\epsilon_t$ 围绕确定性更新 $u_t$ 引入可学习的随机引导,并以变分后验 $q_\phi(\tau|x,y)$ 训练、先验 $p_\theta(\tau|x)$ 推理,桥接两条原本分离的研究线。

核心方法

GRAM 的直觉是:与其让潜状态沿一条固定路径精炼,不如在每一步把确定性更新 $u_t$ 与一个可学习的随机扰动 $\epsilon_t$ 相加,使重复计算定义出潜在推理轨迹的分布。技术上,模型先用编码器把输入 $x$ 映射为嵌入 $e_x=f_{enc}(x)$ 并在整个递归中复用;潜状态采用分层结构 $z=(h,l)$,高层 $h$ 每次转移更新一次承载抽象推理状态,低层 $l$ 在单次转移内被 $f_L$ 精炼 $K$ 次。一个监督步由 $T$ 次随机潜转移组成,解码器在每步末端产生预测并施加损失;$N_{sup}$ 个监督步首尾相接形成完整递归。训练用变分后验采样、最大化 ELBO 的截断近似;推理时用先验并行采样多条轨迹并解码。

核心创新是把递归的潜转移写成 $z_t=u_t+\epsilon_t$,其中 $u_t=f_H(h_{t-1},l_t)$ 是确定性高层更新,而 $\epsilon_t\sim\mathcal{N}(\mu_\theta(u_t),\sigma^2_\theta(u_t)I)$ 是从状态相关高斯采样的'可学习随机引导'。均值 $\mu_\theta$ 编码轨迹被引导的方向,方差 $\sigma^2_\theta$ 控制探索强度。与传统 RRM 的本质区别在于:随机性不是噪声而是结构化信号,且仅在慢变的高层引入(低层精炼保持确定性),既能维持多假设又不丢失确定性精炼的成果。训练上,GRAM 是潜变量生成模型,用变分后验 $q_\phi(\tau|x,y)$(可看到目标 $y$)采样、用先验 $p_\theta(\tau|x)$ 推理,并通过轨迹级 ELBO 对齐两者,使同一框架同时支持 $p(y|x)$ 推理与 $p(x)$ 生成。

方法步骤详情

步骤一:编码器计算 $e_x=f_{enc}(x)$。步骤二:从固定初始状态 $z_0$ 开始一次监督步,内层先做 $K$ 次低层精炼 $l_{t,k}=f_L(h_{t-1},l_{t,k-1},e_x)$。步骤三:用 $l_t$ 做高层确定性更新 $u_t=f_H(h_{t-1},l_t)$,采样 $\epsilon_t\sim\mathcal{N}(\mu_\theta(u_t),\sigma^2_\theta I)$,置 $h_t=u_t+\epsilon_t$。步骤四:重复 $T$ 次转移后解码器读 $h_T$ 得预测并用交叉熵监督。步骤五:把当前步终端状态作为下一监督步初值,递归 $N_{sup}$ 步,配合深度监督与截断梯度(仅 $z_{T-1}\to z_T$ 回传)。步骤六训练时从后验采样噪声并最大化截断 ELBO=重建项$-$KL;步骤七推理时从先验并行采样 $N$ 条轨迹,用 LPRM 价值头或多数投票选最优候选。

技术新颖性

新颖性有三层。其一,首次把递归推理形式化为潜变量生成过程,解通过对随机推理轨迹边际化得到,引入变分后验对齐训练先验。其二,提出宽度推理时扩展——通过并行采样 $N$ 条潜在轨迹并用学到的 LPRM 预测终端价值择优,突破了单纯靠深度扩展带来的延迟瓶颈。其三,消融证明收益来自变分框架而非随机性本身:把随机解码或随机初始化直接套到 TRM 上毫无帮助(N-Queens 仅 71.66%/71.82%),而移除引导($\mathcal{N}(0,\sigma^2 I)$)在 N-Queens 骤降到 50.27%、移除随机性($\mathcal{N}(\mu_\theta,0)$)直接归零,说明'可学习随机引导'是关键创新而非简单加噪。这种把概率潜状态空间视角与递归推理嫁接的设计在文献中尚属首次。

GRAM Architecture
Figure 2: GRAM Architecture

实验结果

在 Sudoku-Extreme 上 GRAM 达 97.0%,超过 HRM 的 55.0%、TRM 的 87.4%、Looped TF 的 61.25%(图3)。图4左显示 GRAM 在 $N=20$ 样本、16 次迭代下(97.0%)就超过所有确定性基线在 320 次迭代的表现(TRM 90.5%),证明宽度扩展绕开了深度扩展的延迟瓶颈。在 N-Queens(表1)GRAM 达 99.7% 准确率/90.3% 覆盖率,远超 TRM(66.8%/36.1%)和 HRM(78.7%/26.7%);图4右显示随解数增加确定性基线急剧退化而 GRAM 保持稳定。图着色上 GRAM 冲突边仅 2.7 和 3.3,远优于 AR 的 19.0 和 61.3。无条件生成上 Sudoku 有效率达 99.05%(10.9M 参数/16 步,胜过 D3PM 的 55.1M/1000 步);MNIST 的 FID 随推理步数从 84.08 单调降到 73.34(8$\to$256 步),证明递归精炼可迁移到生成。

Evaluation on N-Queens and Graph Coloring benchmarks
Table 1: Evaluation on N-Queens and Graph Coloring benchmarks
Unconditional generation results on binarized MNIST
Table 2: Unconditional generation results on binarized MNIST
Ablation study on Sudoku-Extreme and N-Queens
Table 3: Ablation study on Sudoku-Extreme and N-Queens
Performance on puzzle benchmarks
Figure 3: Performance on puzzle benchmarks
Inference-time scaling on Sudoku-Extreme and N-Queens accuracy
Figure 4: Inference-time scaling on Sudoku-Extreme and N-Queens accuracy
Unconditional Sudoku generation
Figure 5: Unconditional Sudoku generation
Visualization of the generation process and samples
Figure 6: Visualization of the generation process and samples
Qualitative examples of unconditional Sudoku generation by GRAM
Figure 7: Qualitative examples of unconditional Sudoku generation by GRAM
查看结构化数据
任务指标本文基线提升
Sudoku-Extreme 准确率 (%) GRAM 97.0% TRM 87.4% / HRM 55.0% / Looped TF 61.25% 较最强确定性基线 TRM 提升 9.6 个百分点,且 N=20 样本 16 步即超 TRM 320 步
ARC-AGI-1 准确率 (%) GRAM 66.7% TRM 34.5% / HRM 9.7% 较 TRM 提升约 32 个百分点,约翻倍
ARC-AGI-2 准确率 (%) GRAM 61.3% TRM 44.6% 提升约 16.7 个百分点
N-Queens (8x8) 准确率 / 覆盖率 (%) 99.7 / 90.3 TRM 66.8 / 36.1 ; HRM 78.7 / 26.7 覆盖率较 TRM 提升 54.2 个百分点,证明能覆盖多解
Graph Coloring (8/10 顶点) 冲突边数 (↓) 2.7 / 3.3 AR 19.0 / 61.3 递归精炼使约束满足更锐利,冲突边大幅下降
无条件 Sudoku 生成 有效率 (%) GRAM 99.05% (10.9M 参数, 16 步) D3PM 用 55.1M 参数、1000 步 以约 1/5 参数、1/60 步数达到更高有效率
binarized MNIST 无条件生成 FID (↓) / IS (↑) FID 73.34 / IS 2.04 (256 步) TRM FID 303.29 (模式坍缩); D3PM FID 74.03 修复 TRM 模式坍缩,FID 随推理步数单调改善并略胜 D3PM

局限与改进

作者明确承认:深度监督的顺序性使训练效率低于 Transformer,成为把 GRAM 扩展到大型基础模型的重要障碍;实验也只在受控的探针任务上与确定性递归基线对比,刻意未与前沿通用 LLM(o3-mini、GPT-5.2、Grok-4)做受控比较,这些大模型分数仅作难度参考。我的观察补充:任务规模小(最大约 12.6M 参数),未见在自然语言推理、代码或数学竞赛上的验证;无条件 MNIST 的 IS 仅约 2.0、FID 约 73,远低于现代扩散模型,生成质量仅是'可比拟 D3PM'而非领先;论文未报告训练成本/能耗与样本效率,变分推断的 KL 权重调度等关键超参藏在附录,复现门槛偏高。

独立分析的弱点

弱点一:训练效率瓶颈——深度监督要求顺序执行 $N_{sup}$ 步且每步回传,难以并行,改进方向是用异步/推测解码或把监督步蒸馏为更少步。弱点二:仅在小规模符号任务验证,泛化到语言、代码、多模态未证;应在大规模预训练语料上预训练 GRAM 骨干再迁移。弱点三:无条件生成质量偏低(MNIST FID 约 73),方差调度 $\sigma_\theta$ 易导致模式坍缩或过散;可借鉴扩散模型的噪声表和分类器自由引导。弱点四:宽度扩展需 $N$ 倍前向与一个 LPRM,推理成本随 $N$ 线性增长,缺少自适应早停或轨迹剪枝;可引入置信度门控。弱点五:变分后验与先验的间隙未量化,KL 项是否被低估尚不清楚,可加入事后重加权。

未来方向

作者提出的方向:把递归生成扩展到层级递归的无条件生成新范式,并解决深度监督的训练效率问题以迈向更大基础模型。基于成果可延伸:其一,把 GRAM 与 Transformer 解码头结合做自然语言/代码推理,验证概率多轨迹思路在开放域的价值;其二,研究宽度扩展与深度扩展的联合自适应调度,用强化学习学习'何时采多少条轨迹';其三,把 LPRM 升级为过程奖励并结合树搜索(如 MCTS),在潜空间做规划;其四,探索 GRAM 作为世界模型或 model-based RL 的概率动力学模块;其五,用 GRAM 的潜轨迹做可解释性分析,理解'推理路径多样性'与泛化的关系。这些都是把递归概率推理推向实用的下一步。

复现评估

论文提供了项目主页链接(ahn-ml.github.io/gram-website)以及详细的附录 B.2 训练超参、附录 C 的实验配置,方法描述(公式 4-14、图2)足够精确,理论上可复现。但代码仓库在正文未见明确开源声明,关键组件(LPRM 训练见附录 A.2、ACT 细节见附录 A.1、截断 ELBO 近似见附录 A.3)分散在附录,需仔细整合。任务规模虽小(7M-12.6M 参数、Sudoku/ARC/N-Queens 等),但 ARC-AGI 训练成本作者明确指出'过高'以致 Looped TF 在 ARC 上缺数据,单卡复现不一定可行。变分推断对超参敏感,复现者需自行调 KL 权重与方差初始化,整体难度中等偏上。