备注好友：
方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
思维链（CoT） Prompt 已成为增强语言模型推理能力的一种强大技术。然而，生成长且正确的CoT轨迹颇具挑战。最近的研究表明，循环Transformer具有出色的长度泛化能力，但其有限的通用性和适应性阻碍了它们成为自回归解决方案的替代品。
为了更好地利用循环Transformer的优势，作者提出了RELAY（通过循环对齐迭代推理）。
具体来说，作者将思维链（CoT）推理的步骤与循环迭代对齐，并在训练循环Transformer时引入中间监督。
这种额外的迭代级监督不仅保留了循环Transformer的长度泛化能力，还使其能够预测未见数据的CoT推理步骤。
因此，作者利用这种循环Transformer生成超过训练长度的复杂问题的准确推理链，然后用其微调自回归模型。
作者进行了广泛的实验，结果表明作者的方法有效，自回归模型的表现显著提升。
代码将在https://github.com/qifanyu/RELAY发布。
- 引言
推理在塑造有效的决策过程和指导人工智能系统中的问题解决策略方面发挥着核心作用。对于大语言模型（LLMs）而言，实现推理的最有效方法是通过思维链，该方法逐个生成所有中间步骤的 Token ，直到最终答案得出。然而，利用LLMs生成正确的推理过程具有挑战性。一方面，思维链过程可能非常长，有时与 Prompt 长度呈多项式增长。当推理长度超过训练数据长度时，会遇到长度泛化问题，导致准确性显著下降。另一方面，网络数据通常存在噪声，从错误的轨迹中学习可能导致错误答案。虽然合成数据可以缓解这一问题，但生成和整理这些数据需要大量的人力和专业知识。
最近，一种名为循环Transformer 的替代框架引起了关注。通常，循环Transformer是一种带有跨块参数共享的标准Transformer模型，类似于AlBERT。在该框架中，输入 Prompt （即问题）通过重复迭代相同的块进行处理，迭代次数根据问题复杂度自适应地确定。图1展示了这一过程的示意图。一些初步结果表明，循环Transformer模型具有更好的长度泛化能力，部分原因是因为问题复杂度（例如问题长度）的增加不如思维链步骤那样显著。
然而，这种方法的成功也伴随着一些实际限制。虽然对于推理任务，确定合适的循环迭代次数是可行的，但在一般任务，如翻译和摘要中，这变得成问题。此外，尽管循环Transformer可以处理特定的推理任务，但它是否具有在单个模型中管理多个推理任务的能力尚不清楚。
鉴于这些担忧，一个自然的问题随之产生：如果循环Transformer是一种通用推理器，作者能否探索将其能力整合到标准自回归模型的思维链框架中的方法？这种整合将使作者能够利用循环Transformer在复杂推理问题上的强大性能，同时保持自回归模型在多样化语言任务中表现卓越的灵活性。
本文介绍了RELAY（通过循环对齐进行推理）这一新颖框架，该框架利用循环Transformer的优越能力，帮助自回归模型处理更长的推理链。其核心方法集中在两个关键创新上。首先，作者通过实证证明，单个循环Transformer模型可以作为多个任务的一般推理器，同时保持强大的长度泛化能力。其次，作者提出了循环Transformer与思维链推理步骤之间的迭代对齐，使循环模型能够为超出训练长度的问题生成准确的推理链。这些生成的推理链可以作为自回归模型的训练数据，从而在两种架构范式之间搭建桥梁。
作者进行了广泛的实验，证明了作者的方法通过高质量生成的推理链，显著提升了自回归Transformer的推理能力。
- 相关研究工作
2.1. 基于思维链的自动回归大语言模型
思维链（CoT）已成为一种强大的技术，用于增强语言模型的推理能力，这在实证研究（Wei等，2022；Khot等，2022）和理论研究（Feng等，2024；Merrill & Sabharwal，2024）中都有所体现，尤其是在最新模型如OpenAI O1、DeepSeek r1（DeepSeek-AI等，2025）和Qwen QwQ中。通过逐个生成中间推理步骤的 Token ，这些模型有效地将复杂问题分解成一系列子过程。然而，仍存在两个关键挑战。首先，获取高质量的CoT训练数据既耗时又费力（Lightman等，2024），尤其是在需要复杂推理链的问题中。其次，生成和理解扩展推理序列可能会遇到问题。
2.2 循环 Transformer
近年来，关于循环Transformer的研究取得了显著进展。Dehghani等人（2019年）和Lan等人（2020年）的早期研究证明了在监督学习和BERT预训练中跨层参数共享的有效性。自此，这一研究方向在理论和实践两个方向上得到了扩展。在理论方面，Giannou等人（2023年）和Xu与Sato（2024年）确立了循环Transformer的基本特性，证明了其图灵完备性和描述其近似能力。Gatmiry等人（2024年）进一步深化了这一理解，展示了如何将归纳偏差融入迭代算法的学习，特别是在多步梯度下降的上下文学习背景下。从实证角度来看，循环Transformer在各种应用中表现出了良好的前景。
Yang等人（2024年）展示了它们在数据拟合任务中的参数效率，而de Luca与Fountoulakis（2024年）以及Chen等人（2024年）揭示了它们在图算法模拟和上下文学习增强方面的潜力。值得注意的是，Fan等人（2024年）在RASP-L任务中建立了它们在长度泛化能力方面的优势。在算法学习领域，Gao等人（2024年）引入了AlgoFormer框架，这是一个利用循环Transformer进行算法表示和学习的框架。尽管这些工作广泛探讨了循环Transformer的各个方面，但作者的研究采取了不同的方向。作者特别关注利用循环Transformer更好的长度泛化能力来帮助标准自回归Transformer。
2.3 长度泛化方法
Transformers在处理更长序列的能力上受到其位置编码的影响（Press等，2022年）。近期的研究主要追求两个方向以增强大语言模型（LLM）的长度泛化能力。第一个方向专注于开发高级相对位置编码方案，而第二个方向则探索通过索引粒度调整（Chen等，2023年；Peng等，2024年）和策略性索引移动对位置表示进行修改。这些研究与本论文的核心贡献是独立的。
另一条并行的研究线关注通过更好的训练数据来提升LLM的推理能力。这些方法通常利用可访问的标签或奖励来生成和筛选推理步骤，选择那些能得出正确解决方案或高奖励的步骤。然而，LLM倾向于在偶然得到正确答案的同时，生成错误或多余的中间推理步骤，这一关键局限性显著制约了LLM在复杂推理任务上的微调效果（Xia等，2024年；Zhou等，2023年）。
- 研究方法
3.1 符号表示法
任何推理任务都可以分解为三个组成部分：问题 Token 、推理 Token （即思维链步骤）以及答案 Token 。设问题 Token 序列表示为
。思维链（CoT）过程生成一系列中间推理 Token
，其中
和
分别表示问题 Token 和推理 Token 的数量。在本工作中，作者关注一个简单设置，其中问题的答案由单个 Token 表示，记为
。CoT 自回归生成。对于自回归生成，问题序列
到答案
的映射是通过逐个生成中间 Token
来实现的。
循环模型。与通过生成显式 Token 来获取答案的自回归模型不同，循环模型通过在表示空间中执行相同的函数（例如，多层Transformer块）
次，将输入序列
隐式地映射到最终答案
。迭代次数
取决于问题的复杂性。正向过程包括三个步骤：首先， Token 序列
通过嵌入函数
映射到嵌入向量。
其中
属于
，且
代表隐藏维度。其次，通过变换
，嵌入向量将迭代地进行细化。
在本文中，函数
通常是一个Transformer模型，迭代次数
根据问题长度自适应确定。最后，通过基于最后一层表示的最终答案预测头预测答案。
此设计使得模型能够通过在表示空间中迭代优化来实现隐式推理，其中每个迭代过程都能自动捕捉推理过程中的不同方面。
作为对比，
自回归生成通过首先以自回归方式生成一系列中间推理 Token
来推导最终输出
，其中长度
可以随着输入长度
的增加而多项式增长（即
）。
这样的可变长度和可能较大的
对位置编码正确反映注意力关系提出了挑战，导致长序列推理的准确性降低。相比之下，循环模型采取了根本不同的方法。它直接处理输入序列
并生成输出
。网络只需处理
而不是
，从而减轻了推理链中的长序列问题。
3.2 单推理任务中的长度泛化
在介绍作者的RELAY框架之前，作者首先通过实证研究证明了循环Transformer相较于标准自回归模型在长度泛化能力方面的优越性。这一分析构成了作者提出框架的基础和动力。
任务描述。为了验证不同方法的能力，作者采用了来自Feng等人（2024年）的三个代表性任务，包括算术，一个数学任务，以及两个动态规划（DP）问题：编辑距离（ED）和最长递增子序列（LIS）。这些任务因其多样化的解决问题模式和不同复杂程度而被选中，并且它们可以通过思维链推理过程来解决，从而得出最终答案。性能评估基于两个模型最终答案的准确性。这些任务的详细描述见附录A。
实验设置。对于每个任务，作者构建了一个包含100万训练样本和10万测试样本的数据集。对于算术任务，问题复杂度定义为操作符的数量。对于编辑距离（ED）任务，问题复杂度对应于每对字符串中较短字符串的长度。对于最长递增子序列（LIS）任务，作者将问题复杂度定义为
，其中
是输入序列的长度，因为作者的数据集按照每10个数字一个推理步骤进行结构化。训练数据集的长度分别为算术任务的
，编辑距离任务的30和最长递增子序列任务的100。为了评估模型的泛化能力，测试数据集的长度设置为算术任务的[15,25]，编辑距离任务的[30,40]和最长递增子序列任务的[100,120]。
对于自回归CoT模型，作者采用了仅解码器的标准Transformer语言模型。对于循环模型，作者使用仅编码器的Transformer，并配备了双向注意力机制。为了应对不同的问题复杂度，作者在循环模型中实现了动态迭代控制，将循环迭代次数设置为问题复杂度。所有模型的架构配置保持一致，包括3层、256维的隐藏状态和4个注意力头。对于位置编码，作者采用RoPE（Su等人，2024）对所有模型变体进行编码，以增强训练和测试用例的序列编码。
结果。图2展示了两种模型的比较性能。在训练分布内（例如算术：
运算符），循环Transformer和自回归CoT Transformer均实现了完美的准确率。然而，当在超出训练长度的问题上进行测试时，两种模型的表现出现了显著差异：自回归CoT Transformer的性能明显下降，而循环Transformer在所有长度范围内均保持了优越的性能。这证明了循环Transformer在长度泛化方面的能力。
在循环式Transformer在最终答案预测方面表现出色，但在中间计算过程中的可解释性方面存在不足。此外，它们的设计理念可能在处理通用语言任务时遇到困难，因为确定循环迭代次数在推理问题之外变得具有挑战性。这项工作旨在利用循环模型的准确推理预测来指导自回归思维链（CoT）模型的训练，以更好地处理长序列推理。
3.3 循环增强的思维链推理
直接利用经过良好训练的循环模型来提升自回归CoT模型的有效方法是将该模型用作验证器。当问题提出时，两个模型都会生成一个最终答案，如果两个答案匹配，则信任CoT的输出。然而，在实际应用中，这种方法往往失败，因为即使在达到正确最终答案的情况下，CoT模型也能产生错误的推理路径（参见第4.3节），这使得仅依赖最终答案的准确性作为指导信号变得不可靠。
作者的主要见解是，可以建立循环Transformer的迭代结构与CoT推理的逐步特性之间的对应关系。如图3所示，与CoT中逐步的 Token 生成不同，循环模型在每次迭代中同时更新其表示，而这种迭代的次数自然对应于推理轮次的数量。这种结构上的相似性为训练循环模型在并行生成每轮对应的CoT Token 的同时，保持其预测最终答案的能力开辟了可能性。基于这一见解，作者提出了RELAY（通过循环对齐迭代推理），这是一个两阶段框架，它连接了循环模型和自回归模型。
第一阶段：训练与CoT对齐的循环模型。在第一阶段，作者训练循环模型生成与CoT步骤对齐的中间推理过程。为了形式化这种对齐，假设作者有一个包含
轮的推理链。给定推理 Token
，令
表示第
轮推理的开始 Token 位置，其中每一轮都包含有效的推理 Token
。以算术推理为例，考虑表示完整推理链的 Token 序列，
。这个序列可以自然地分为
轮，使用等号作为分隔符：给定输入问题
，第一轮对应于
，第二轮对应于
，第三轮（最后一轮）呈现最终答案
。尽管轮数与循环模型的迭代次数相匹配，但由于不同推理步骤中 Token 长度的不匹配，一个关键挑战出现了。例如，涉及复杂表达式（例如，
的早期步骤通常比后期步骤（例如，
需要更多的 Token 。这种可变长度的特性对循环模型构成了挑战，该模型需要在迭代过程中保持大小为
的固定长度表示。
为了解决长度不匹配问题同时保持循环模型并行处理的能力，作者采用了右对齐填充策略。对于第
次迭代，作者通过右对齐真实推理 Token
，并用
Token 填充剩余的左侧位置，构建一个长度为
的固定长度序列
。固定长度是根据所有推理轮次及原始输入问题的最大长度确定的（请注意，推理轮次的长度通常不会超过输入问题的长度；否则，每一轮可以进一步细分为更短的轮次）。为了跟踪有效推理 Token 和填充边界，作者引入了一个二进制 Mask ：
元素级乘法
确保仅在有效的推理 Token 和最后一个
Token 上计算损失。
为了得到最终答案，作者引入了答案预测损失，以确保正确预测最终结果。
表示真实答案。总训练损失随后计算如下：
表示有效推理 Token 的位置以及最后一个
Token 的位置
。采用这种对齐策略，作者训练循环模型在每次迭代中预测相应的CoT Token ，从而使其能够生成与CoT对齐的中间输出。具体来说，在每次迭代
中，作者训练模型通过一个中间预测头来预测有效推理 Token 以及最后一个
Token 。
在中间推理步骤中，作者忽略所有先前的
Token ，除了最后一个，因为它们对推理过程没有影响。这部分损失的权重由超参数
决定，该超参数平衡了两个目标。这种设计使得循环模型能够准确预测答案，并提供可解释的中间推理步骤，这些步骤可以有效地指导第二阶段的自回归模型。
第二阶段：增强自回归CoT模型。在第二阶段，作者利用训练好的循环模型通过一个系统化的过程来增强自回归CoT模型。
首先，作者利用第一阶段训练好的循环模型来为不断增复杂的难题生成推理演示。对于每个问题
，作者得到：
表示第一阶段训练好的循环模型，
代表跨迭代生成的推理 Token ，而
是预测的答案。
作者随后利用这些演示来微调一个自回归模型。对于超出原始训练范围的问题长度，作者使用循环模型生成一个全面的推理演示数据集。随后，将新生成数据与包含初始训练长度内问题的原始训练数据集合并。这个综合数据集涵盖了原始和扩展的问题长度，然后用来在单一步骤中微调自回归模型。这种方法使模型在保留原有推理能力的同时，能够有效应对更复杂、更长的任务，利用演示提供的结构化见解。
与合成数据生成方法的比较。为了有效引导大语言模型（LLMs）处理复杂问题，先前的研究探讨了合成数据生成方法，该方法中，人工标注者根据对任务及其解决方案过程的理解构建数据生成流程。
此方法要求标注者具备三个方面的全面知识：（1）问题构建，（2）问题解决策略，（3） Pipeline 开发技能。虽然这种方法有效，但它为跨不同领域的部署设置了较高的门槛，因为找到在所有这三个领域都表现优异的专家是一项挑战。
相比之下，RELAY降低了这些要求。作者的方法遵循一个更自动化的流程：训练数据
具有强大泛化能力的循环模型
更长的问题描述
自动推理生成
自回归的CoT模型训练。人类的作用主要限于更长的问题描述，消除了对解决方案策略和 Pipeline 开发的专家知识需求。这种减少人类专业知识要求使作者的方法更实用，并能够跨不同领域进行扩展。此外，通过利用循环模型的内在泛化能力，而不是手动设计的规则，作者的方法有可能捕捉到在手工制作的 Pipeline 中可能被忽视的更细微的推理模式。
4.1. 多任务训练
在3.2节中讨论的单任务评估之后，作者将分析扩展到多任务学习环境，以探索三种模型的通用推理能力：具有显式CoT对齐的循环模型、自回归CoT模型和普通的循环模型。在这个设置中，作者联合训练模型在三个代表性的推理任务上：算术、编辑距离（ED）和最长递增子序列（LIS），每个任务都需要多步推理才能得出准确的最终答案。这种设置使得可以全面比较模型在多样化任务中有效泛化的能力。
关于任务的详细描述，包括示例输入、预期答案以及相应的思维链（CoT）推理步骤，请参阅附录A。
实验设置。作者对三个任务进行了实验：算术、编辑距离（ED）和最长递增子序列（LIS），以评估具有显式CoT对齐的循环模型与自回归CoT模型以及传统循环模型的泛化能力。训练数据集保持了与第3.2节相同的问题长度：算术的运算符数量不超过15，ED的输入字符串长度不超过30，LIS的序列长度不超过100。同样，测试数据集分别使用扩展的问题长度[15,25]、[30,40]和[100,120]构建，用于评估长度泛化能力。
- 实验
本节通过一系列旨在解决四个关键研究问题的实验，全面评估了作者RELAY框架的实证效果。
· 问题1：具有显式CoT对齐的循环模型在跨不同任务中作为通用推理器的有效性如何？（第4.1节） · 问题2：与自回归CoT模型相比，具有显式CoT对齐的循环模型在长度泛化方面表现出哪些优势？（第4.1节） 在本设置中，每个模型——具有显式CoT对齐的循环模型、自回归CoT模型和原始循环模型——都是通过在输入序列前添加特定任务的 Token （[ARI]、[ED]、[LIS]）来区分任务，并联合训练在所有三个任务上的。所有模型均在相同的指标下进行评估，仅考虑最终答案的准确性。
结果。图4展示了三种模型在三个任务上的比较性能。所有模型在训练分布的所有任务上都实现了接近100%的准确率，这表明具有显式CoT对齐的循环模型可以作为一款通用推理引擎，能够处理各种需要不同任务的复杂场景。
然而，对于训练范围之外的长度问题，具有显式CoT对齐的循环模型和原始循环模型在性能上显著优于自回归CoT模型，展示了基于循环架构在处理需要泛化到更长时间输入的任务上的优越性。此外，具有显式CoT对齐的循环模型不仅保持了原始循环模型强大的长度泛化能力，还显著超越了它，得益于CoT推理步骤与循环迭代之间的显式对齐。这种对齐提供了结构指导，增强了模型在较长长度上的推理能力以及生成显式中间CoT推理链的能力，使其既准确又可解释。这些结果将具有显式CoT对齐的循环模型确立为既稳健的推理框架，又是应对长度泛化挑战的普遍有效解决方案，在众多任务中均优于标准自回归CoT模型。（问题2）
4.2. 通过RELAY生成的CoT数据增强自回归模型
在本节中，作者采用在第4.1节中训练的具有显式CoT对齐的循环模型，通过有效数据生成来提升自回归CoT模型的性能。具体来说，作者利用其生成超出训练长度的复杂问题准确推理链的能力。这些推理链作为高质量数据，随后被用于微调自回归CoT模型。
实验设置。首先，作者采用带有显式CoT对齐的循环模型来生成针对更高复杂度问题的CoT推理链，分别针对算术、ED和LIS任务，问题长度为[15, 25]、[30, 40]和[100, 120]。然后，这些新生成数据与原始训练数据集合并，原始训练数据集包含初始训练长度的问题。关于合并数据集时不同问题长度的样本比例的详细信息，请参阅附录D.2。
接下来，作者在增强数据集上对自回归CoT模型进行单阶段微调。这一微调过程建立在第4.1节中提到的训练良好的自回归CoT模型之上，保留了相同的模型结构，同时使用增强数据集更新权重。这个过程使得模型能够融入更长的CoT链，从而增强其在较长的序列上的推理能力。
结果。图4展示了RELAY增强的自动回归CoT模型在三个任务中针对问题长度的准确率曲线。与 Baseline 自动回归CoT模型相比，使用由RELAY生成的数据（即通过具有显式CoT对齐的循环模型）微调的自动回归CoT模型在超过原始训练长度的问题上表现出了显著提升。值得注意的是，在某些情况下，其性能接近甚至略超具有显式CoT对齐的循环模型，而始终优于 Baseline 自动回归CoT模型。
这些结果表明，作者的RELAY框架通过利用循环模型的长度泛化能力以及显式的CoT对齐，有效地提升了自回归模型的总体性能。通过生成高质量的CoT推理数据，RELAY使得自回归CoT模型能够更好地处理其原始训练范围之外的难题，而无需改变其架构。（问题3）
4.3 评估由RELAY生成的中间推理步骤的可靠性
本节旨在展示通过具有显式CoT对齐的循环模型生成的CoT链的可靠性，与自回归CoT模型的自我生成数据相比。具体来说，作者强调自回归CoT模型生成的数据，即便最终结果是正确的，通常也包含错误的中间步骤，这就是为什么利用这些数据无法在较长且复杂的复杂问题中提升模型性能的原因。相比之下，通过具有显式CoT对齐的循环模型生成的数据则避免了这些问题，因为它确保了准确的中间推理步骤和最终答案，从而能够有效地微调自回归模型，显著提升其性能。
实验设置。作者评估了两种生成数据的有效性：通过具有显式CoT对齐的循环模型和自回归CoT模型。这种评估侧重于两个指标：(1)命中矩阵和(2)比特准确度，这些指标提供了对推理步骤可靠性的详细视角。
对于击中矩阵，作者选择LIS任务作为例子，因其推理步骤具有结构化特点。LIS任务的中间推理步骤遵循
矩阵格式，其中
代表100oT步骤的数量以及循环模型的迭代次数，而11代表每步的 Token 数（根据作者的数据集规定的10个数字加上一个分隔符$
$）。这种结构化格式使得LIS任务特别适合评估和可视化推理步骤的可靠性，为每个位置上匹配 GT 推理步骤的 Token 比例提供了一个直观的表示。
在所有三个不同长度的任务中均提供了比特精度，评估了整个推理步骤中按 Token 计数的全局精度。对两种不同类型生成数据进行微调的自回归CoT模型进行了比较。
为了验证自回归CoT模型生成的数据，作者在两种并行设置下进行了以下实验，分别使用普通循环模型或真实答案作为验证器。实验步骤如下：(1) 使用自回归CoT模型为长问题长度生成CoT链。(2) 通过循环模型或真实答案筛选这些数据，仅保留最终答案匹配的部分。(3) 筛选后的数据用于微调自回归CoT模型。微调过程旨在提升模型生成推理轨迹并针对更长问题达到正确最终答案的能力。
同时，通过具有显式CoT对齐的循环模型生成的数据被用作相同初始自回归CoT模型预训练权重的微调数据集，在相同的微调参数和受控的不同长度样本比例下。这两种方法在三个不同长度的任务中进行了评估，以评估性能提升。
结果。作者评估了LIS任务中问题长度为105的命中精度矩阵，这对应于
步，形成一个
的矩阵（附录B中也提供了问题长度为101的结果）。如图5所示，具有显式CoT对齐的循环模型生成数据在所有位置上均实现了持续的高 Token 精度，其中大多数值接近
，这证明了其生成高质量和可靠数据的能力。（Q4）相比之下，自回归CoT模型生成数据只在最初几个位置上具有高 Token 精度，而在后续步骤中精度稳步下降。尽管每个步骤末尾的分隔符 Token
达到了高精度，但这仅仅意味着自回归CoT模型只捕捉到了推理过程的基本格式，而无法预测准确的 Token ，这表明其在长时间问题的推理过程中维持准确预测的能力有限。
此外，作者在图4中提供了使用自生成数据进行微调的自动回归CoT模型的最终答案准确性， Token 为“AR-CoT+S·SelfChains & Loop/GT Answers”，分别对应于通过循环模型过滤的数据或真实答案，这仅比 Baseline 模型略有提升。
本文介绍了RELAY（通过循环对齐迭代推理），这是一个通过结合循环和自回归Transformer来增强思维链推理的框架。
作者的贡献表明：
（1）循环Transformer可以作为具有强大长度泛化能力的通用推理器；
（2）迭代对齐使得推理链生成能够在训练长度之外实现准确；
（3）RELAY通过高质量生成的推理链提升了自回归模型。未来的工作可以探讨循环Transformer长度泛化的理论基础，并将RELAY扩展到更广泛的语言任务中。
参考
[1]. Enhancing Auto-regressive Chain-of-Thought through Loop-Aligned Reasoning .