论文系统地研究了循环深度Transformer(Recurrent-Depth Transformer,或称Looped Transformer)在隐式推理(单次前向传播,无显式思维链)任务上的组合泛化能力。作者通过可控的合成多跳推理实验,证明了该架构能够有效解决标准Transformer在此类任务上的两大核心缺陷:系统性泛化和深度外推。
1. 研究背景与问题
背景:标准Transformer(如GPT)的LLM虽存储大量知识,但在隐式多跳推理(如“A的B的C是谁”)中表现不佳。
根本原因:知识分布在不同层,且层间参数不共享。模型难以将分布在不同深度的知识灵活组合以应对训练中未见过的组合或更深的推理链。
本文核心问题:引入参数共享(深度循环)能否让Transformer克服上述组合泛化瓶颈?
2. 核心解决方案:循环深度Transformer
架构设计:
将同一个Transformer层堆栈(如L=4层)重复应用R次(循环迭代)。
有效深度为
D = L × R,但参数量仅相当于L层。关键初始化:将输出投影矩阵初始化为零,使每个模块初始为恒等映射,保证训练稳定性。
训练策略:
固定迭代:所有样本使用相同的R。
动态迭代:每个批次从泊松分布中采样R(有上下限),更贴近现实(任务复杂度未知)。
3. 任务设计与泛化挑战
构建基于知识图谱的多跳推理任务,定义三类泛化场景:
| 泛化类型 | 定义 | 难度 |
|---|---|---|
| 分布内(ID)泛化 | 在训练见过的原子知识范围内,推理未见过的事实组合。 | 基础(但Vanilla仍需长训) |
| 系统性泛化 | 组合训练中从未用于任何推理的原子知识(OOD知识组合)。 | 高(Vanilla完全失败) |
| 深度外推 | 推理深度(跳数)超过训练最大深度(如训练≤5跳,测试10跳)。 | 高(Vanilla失败) |
4. 主要发现与贡献
发现一:系统性泛化能力涌现(三阶段顿悟过程)
对比结果:
标准Transformer(R=1):系统性泛化准确率为0。
循环深度Transformer(R≥2):成功实现系统性泛化,且R越大收敛越快。
内在机制(三阶段动态):
记忆阶段:模型过拟合训练集,仅能记忆答案。
ID泛化阶段:经历“顿悟”,学会在分布内组合知识。
系统性泛化阶段:在ID泛化之后,突然获得组合全新知识的能力。
机理验证(Logit Lens):在阶段1,模型直接预测目标,无法解码中间桥梁实体;阶段2,桥梁实体变得可解码,模型开始按步骤推理;阶段3,该推理路径对OOD知识也生效。
发现二:深度外推能力与推理时计算缩放
核心发现:训练时R越大,模型能学习的基础递归深度越深。但在相同R下,测试时增加推理迭代次数可以立即解锁更深的泛化能力(例如,R=8模型在推理时用更多迭代,可从12跳泛化至24跳)。
训练策略影响:
固定迭代:训练R的上限决定了外推的潜力范围。
动态迭代:在相同最大R预算下,表现优于固定迭代,能更有效地利用训练时接触的复杂度范围。
缩放限制——过度思考:
推理迭代并非越多越好。超过某一点后,模型性能会下降。
机理分析(Logit Margin):随迭代增加,模型对正确预测的置信度先升后降。动态迭代模型比固定迭代模型更鲁棒,衰减更慢。任务越复杂,置信度峰值越低,越容易受过度思考影响。
发现三:实用的推理效率优化(自适应停带)
提出结合KL散度和预测熵的自适应停带准则,可动态决定每个样本所需的推理迭代次数,在保证性能的同时节约计算资源。
相比仅用KL散度的方法,该准则能更准确地在模型尚未确定(熵高)时继续迭代,避免过早停止。
5. 局限性与未来方向(作者自述)
可控环境:实验基于合成数据和专用token,排除了自然语言的歧义性和表面形式变化,结果向真实LLM的迁移需谨慎。
规模限制:模型规模远小于前沿LLM(如GPT-4),更大规模上的表现有待验证。
任务类型:仅研究了基于知识图谱链式推理,未覆盖所有推理类型。
未解决问题:过度思考在隐式推理中仍未完全解决,尤其是在极深推理任务上。
通过引入深度循环(参数共享)机制,让Transformer能以“顿悟”方式学会组合全新知识(系统性泛化),并能通过增加推理时的迭代次数“算得更深”(深度外推),同时揭示了过度思考这一关键瓶颈。
这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
项目地址在这里,如下所示:
摘要
我们研究隐式推理,即在单次前向传递中组合知识或规则的能力。虽然基于Transformer的大型语言模型存储了大量事实性知识和规则,但它们往往无法组合这些知识进行隐式多跳推理,这表明它们在其参数化知识上缺乏组合泛化能力。为解决此局限性,我们研究了循环深度Transformer,它允许在相同的Transformer层上进行迭代计算。我们在隐式推理场景下研究了两个组合泛化挑战:系统性泛化,即组合在训练期间从未用于组合的知识;以及深度外推,即从有限的推理深度(例如,训练最多5跳)泛化到更深的组合(例如,10跳)。通过使用从头开始训练的模型进行受控研究,我们表明,尽管标准Transformer在这两个泛化挑战上均表现不佳,但循环深度Transformer可以有效地实现这种泛化。对于系统性泛化,我们发现这种能力通过一个三阶段的顿悟(grokking)过程出现,从记忆过渡到分布内泛化,最后到系统性泛化,并得到了机制性分析的支持。对于深度外推,我们表明,通过扩展推理时的循环次数,可以解锁超越训练深度的泛化能力,更多的迭代次数能够实现更深的推理。我们进一步研究了训练策略如何影响外推,为训练循环深度Transformer提供了指导,并确定了一个关键限制,即过度思考(overthinking),其中过度的循环会降低预测质量,并限制对非常深层的组合的泛化。
1 引言
众所周知,大型语言模型(LLMs)(Brown等人,2020)在预训练期间会获取大量事实性知识,并将其存储在参数中(Geva等人,2023)。然而,这些知识如何被有效地组合用于推理,我们对此的了解仍然较少(Dziri等人,2023;Press等人,2023)。特别是,最近的研究表明,基于Transformer的LLMs在隐式推理(即单次前向传递中不带有显式思维链(CoT)的推理(Wei等人,2022))方面存在困难。这种失败揭示了Transformer的一个根本局限性:尽管存储了丰富的知识,但它们往往无法灵活地组合这些知识来解决新问题。这种局限性对泛化具有重要影响,因为许多任务需要以训练期间未见的新颖方式组合多个已知知识片段(Lake & Baroni, 2018;Berglund等人,2023)。
为什么Transformer在隐式推理中难以组合其参数化知识?考虑一个查询,例如“Imagine的表演者的配偶是”。先前的研究表明,Transformer通过链接两个事实来解决这个问题:首先在较浅层检索到Imagine的表演者是John Lennon,然后在较深层检索到John Lennon的配偶是Yoko Ono(Biran等人,2024;Wang等人,2024a;Yang等人,2024b)。然而,由于知识分布在Transformer的不同层中,无法保证特定查询所需的事实能被正确访问。例如,如果事实“John Lennon的配偶是Yoko Ono”仅存储在浅层,那么更深的层无法访问它,因为参数在各层之间不共享。虽然Transformer可以通过训练学会正确地组合此类知识(Wang等人,2024a;Yao等人,2025),但它们无法组合泛化到不熟悉的组合或更深的递归组合。
为解决这一局限性,我们在Transformer中引入了深度循环,允许将同一组层迭代应用。输入序列由一个共享的Transformer模块处理多次,每次迭代的输出作为下一次迭代的输入。与知识绑定于特定层的标准Transformer不同,循环机制使得在单次前向过程中能够更灵活地访问和组合参数化知识。这类模型,即循环深度Transformer或循环Transformer(looped transformers),作为一种有前景的架构近来受到了关注(Dehghani等人,2019;Geiping等人,2025;Zhu等人,2025)。虽然先前的研究表明循环深度Transformer可以改善长度泛化(Bansal等人,2022;Fan等人,2025),但它们是否能够在参数化知识推理中克服组合泛化的局限性尚不清楚。
在本文中,我们系统地研究了循环深度Transformer是否能够隐式地组合其参数化知识。通过构建合成数据集,我们从头开始训练模型以学习隐式推理。与在海量、不透明的网络规模语料上训练的LLM不同,这种设置使我们能够控制数据,并减轻预训练引入的混杂偏差。具体来说,我们定义了两种挑战:系统性泛化(组合在训练期间未在任何组合中使用的知识)和深度外推(例如,在5跳推理上训练,在10跳上评估)。
我们的主要发现有两方面。首先,循环深度Transformer表现出强大的系统性泛化能力,而标准Transformer则无法做到。我们表明,这种能力通过一个急剧的三阶段顿悟过程出现,该过程从记忆过渡到分布内泛化,最后到系统性泛化。我们还通过模型在不同训练阶段的内部激活证据支持了这一发现。
其次,循环深度Transformer能够实现深度外推,泛化到训练期间未见过的推理深度,因为推理时计算(即循环迭代次数)增加。我们进一步发现,训练时的循环策略对外推性能起着关键作用,其中动态循环实现了最强的泛化。尽管取得了这些进展,我们还是发现了一个关键限制:循环深度Transformer会遭受过度思考(Bansal等人,2022)的困扰,这会降低性能并限制对极深层递归的泛化。
2 相关工作
有几项小规模研究在合成任务上预训练循环或循环深度Transformer,以便在受控环境中更好地理解其行为。我们的工作与这些研究最为一致,我们能够在其中清晰地将性能和泛化的差异归因于特定的架构选择。Yang等人(2024a)展示了“循环”Transformer模块如何有助于更好地模拟学习算法,如用于上下文线性回归、2层神经网络和决策树的梯度下降。Fan等人(2025)表明,这种循环Transformer在奇偶校验和二进制加法等算法任务上提供了优越的长度泛化能力。
Saunshi等人(2025)使用Pile数据集(Gao等人,2020)的2500亿个token进行了更大规模的预训练,并发现具有相同有效深度的循环版本Transformer模型在推理方面具有更强的归纳偏置,但代价是记忆能力和困惑度。基于这些结果,他们提出了一种正则化项,鼓励某些层彼此更接近,从而改善推理和事实回忆之间的权衡。
相对于其他工作,我们的针对性设置对训练动态和模型行为产生了独特的见解。我们展示了通过循环实现权重重用如何解决标准Transformer已知难以处理的系统性组合问题,并且通过增加推理时的循环次数,多跳组合中的外推是可能的。虽然Fan等人(2025)提出了用于长度泛化的循环架构,但他们假设了基于样本复杂度的神谕(oracle)训练迭代次数。我们认为我们的设置更接近现实世界场景,其中任务复杂度不能轻易通过启发式方法(如输入长度)来估计。在没有先验任务复杂度假设的情况下,我们在训练模型时面临不同的挑战。我们分析了如何最好地应用循环深度等方法以及需要避免的常见陷阱,这有助于为未来更稳健的隐式推理模型提供信息。我们在附录D中讨论了其他相关工作。
3 任务定义
我们使用一个合成的多跳推理任务正式定义我们的隐式推理设置,并在该定义下分类了三种泛化挑战:分布内泛化、系统性泛化和深度外推(图2)。后两者可以看作是分布外(OOD)泛化。这类任务已被证明对标准Transformer难以学习(Yao等人,2025),凸显了它们在组合参数化知识进行推理方面的局限性(Allen-Zhu & Li, 2023;Yang等人,2024b)。
3.1 任务定义
3.2 泛化挑战
4 循环深度Transformer
此设计基于具有共享参数的深层网络已知的不稳定性(Agarwala & Schoenholz, 2022),这在循环深度Transformer中尤为突出(Saunshi等人,2025)。遵循Zhang等人(2019)的方法,这种初始化支持在循环迭代无限展开下的稳定优化。
循环迭代的停止策略。训练循环Transformer需要一种停止策略来确定输入前向传递中的循环迭代次数。我们考虑两种停止策略:固定迭代和动态迭代。固定迭代确定所有训练实例的循环迭代次数为相同的固定值。动态迭代策略则为每个训练批次独立采样循环迭代次数。具体来说,对于动态模型,我们采样
5 系统性泛化
在本节中,我们研究系统性泛化,即模型是否能够组合在训练期间未组合用于多跳任务的参数化知识。我们专注于2跳任务,因为Wang等人(2024a)表明标准Transformer在此简单任务上已存在困难。
5.1 实验设置
5.2 结果
循环深度Transformer能够进行系统性泛化,而标准Transformer则不能。在图3左侧,我们绘制了OOD准确率随训练轮次变化的曲线。我们发现标准Transformer(即 R=1)在需要组合不熟悉的原子事实时完全失败,而即使是最简单的 R=2 循环也能实现非平凡的泛化性能。增加训练迭代次数进一步加速了收敛,例如 R=4 在2k轮次收敛,而 R=2 需要7k轮次。这种加速不仅在训练步骤上,在绝对挂钟时间上也成立(图3,中间)。
系统性泛化通过三阶段顿悟动态出现。我们进一步分析了 R=4 模型的训练动态(图3右侧),以理解系统性泛化是如何出现的。我们观察到一个三阶段动态:在第一阶段,模型过拟合训练集,只有训练准确率在提高。在第二阶段,在经过长时间的记忆化训练后,分布内泛化出现,这种现象被称为顿悟(grokking)。在最后阶段,系统性泛化仅在模型达到接近完美的分布内准确率后才出现,这发生在比训练过拟合晚得多的时间点(例如,10^4 轮次对比 10^2 轮次)。
图3:循环深度模型在训练轮次和挂钟时间上的准确率曲线。左图:使用 (R in {1, 2, 4, 8}) 训练的模型在测试OOD上的准确率,相对于训练轮次绘制。曲线使用100轮次的滚动平均值进行平滑,阴影表示标准差。中图:相同模型在测试OOD上的准确率,相对于训练挂钟时间(小时)绘制。右图:(R = 4) 模型在训练、ID和OOD测试集样本上的准确率,相对于训练轮次绘制。
图4:使用logit透镜在相应词符位置预测桥接实体和目标实体的准确率,比较循环深度((R = 2))和8层标准Transformer。
顿悟标志着从记忆化到系统性泛化的转变。我们首先关注循环深度模型(图4左侧面板),它在三个阶段表现出不同的机制。在第一阶段,模型在没有可靠解码桥接实体的情况下预测目标,表明是记忆化。在第二阶段,桥接实体变得可解码,随后在更深的有效深度上对分布内数据的目标预测正确。仅在第三阶段,模型在OOD输入上成功,标志着从死记硬背到系统性组合的转变。
6 深度外推
在本节中,我们研究深度外推,即模型在组合其参数化知识时,能否执行比训练期间所见更深层次的递归。
6.1 实验设置
课程训练。与2跳场景不同,学习k跳任务通常需要按照跳数深度(k)对模型进行从易到难的课程训练,如Yao等人(2025)所建议。具体来说,我们首先在原子事实和2跳事实上进行训练,直到在留出的2跳测试集上达到95%的准确率阈值。随后,在课程的下一阶段,将3跳数据纳入训练,直到在留出的3跳测试集上同样达到95%的准确率。此过程对每个跳数级别 k∈{2,…,Nmax} 依次重复。为防止遗忘,在每个阶段,我们都会在所有先前引入的事实上联合训练(例如,在 k=5 阶段,我们在原子事实以及2至5跳事实上同时训练)。
由于采用这种基于阈值的课程机制,模型能力之外的训练事实永远不会被暴露。具体而言,如果模型在k跳查询上未能达到阈值以上的准确率,则训练终止,并且 (k+1) 跳数据将不会被引入。我们将模型所能达到的最大此类k值定义为其可学习递归深度。
数据集。我们通过实例化一个知识图谱来构建数据集,该图谱包含 ∣E∣=200∣ 个实体和 ∣R∣=10 个关系,每个实体的平均出度为10。我们还对知识图谱施加了排列约束,以避免模型学习到捷径解,详细说明见附录B。我们为每个 k∈[2,Kmax](其中 Kmax=40)预生成了2k个原子事实和15k个k跳推断事实。在训练期间,这些事实按照上述课程设置逐步引入。
对于每个模型,我们在750个留出的k跳事实上进行评估。跳数k达到模型可学习递归深度的事实构成分布内测试集,而超出该深度的事实则构成外推测试集。此划分因模型而异,反映了各模型的最大可实现推理深度。
模型。我们沿用第5.1节中的模型设置,但固定迭代的循环次数改为 R∈{1,2,3,4,5,6,7,8}。此处我们使用无位置嵌入(NoPE)(Kazemnejad等人,2023;Wang等人,2024b),初步研究表明其具有更好的泛化性能。除 L=4 的结果外,我们还在附录F和附录G中分别展示了不同模型大小、动态训练循环以及不同随机种子初始化下的实验结果。
6.2 分布内泛化
增加训练时的迭代次数可以提高循环Transformer的可学习递归深度。观察图5的蓝色区域,我们发现增加训练循环迭代次数相应地提高了ID泛化能力。这与先前的发现(Wang等人,2024a;Yao等人,2025)一致,即Transformer的可学习递归深度受其层数深度的限制,而我们证明了对于循环Transformer,增加训练循环迭代次数也可以增加其“有效深度”,而无需依赖额外的参数。使用动态迭代训练进一步增加了相对于固定迭代的可学习递归深度,这表明固定迭代并非最优设计选择。有趣的是,更多的循环迭代并不总是转化为更大的可学习深度(例如,R=7和R=8都最多学习到16跳任务)。
从长时间训练到快速学习的相变。我们观察到,模型需要一个长时间的训练阶段来掌握低跳任务,之后它们会迅速泛化到复杂得多的样本(图6)。我们通过绘制训练步骤与在分布内数据上实现的组合复杂度,来说明使用动态循环训练的模型。这表明主要困难在于发现底层的组合规则。一旦这种规则被内化,模型就可以快速将其扩展到复杂度高得多的样本。在图6中,我们观察到模型需要超过130万步来顿悟高达4跳的训练样本,但很快在很少的额外步骤内就学会了高达19跳的样本。除此之外,对于更复杂的样本,虽然每个新的跳数都需要额外的训练步骤才能达到 95% 的阈值,但模型在20、21和22跳上仍然在每跳不到8k额外步骤内实现了强泛化(>90%),这与从4跳到19跳每跳所需步骤相当。通过加载一个训练好的检查点(例如,第20跳),并在仅包含21跳样本的数据混合(而不是我们课程中包含所有先前阶段样本的数据混合)上继续训练,在新的划分上可以在仅50个额外训练步骤内实现超过 90% 的泛化。
图6:首次泛化到每个跳数复杂度所需的累积梯度更新次数。
6.3 深度外推
扩展推理时迭代次数可以解锁深度外推。在图5中,我们观察到当使用与训练时相同数量的循环迭代时,所有模型都难以泛化到比训练时见过的复杂度更高的任务。然而,当我们增加推理时的迭代次数时,这种限制立即得到缓解,更多的迭代次数使得模型能够泛化到越来越难的任务。值得注意的是,这种缩放效应仅在 R>4 时出现,这表明足够的训练时迭代次数是从增加的推理时计算中受益的先决条件。
训练迭代策略对外推的影响。上述结果描述了在课程设置下每个模型能达到的最大推理深度,但它们并没有区分这些差异(例如,R=6 泛化到17跳 vs. R=8 泛化到24跳)是由于更多的训练迭代次数,还是由于接触到了更复杂的训练数据(例如,R=6 训练到13跳,而 R=8 训练到16跳)。为了隔离训练迭代策略的影响,我们在相同数据(最多12跳)上训练所有模型,并评估在12-24跳任务上的外推能力。
从图7中,我们首先发现,在固定迭代模型中,增加训练时的迭代次数在缩放推理时迭代次数时能显著改善外推(例如,R=6 外推到14跳,而 R=8 达到19跳)。其次,在相同训练数据下,动态迭代策略达到了与 R=8 模型相当的外推性能(两者都达到19跳)。这与图5形成对比,在图5中动态策略泛化到的复杂度显著高于 R=8。
这些结果表明,训练期间使用的最大迭代次数决定了外推范围(即模型可以泛化到超出训练复杂度的距离),而动态迭代有效地利用了此范围,因为它实现了更大的可学习递归深度。因此,当训练数据足够复杂时,在相同的最大迭代预算下,应优先选择动态迭代而非固定策略。
图7:使用固定循环 (R in {5,6,7,8}) 和动态循环训练的模型,在最多12跳样本上的准确率。
这些结果表明,训练期间使用的最大迭代次数决定了外推范围(即模型可以泛化到超出训练复杂度的距离),而动态迭代有效地利用了此范围,因为它实现了更大的可学习递归深度。因此,当训练数据足够复杂时,在相同的最大迭代预算下,应优先选择动态迭代而非固定策略。
缩放推理时迭代次数受限于过度思考。尽管有强大的泛化性能,我们观察到当使用过大的推理时迭代次数时,性能会下降。例如,对于R=dynamic模型,将迭代次数增加到超过15次并不会改善OOD性能(图5),这种现象被称为过度思考(Bansal等人,2022)。我们沿两个关键轴对此进行研究:(1)推理迭代次数,即增加推理迭代次数如何影响模型的预测置信度;(2)任务复杂度,即增加任务复杂度如何影响此置信度。
在图8中,我们分析了logit边际,定义为正确实体的logit与最强竞争词符的logit之间的差值,并有两个观察结果。首先,在所有模型和任务中,边际随推理迭代次数增加而增加,直到达到峰值,然后随着迭代继续进行而持续下降。这表明无论任务复杂度如何,过度思考都随着迭代次数的增加而产生。值得注意的是,动态模型相比固定迭代模型表现出慢得多的边际衰减,表明动态迭代对过度思考更为鲁棒。其次,在所有模型中,峰值边际随着任务复杂度的增加而降低。这意味着对于更复杂的任务,模型的预测本身就不太自信,因此在额外迭代下更容易受到性能下降的影响。结果,增加推理迭代次数的好处对于高度复杂的任务会减弱。请注意,尽管先前的工作通过在每个迭代中注入输入信息来缓解过度思考(Bansal等人,2022;Geiping等人,2025),但我们发现这些方法并不能解决隐式推理中的这个问题。
图9:使用自适应停带与样本复杂度的平均循环迭代次数对比。
图10:基于KL散度和熵(我们的方法)与仅基于KL散度(Geiping等人,2025)的自适应停带比较。
7 结论
我们研究了循环深度Transformer是否能够组合使用参数化知识进行隐式多跳推理,这是一项对标准Transformer具有挑战性的任务。通过在从头开始训练的模型上进行受控实验,我们表明循环深度Transformer成功地解决了这两个组合泛化挑战。系统性泛化通过一个三阶段的顿悟动态出现,模型从死记硬背过渡到可泛化的解决方案。深度外推通过扩展推理时计算实现,额外的迭代允许更大的推理深度,尽管潜在的过度思考限制了在高度复杂任务上的性能。总体而言,我们的结果凸显了循环深度Transformer作为一种有前景的架构,用于参数化知识上的组合推理。