生成式递归推理
传统神经推理系统在实现扩展计算时通常采用自回归序列扩展,而递归推理模型(Recursive Reasoning Models, RRMs) 通过迭代潜在状态精炼与共享转换函数提供了一种替代方案。然而现有 RRMs 大多是确定性的,仅沿单一潜在轨迹收敛至单一预测,缺乏处理多假设和替代策略的能力。 本文提出生成式递归推理模型(Generative Recursive reAsoning Models, GRAM),将递归潜在推理转化为概率多轨迹计算。GRAM 将推理建模为随机潜在轨迹,支持多假设与并行推理路径,并通过递归深度和平行轨迹采样实现推理时扩展。该模型是一个潜在变量生成模型:给定输入时支持条件推理 $p\theta(y \mid x)$,输入缺失或固定时可进行无条件生成 $p\theta(x)$。训练采用摊销变分推理(amortized variational inference),使模型能够高效学习。 在结构化推理与多解约束满足任务上,GRAM 显著优于确定性循环与递归基线,并展现出高质量的无条件生成能力。实验验证了概率多轨迹计算在提升推理鲁棒性和多样性方面的优势。
论文精读
TL;DR GRAM 把递归推理从确定性单轨迹转为概率多轨迹计算,通过随机潜在轨迹和并行采样支持多假设推理与无条件生成,显著提升结构化推理性能。
问题
递归推理模型 (RRM) 通过共享的转移函数在隐空间中进行迭代精炼,为长时间推理提供了自回归生成之外的扩展计算范式。然而,现有 RRM(如 Looped Transformer)本质上是 确定性的,只能沿单一隐轨迹收敛到唯一解,无法表示推理过程中的不确定性或多解可能性。
这一局限性在需要探索多个候选解的任务中尤为突出:
- 对 约束满足(如数独、图着色)和 规划 问题,推理往往需要回溯或并行考虑多种假设;
- 确定性模型只能依赖深度扩展(更多迭代步),但单条轨迹的探索能力有限,且无法通过并行采样来提升性能;
- 训练目标(如重构或分类损失)不提供对隐状态分布的建模,使模型难以捕获 多模态后验,也无法作为生成模型进行无条件采样。
将隐推理过程转化为 概率多轨迹计算 是一个重要技术挑战:它要求模型既能维持迭代推理的高效性,又要引入可学习的随机性来生成多样化的推理路径,并在训练中有效地近似后验分布。这直接关系到 推理时扩展 能力的本质提升——从单纯的深度堆叠走向“深度+并行采样”的双维度 scaling,同时赋予模型从数据中学习联合分布的能力,进而支撑 条件推理 p_θ(y|x) 与 无条件生成 p_θ(x) 的统一。
这一思路类似于 大语言模型 利用温度采样或束搜索来生成多个候选回答,但在隐式推理器的内部循环中实现,可避免每一步都展开显式的 token 序列,提升计算效率并直接针对推理过程建模。
核心洞察
- 生成式递归推理将推理过程建模为随机潜在轨迹,允许并行采样多条推理路径,打破了传统递归推理模型收敛到单一预测的限制。GRAM 通过**概率性多轨迹计算**(probabilistic multi-trajectory computation)在推理时同时扩展递归深度和轨迹宽度,显著提升了对多解约束满足和结构化推理任务的求解能力与覆盖率。相比确定性 RRM,它引入了推理的不确定性建模,使系统能生成多种合理的候选解而非仅输出一个最优解,这对于需要探索多种可能性的规划与搜索问题尤为关键。
- GRAM 将条件推理与无条件生成统一于同一个潜在变量生成模型中,通过 p_θ(y|x) 和 p_θ(x) 分别支持解题与数据生成两种模式。这种一体化设计使得递归推理模型不仅能完成推理任务,还能成为结构化的生成器,如无条件生成 Sudoku 谜题或图着色实例。在训练与推理中,它利用**摊销变分推理**(amortized variational inference)高效学习随机递归过程的近似后验,并通过**潜在过程奖励模型**(LPRM)等机制稳定训练,为构建兼具推理与生成能力的通用神经推理架构提供了可复用的工程范式。
方法
输入与任务设定
GRAM 支持两种模式:条件推理 p_θ(y|x)(给定输入 x 预测目标 y)和无条件生成 p_θ(x)(无输入或输入固定)。输入 x 可以是图像、结构化问题(如 Sudoku 谜面)或约束描述;输出 y 为离散解(如类别、序列)。
核心架构
整个框架是一个隐变量生成模型,由三个关键模块串联:
- 编码器:将输入
x(若为图像则用 ViT 分块编码)映射为初始隐状态z_0,并推断后验分布的参数(均值、方差),用于采样初始隐变量。 - 递归核心(Transition Function):一个共享权重的循环模块(如 Gated Recurrent Unit 或 Transformer 层),在每步
t接收上一步隐变量z_{t-1},输出新的隐状态提议h_t,再通过随机采样(重参数化)产生z_t,形成多轨迹随机隐路径。递归深度T可自适应调整(Adaptive Computation Time)。 - 解码器:将最终隐状态
z_T映射为输出分布(条件模式下为p_θ(y|z_T),生成模式下可解码出完整数据x)。
训练方法
采用摊销变分推断最大化证据下界(ELBO):
- 重建项:解码器生成目标
y(或x)的对数似然; - KL 散度项:使后验分布
q(z_{1:T}|x,y)接近先验(无条件生成时先验为标准高斯); - 辅助目标:针对离散输出,引入一个潜在过程奖励模型(Latent Process Reward Model, LPRM) 作为替代训练信号,缓解梯度稀疏问题。训练时教师强制(teacher forcing)提供完整序列
y以加速收敛。
推理时扩展
推理时可通过两种正交方式提升性能:
- 深度扩展:增大递归步数
T,允许更长的内部迭代思考; - 并行采样:独立运行
K条隐轨迹(不同随机种子),聚合所有轨迹的预测(如投票或均值)。这种并行的多假设生成天然适合约束满足任务,能覆盖更多可行解。
与同类方法的差异
GRAM 将传统确定性递归推理(如 Deep Thinking 或 Recurrent Networks 的单轨迹收敛)重构为生成式多轨迹随机过程,使得推理不再是单一确定性映射,而是可并行探索多条解路径,同时具备了无条件生成的能力——这是先前递归推理模型(如 RRN、Deep Thinking 变体)所缺失的。
实验
实验设计
GRAM 在两类推理任务上验证:结构化难题(Sudoku-Extreme、ARC-AGI 上的 Looped TF)和多解约束满足(N-Queens、图着色)。对比对象包括确定性递归推理模型(RRM)、循环神经网络(如 LSTM/GRU)等基线。训练采用摊销变分推断,推理时通过控制递归深度和并行轨迹采样数量来缩放计算量。
关键发现
- GRAM 在所有结构化推理任务上超越确定性基线,尤其在 Sudoku-Extreme 上优势明显,参数扩展也无法使循环模型追上。
- 在 N-Queens 和图着色等多解问题中,GRAM 不仅给出正确解答,还能通过多条随机潜在轨迹覆盖不同解,解覆盖率显著优于单一轨迹方法。
- 无条件生成实验表明:固定或缺失输入下,GRAM 可从 p_θ(x) 直接生成合理样本(如有效数独盘面),证明其内部表征了任务的结构先验。
- 推理时并行采样可稳定提升性能,且与数据增强互补。
基线对比解读
GRAM 的核心差异是将推理定义为潜在变量的概率生成过程,而非确定性映射。这使模型能同时探索多个假设路径,避免早期错误固化。相比之下,RRM 或循环网络一旦沿单条轨迹收敛,便无法重新审视其他可能性。该设计在需要探索的解空间(如约束满足)中带来本质优势——GRAM 的并行采样相当于隐式搜索,增加了有效推理计算量,而确定性模型仅增加深度无法改变单轨迹瓶颈。这一思路为构建可扩展、多样化的神经推理系统提供了新范式:将推理计算表达为潜在轨迹上的蒙特卡洛采样,在固定参数下通过增加采样数提升性能。
行业影响
落地场景
GRAM 的 概率多轨迹推理 能力直击需多步探索与备选方案生成的工业场景。例如自动化定理证明、复杂排程与规划、程序合成、分子构象生成等任务,要求模型不只给出单一答案,而是产出多样化的有效解。在创意设计辅助(如平面布局、电路布线)或教育谜题生成中,GRAM 的无条件生成功能可直接制造符合约束的新问题实例,降低人工设计成本。
商业价值
推理时间缩放(调节递归深度与并行轨迹数)赋予 GRAM 灵活的精度-成本取舍,允许按业务需求分配推理预算。相比确定性推理器,其多假设采样大幅降低单点失败率,提升复杂任务的一次性正确率,从而减少人工校验开销。在多约束推荐或配置器场景中,提供多个可行方案能显著改善用户体验,形成产品差异化。长期看,将推理从自回归序列预测解耦为潜在空间迭代,有望突破现有模型的推理上限,解锁高附加值应用。
与现有产品/工作流的接口
GRAM 可作为可微分的神经推理模块,通过输入/输出标准化接口(如 token 序列或图结构)集成进现有 AI 管线。在大语言模型(LLM)架构中,可将其用做即插即用的推理规划器,专门处理需严格约束解析的子任务,与检索增强生成(RAG)或工具调用流水线互补。其潜在轨迹具备一定可解释性,便于调试与信任构建。训练采用摊销变分推断,与端到端训练框架自然兼容。
具体落地 use case:
- 电商复杂约束选品助手:当用户提出多维约束(如“价格<1000,重量<1.5kg,续航>10h,屏幕>13寸”)且无法同时满足时,GRAM 可在潜在空间生成多种部分符合的方案并按概率排序,帮助用户在不可兼得的指标间权衡,而传统过滤器只能返回空结果。
- 药物分子生成:给定靶点蛋白的结合口袋,GRAM 可并行采样多种满足化学键合与空间位阻约束的小分子骨架,加速苗头化合物筛选阶段,提供多样化的先导化合物选项。
局限
- **推理时计算开销较大**:GRAM 通过多轨迹并行采样实现推理时扩展,每条轨迹都需要运行完整的递归迭代过程,导致计算量随轨迹数线性增长。尽管论文通过 Adaptive Computation Time 控制深度,但在实际部署中,尤其对实时性要求高的任务,并行采样的计算成本可能难以承受,且采样数量与性能提升的关系并不总是成正比,需要精细的预算权衡。
- **任务复杂度有限,泛化性待验证**:实验集中在 ARC-AGI 等 puzzle 类任务和 N-Queens 等约束满足问题,这些任务虽能体现结构化推理,但规模较小、模式相对固定。对于大规模真实场景(如代码生成、数学证明、复杂规划),GRAM 能否保持优势并稳定训练尚不明确,其概率化推理的收益可能被环境噪声淹没,需要更多样的 benchmark 验证。
- **训练复杂度与先验依赖**:GRAM 使用 amortized variational inference 训练,涉及对潜在轨迹的近似后验推断,需要仔细设计网络结构、先验分布和损失函数。论文中展示了训练稳定性问题(如 Sudoku-Extreme 任务上参数缩放不提升),暗示其训练对超参数敏感,且无条件生成时仍需条件训练的 checkpoint 做初始化,表明训练范式仍有改进空间。