Reasoning with Latent Thoughts: On the Power of Looped Transformers

作者/机构: Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, Sashank J. Reddi (Google Research, Toyota Technological Institute at Chicago)

A1 主要贡献

本文探讨了大型语言模型中的推理能力，并提出了一个核心论点：许多推理问题需要的是巨大的深度，而非必然需要大量的参数。这一观点为循环模型（looped models）在推理任务中的应用开辟了新的可能性。

研究目标与核心观点：
传统观点认为，大型语言模型的推理能力主要源于其庞大的参数量和深度，这与各种缩放定律相符。然而，近期研究开始质疑这一看法，指出深度可能比参数量更为关键。本文进一步强化了这一观点，主张推理能力的核心在于深度，而参数量可以相对较少。为了实现“大深度、少参数”的目标，本文提出循环模型是理想的解决方案，即通过迭代应用同一个参数较少的函数模块来构建有效深度。

本文提出的四个核心论证（Claims）：
1. Claim 1：许多推理问题需要深度但非必要大量参数，因此可以通过循环模型解决。 通过在加法、p-跳归纳和数学应用题等合成推理任务上的实验，本文发现一个 k 层模型循环 L 次（表示为 (k ⊗ L)）的性能几乎与一个 kL 层的非循环模型相当，且远优于一个 k 层模型。理论分析也表明，许多推理问题可通过迭代算法解决，因此循环模型能以近乎最优的深度高效解决这些问题。
2. Claim 2：在语言建模中，循环模型尽管在困惑度（perplexity）和记忆方面不如计算量相同的非循环模型，但其具有促进推理能力的归纳偏置。 实验表明，(k ⊗ L) 循环模型在预训练后的困惑度高于同等计算量的 (kL ⊗ 1) 模型，但在需要推理的下游任务中，其性能非常接近甚至超过后者。这揭示了循环模型在推理和记忆任务之间的性能差异。
3. Claim 3：循环模型能够生成“隐性思维”（latent thoughts），并理论上可以模拟“思维链”（Chain-of-Thought, CoT）推理。 CoT 推理通过生成中间步骤来提升性能，本质上是一种每次迭代生成一个“思维”词元的循环模型。本文认为循环模型更为强大，因其每次迭代可以并行生成多个“隐性思维”，并从理论上证明了循环模型模拟CoT的能力。
4. Claim 4：受循环模型启发设计的正则化方法可以利用其归纳偏置来提升推理能力。 本文提出一种正则化方案，通过鼓励模型中不同层的权重相似，从而在不牺牲困惑度的情况下，继承循环模型对推理能力的促进作用。

具体贡献总结：
* 模型对比框架：系统性地比较了 k 层模型循环 L 次的 (k ⊗ L) 模型、参数量相同的 k 层 (k ⊗ 1) 模型，以及计算量（FLOPs）和有效深度相同的 kL 层 (kL ⊗ 1) 模型。
* 合成任务验证：在n元加法、p-跳归纳和i-GSM数学题等任务上，实验证明 (k ⊗ L) 模型在参数少得多的情况下，性能可以媲美甚至超越 (kL ⊗ 1) 模型。
* 语言模型中的归纳偏置：在1B参数规模的语言模型上，发现循环模型具有显著的推理归纳偏置，尽管其困惑度较高。通过困惑度-下游任务性能图谱，直观展示了这一偏置。
* 与CoT的联系与缩放行为：揭示了循环模型的性能会随着循环次数的增加而提升，类似于CoT的推理时缩放行为，并理论证明了循环模型可以模拟CoT。
* 循环启发式正则化：提出了一种新颖的正则化方法，能使标准模型获得循环模型的推理优势，同时保持较低的困惑度。

A2 方法细节

2 循环模型在简单推理任务上的应用

初步假设验证。我们首先在一系列以程序化方式构建的任务上，探索循环模型有助于推理任务的假设。我们考虑的示例性推理任务包括：n元加法、测试模型回溯p步能力的p-跳归纳头，以及由综合构建的小学数学问题组成的i-GSM。尽管这些任务显然不能涵盖所有推理问题的范畴，但它们为循环模型提供了有用的见解，并为第5节的理论结果提供了基础。

循环模型的定义。虽然文献中已提出多种循环模型的变体【31, Zhenzhong Lan et al., ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS, ICLR 2020】【11, Mostafa Dehghani et al., Universal transformers, ICLR 2018】【22, Angeliki Giannou et al., Looped transformers as programmable computers, ICML 2023】【72, Liu Yang et al., Looped transformers are better at learning learning algorithms, arXiv 2023】【43, Amirkeivan Mohtashami et al., Cotformer: More tokens with attention make up for less depth, arXiv 2023】，但为了简化探索，我们使用了最基础的版本。对于任何序列到序列函数 $f$，我们用 $f^{(L)} = f \circ f \circ \dots \circ f$ 表示函数 $f$ 循环 $L$ 次。通常，循环机制独立于 $f$ 的架构选择。在本文的其余部分，我们通常用 $f$ 表示模型的k层Transformer骨干。因此，$f$ 循环 $L$ 次与一个 $kL$ 层模型相同，其中所有 $L$ 个连续的 $k$ 层块之间共享权重。我们用符号 $(k \otimes L)$ 来表示这种循环模型。请参考图1，以简洁地了解循环机制。第5.1节提供了用于理论分析的循环Transformer的更正式定义。

2.1 简单推理问题实验

n元加法。我们考虑将n个各有3位数的数字相加的问题。加法在文献中很受欢迎，用于研究Transformer的推理方面，例如使用草稿纸【45, Maxwell Nye et al., Show your work: Scratchpads for intermediate computation with language models, arXiv 2021】，思维链推理【32, Nayoung Lee et al., Teaching arithmetic to small transformers, ICLR 2024】【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】和长度泛化【7, Hanseul Cho et al., Position coupling: Leveraging task structure for improved length generalization of transformers, arXiv 2024】。其受欢迎的原因之一是加法可以有算法解，这是许多推理问题的特点。在我们的实验中，我们在操作数 $n \in {2, 4, 8, 16, 32}$ 的均匀混合上进行训练，并从[0, 999]之间随机均匀采样每个3位数的操作数。我们直接在输入-输出对上训练所有模型，没有任何思维链步骤。以下是 $n=4$ 的一个例子：

$\text{Input: 315 + 120 + 045 + 824 =" ; Output =1304".}$

n元加法训练与比较。我们训练了一个标准的基于Transformer的基线模型（12 ⊗ 1），该模型有12层。有关训练设置的详细信息，请参阅附录A.1。我们还训练了（k ⊗ 12/k）循环模型和等参数的（k ⊗ 1）基线模型进行比较，并改变 $k \in {2, 3, 4, 6}$。所有训练好的模型最终在每个 $n \in {8, 16, 24, 32}$ 上分别进行评估，以衡量在日益困难的问题上的准确性。结果呈现在表1中。

n元加法结果分析。我们发现，尽管较浅的基线模型（k⊗1）随着k的减小而性能下降，但循环模型（k⊗12/k）表现非常出色，几乎与等效FLOPs的（12⊗1）基线模型相匹配。实际上，即使是一个1层网络循环12次也能够解决这个问题，尽管它只使用了基线模型1/12的参数。这表明加法问题主要需要深度，而不一定需要更多的参数。

p-跳归纳任务。p-跳问题是Sanford等人【54, Clayton Sanford et al., Transformers, parallel computation, and logarithmic depth, arXiv 2024b】研究的一种合成归纳任务，其灵感来源于Elhage等人【14, Nelson Elhage et al., A mathematical framework for transformer circuits, Transformer Circuits Thread 2021】对归纳头的分析。具体来说，给定一个来自字母表Σ的字母序列 $v = (v_1, \dots, v_n)$，归纳头会尝试找到 $v_n$ 在序列中倒数第二次出现的位置（从左到右），并输出紧随其后的字符。p-跳问题将这个思想推广到顺序跳跃p次。直观上，p-跳问题测试了模型递归回溯并为给定查询检索答案的能力。这让人联想到解决阅读理解类问题所需的推理能力。我们在定义A.1中给出了p-跳问题的正式定义。我们在p-跳问题上进行了循环和非循环模型的实验，字母表大小设为4，序列长度为256，并在16和32之间变化p以控制问题难度。我们的观察结果呈现在表1中。与我们在加法任务上的发现类似，深度合理的循环模型在使用少得多的参数的情况下，表现得和基线模型一样好。

i-GSM (合成小学数学问题)。受Ye等人【74, Tian Ye et al., Physics of language models: Part 2.1, grade-school math and the hidden reasoning process, arXiv 2024】的启发，我们构建了我们自己版本的小学水平数学应用题。虽然我们遵循了Ye等人【74】的许多设计准则，但我们做了一些简化的改动。我们将数学问题生成为一个模7的算术计算的有向无环图（DAG），并将图的深度限制为4。为简单起见，我们保留了问题的符号形式，而没有将它们映射成英文（例如，见表2）。我们训练了深度为1、2、4和8的模型，并在表2中将它们与不同的循环变体进行了比较。答案是模7计算的。因此，随机猜测的基线模型至少会获得14%的准确率。附录A.1有更多细节。值得注意的是，我们再次观察到，深度为k的模型循环L次，其性能与深度为kL的模型相当或更好，并且远远超过了深度为k的非循环模型。这表明，即使是更复杂、更现实的推理问题，也不需要太多参数，并且可以从通过循环增加深度中受益。

3 使用循环模型进行语言建模

在本节中，我们对因果语言模型进行预训练和评估。我们在Pile数据集【17, Leo Gao et al., The pile: An 800gb dataset of diverse text for language modeling, arXiv 2020】的250B个词元上训练模型，并使用一个24层、1B参数的模型进行大多数实验，这受到了Tay等人【67, Yi Tay et al., Ul2: Unifying language learning paradigms, ICLR 2022】设置的启发（更多细节请参考附录A.2）。

3.1 1B语言模型实验

模型设置。对于因果语言建模，我们在标准的GPT2风格的下一个词元预测目标上预训练各种循环模型【48, Alec Radford et al., Language models are unsupervised multitask learners, OpenAI blog 2019】。我们训练具有不同参数预算的模型，以确保研究结果的鲁棒性。我们提醒读者，符号 $(k \otimes L)$ 对应一个 $k$ 层模型循环 $L$ 次。对于每种设置，我们比较3个模型：(a) (24 ⊗ 1)：24层1B模型，(b) (k ⊗ 1)：一个 $k$ 层模型，其其他维度配置与24层模型相同，(c) (k ⊗ 24/k)：一个 $k$ 层模型循环24/k次，以匹配(b)的参数数量和(a)的有效深度/FLOPs。我们对 $k \in {4, 6, 8, 12}$ 进行了实验，以确保研究结果是稳健的。在Pile上预训练后，我们使用k-shot评估来评估模型的验证困惑度和下游基准测试。结果总结在表3中。

评估指标。我们在训练完成后评估模型的困惑度指标。由于越来越多的证据表明，困惑度虽然对训练非常有用，但却是衡量模型质量的一个狭隘标准，因此我们也跟踪更全面的下游评估【34, Percy Liang et al., Holistic evaluation of language models, TMLR 2023】。因此，我们在4个重要的方面评估模型：闭卷问答（closed book QA）、开卷问答（open book QA）、数学应用题和推理基元（reasoning primitives）。这些总共包括19个不同的任务。

• 闭卷问答：这包括像TriviaQA【26, Mandar Joshi et al., TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension, ACL 2017】、TydiQA-NoContext【9, Jonathan H. Clark et al., TyDi QA: A benchmark for information-seeking question answering in typologically diverse languages, TACL 2020】、Natural Questions【30, Tom Kwiatkowski et al., Natural questions: A benchmark for question answering research, TACL 2019】和Web Questions【66, Alon Talmor and Jonathan Berant, The web as a knowledge-base for answering complex questions, NAACL-HLT 2018】等任务，这些任务测试模型在没有任何上下文的情况下回答问题的能力，因此主要衡量语言模型的记忆能力。
• 开卷问答：这包括像TydiQA-GoldP【9, Jonathan H. Clark et al., TyDi QA: A benchmark for information-seeking question answering in typologically diverse languages, TACL 2020】、SquadV2【49, Pranav Rajpurkar et al., Know what you don’t know: Unanswerable questions for SQuAD, ACL 2018】、Drop【12, Dheeru Dua et al., DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs, NAACL 2019】、QuAC【8, Eunsol Choi et al., QuAC: Question answering in context, EMNLP 2018】、CoQA【51, Siva Reddy et al., CoQA: A conversational question answering challenge, TACL 2019】等任务，这些任务评估模型从提供的额外上下文中推断问题答案的能力，类似于阅读理解。
• 数学应用题：为了评估模型的推理能力，我们在Wei等人【71, Jason Wei et al., Chain of thought prompting elicits reasoning in large language models, NeurIPS 2022b】考虑的数学应用题上对它们进行测试。这包括SVAMP【46, Arkil Patel et al., Are nlp models really able to solve simple math word problems?, NAACL-HLT 2021】、ASDiv【41, Shen-Yun Miao et al., A diverse corpus for evaluating and developing english math word problem solvers, ACL 2020】、MAWPS基准【28, Rik Koncel-Kedziorski et al., Mawps: A math word problem repository, NAACL-HLT 2016】等任务。我们报告了预训练模型在这些任务上的5-shot评估结果。
• 推理基元：Saunshi等人【56, Nikunj Saunshi et al., On the inductive bias of stacking towards improving reasoning, arXiv 2024】引入这些数据集来研究堆叠（stacking）对提高推理能力的归纳偏置，通过隔离简单的推理能力。一个这样的基元是深度-k变量赋值，要求模型解决一个长度为k的赋值链。深度-0变量赋值的一个例子是a=1, b=2, c=6, b=?，深度-1变量赋值的一个例子是a=1, b=2, c=a, d=b, d=?。我们使用5-shot评估在深度-0和深度-1问题的数学和代码变体上进行评估。

% Gap 指标。对于上述每个任务组 $G$，我们在表3中报告了该任务组的平均准确率，记为 $\text{Avg}_G$。此外，对于每个层预算 $k$，我们报告了由循环模型覆盖的等参数和等FLOPs模型之间的差距百分比。更具体地说：

$$\% \text{ Gap} = \frac{\operatorname{Avg}_G(k \otimes 24/k) - \operatorname{Avg}_G(k \otimes 1)}{\operatorname{Avg}_G(24 \otimes 1) - \operatorname{Avg}_G(k \otimes 1)}.$$

该指标衡量了循环在弥合等参数和等FLOPs基线之间差距方面的有效性。它隐含地衡量了当模型只有少量参数但通过循环赋予了深度时，不同任务组的表现。% Gap接近0%意味着通过循环提供深度对该任务组没有好处，参数数量是最重要的因素。另一方面，% Gap接近100%意味着参数数量对该任务组的重要性要小得多，而通过循环获得的深度更为关键。

困惑度结果。首先，我们注意到所有(k ⊗ 24/k)循环模型相比于等参数的(k ⊗ 1)基线模型，其困惑度更优，但相比于非循环的24层基线模型，其困惑度更差。对于不同的k值，循环模型仅弥补了等参数和等FLOPs基线之间约34-50%的困惑度差距。这个困惑度差距并不太令人意外，因为循环模型的参数数量比24层基线模型少24/k倍，因此容量较低。这一点在之前的工作中也曾被观察到【31, Zhenzhong Lan et al., ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS, ICLR 2020】【43, Amirkeivan Mohtashami et al., Cotformer: More tokens with attention make up for less depth, arXiv 2023】，这也是循环模型在语言建模中被忽视的主要原因。然而，正如我们稍后将看到的，下游指标描绘了一幅更有趣、更有利的图景。

QA任务结果。我们首先考虑表3中的闭卷问答和开卷问答类别。闭卷问答任务主要测试模型的记忆能力。而开卷问答则测试模型从提供的额外上下文中推断答案的能力。因此，直观上，开卷问答任务需要更多的推理。我们首先注意到，闭卷问答（记忆）的% Gap与困惑度的% Gap非常相似。而开卷问答的% Gap，在所有参数预算下都要高得多。这表明，循环模型在上下文问答方面的改进相对记忆型问答要大得多。

数学问题与推理基元。我们还在表3中展示了数学应用题的% Gap。令人惊讶的是，我们发现当 $k \geq 6$ 时，(k ⊗ 24/k)循环模型几乎可以与基线(24 ⊗ 1)模型相媲美，尽管其参数少了k倍。事实上，12层模型循环两次（34.3）甚至明显优于24层基线模型（29.3），尽管其参数少了50%且困惑度更高；这表明循环对数学推理的提升作用尤为显著。

推理基元结果分析。为了更好地理解对推理的影响，我们请读者注意表3中推理基元的评估结果。结果非常显著：对于所有k值，(k ⊗ 24/k)循环模型在推理基元上的表现都优于等FLOPs基线模型(24 ⊗ 1)。这一点非常令人惊讶，因为这些是合成生成的任务，与预训练数据或模型架构无关。因此，解决这些任务必须依赖于上下文推理，记忆能力在这里无济于事。这些结果清楚地表明，循环模型具有提升推理能力的偏好，尽管其困惑度和记忆能力较差。接下来，我们通过等性能图（isoplots）来形式化这种对推理的归纳偏置。

3.2 对推理的归纳偏置

通过Isoplots形式化归纳偏置。在本节中，我们通过绘制困惑度与下游指标的等性能图（isoplots）来形式化归纳偏置，这一方法由Saunshi等人【55, Nikunj Saunshi et al., Understanding contrastive learning requires incorporating inductive biases, ICML 2022】引入。3.1节表明，循环模型在推理问题上表现出高于预期的性能。然而，由于循环模型的困惑度较差，很难直接在各种模型之间进行比较。一种使模型具有可比性的方法是在相同的验证预训练损失下观察它们的下游性能【37, Hong Liu et al., Same pre-training loss, better downstream: Implicit bias matters for language models, ICML 2023】。Saunshi等人【56, Nikunj Saunshi et al., On the inductive bias of stacking towards improving reasoning, arXiv 2024】提出，将预训练损失与下游指标随训练过程的变化绘制成图，作为研究各种方法归纳偏置的一种方式。对于每个模型，我们从120k步开始，每20k步评估一次对数困惑度和下游指标。我们将这些值绘制在散点图中，并用对数困惑度作为输入、相应下游指标作为输出拟合一个线性函数。请参考图2和图7中的两组等性能图。

研究发现。对于所有k值，我们观察到以下几点：
* 对于闭卷问答任务，(k ⊗ L)循环模型和(k ⊗ 1)基线模型的等性能图（如果外推）非常一致。这表明对数困惑度是基于记忆任务的下游性能的一个非常强的指标。
* 对于开卷问答和数学应用题，循环模型的等性能图线总是高于基线模型。这表明，在相同的对数困惑度下，循环模型在这些需要更多推理的任务上倾向于有更高的评估分数。
* 对于推理基元，循环模型和基线模型之间存在显著差异。循环模型在训练的大部分时间点上似乎都表现良好。

结论。总体而言，这表明循环模型具有改善推理能力的强大归纳偏置。

3.3 中间循环变体及其与渐进式堆叠的关系

与Gradual Stacking的联系。最近，Saunshi等人【56, Nikunj Saunshi et al., On the inductive bias of stacking towards improving reasoning, arXiv 2024】引入了一种名为MidAS的渐进式堆叠（Gong等人【23, Linyuan Gong et al., Efficient training of BERT by progressively stacking, ICML 2019】；Reddi等人【50, Sashank Reddi et al., Efficient training of language models using few-shot learning, ICML 2023】）方法来训练语言模型。该方法在训练过程中通过在每个堆叠操作中复制模型的某些层来逐渐增加模型深度。令人惊讶的是，他们发现MidAS不仅加速了预训练，还在图2的意义上改善了推理——即在相同的困惑度下实现了更好的推理性能。此外，该论文由于层复制操作，建立了MidAS堆叠与循环模型之间的紧密联系，并推测这是产生这种归纳偏置的原因。我们前一节的结果为这一推测提供了有力的证据，表明循环模型也显示出非常相似的归纳偏置，从而进一步加强了堆叠与循环模型之间的联系。为何会产生这种归纳偏置仍然是一个悬而未决的问题，我们认为理解这一点是一个重要的未来研究方向。

中间循环（Middle Looping）。此外，受他们发现的启发，我们探索了中间循环（见图1的图示）——这是循环的一种变体，它在网络的开始和结束部分保持独立的层，并在中间的层块上执行循环。Saunshi等人【56】的高层直觉是，第一层和最后一层在模型中扮演着特殊的角色，因此应与中间层区别对待。在表3中，我们报告了一个中间循环版本的结果，该版本与（12 ⊗ 1）基线等参数，与（24 ⊗ 1）基线等FLOPs，就像（12 ⊗ 2）模型一样。总体而言，我们发现中间循环比默认的（12 ⊗ 2）循环具有更好的困惑度和更均匀的改进（数学应用题除外），因此可能是一种更实用的循环方法。我们将探索最佳循环策略留作未来的工作。

3.4 循环的扩展行为

扩展行为研究。在本节中，我们讨论循环的一个有趣的扩展行为，特别是循环次数对各种评估指标的影响。我们特别感兴趣的是：(a) 准确率作为循环次数的函数，以及 (b) 相同有效深度的循环模型和非循环基线模型的比较。为此，我们预训练了多种形式为 (4 ⊗ L) 的循环语言模型，即4层模型循环L次，其中 $L \in {1, 2, 3, 6, 9, 12}$。为了进行等FLOPs比较，我们还训练了形式为 (4L⊗1) 的基线模型，该模型参数多L倍。对于每个任务组，我们将平均准确率绘制为有效深度 $D = 4L$ 的函数。从图3呈现的结果中，我们观察到以下几点。

观察1：深度与性能的对数关系。在这两种情况下，我们发现所有任务组的准确率都随着循环/深度的增加而持续提高，尽管毫不意外的是，对于循环和非循环模型，随着深度的增加，回报都在递减。有趣的是，对于循环和非循环模型，我们发现可以拟合一个简单的缩放定律，形式如下：

其中D是有效深度，α衡量深度对下游性能的影响。

观察2：循环对推理任务的相对优势。此外，对于每个任务组，我们计算 $\alpha_{\text{loop}}/\alpha_{\text{base}}$ 来评估“通过循环增加深度”相对于“通过额外参数增加深度”的相对影响。我们发现，增加循环次数持续有益，并且对于像开卷问答和数学问题这样的推理任务，循环的相对益处更高。值得注意的是，对于推理基元，循环的影响（1.19倍）甚至高于深度的影响，这进一步巩固了循环模型在推理方面的优势。

隐性思维及与CoT推理的联系。我们以一个重要问题结束本节讨论：为什么我们应该期待循环模型的这种有趣扩展行为？我们认为，循环模型与思维链（CoT）推理有很强的联系。这种联系很有启发性，因为最近关于“思考”的研究表明，CoT可以表现出推理时扩展行为，即随着模型思维链长度的增加，推理数据集上的准确率持续提高；也就是说，产生更多的“思考”会导致更好的推理。

循环模型与CoT的理论联系。为了建立这种联系，我们对CoT推理做了一个简单的观察——它本质上是一个在每次迭代中生成单个“思维”词元的循环模型。然而，循环模型可能更强大，因为它们可以在每次迭代中生成多个“隐性思维”。这可以在图4中看到。我们进一步将这一直觉转化为一个理论结果（见5.4节），说明循环模型实际上可以模拟CoT推理。鉴于与CoT推理的这种联系，相信循环对于更难的推理任务也能很好地扩展是可信的。我们认为，明确利用循环进行推理时扩展是一个非常有前途的未来方向。

4 受循环启发的正则化

动机与方法。在上一节中，我们观察到循环模型可以在困惑度较差的情况下改善推理能力。我们能否利用这一观察来改善推理而不影响困惑度？在这里，我们提出了一个简单的方法：对模型的权重进行正则化，以鼓励它们接近一个循环模型。这可能有两个优点：(a) 模型仍然有自由参数来改善困惑度，(b) 与循环模型的接近性可以继承其理想的归纳偏置。具体来说，如果一个L层模型表示为 $f_0 \circ f_1 \circ \dots \circ f_{L/k-1}$，其中每个 $f_i$ 是一个k层的块，我们添加一个正则化项，使得所有块 $f_i$ 的权重在余弦相似度上与 $f_{i+1}$ 相近。对于一个参数组G（例如，第一个前馈层，或Transformer中的查询矩阵），我们用 $\theta^{(0)}_G, \theta^{(1)}_G, \dots, \theta^{(L-1)}_G$ 来表示所有层中的权重。那么正则化项是：

$$\mathcal{R}_G(k) = \frac{1}{L-k} \sum_{i=0}^{\frac{L}{k}-2} \sum_{j=0}^{k-1} \text{Cosine} \left( \theta_G^{(ik+j)}, \theta_G^{((i+1)k+j)} \right)$$

总损失函数。最终的损失函数是标准交叉熵损失与在所有组上平均的正则化项之和，再乘以一个标量超参数。设G为所有参数组的集合；$G = {\text{Attn-Q, Attn-K}, \dots, \text{FFN-W2}}$

$$\mathcal{L}=\mathcal{L}_{\text {xent }}+\lambda_{\text {reg }}|\mathcal{G}|^{-1} \sum_{G \in \mathcal{G}} \mathcal{R}_G(k)$$

正则化参数的影响。在上述公式中，$\lambda_{reg} = 0$ 将恢复标准训练，而 $\lambda_{reg} \to \infty$ 将收敛到一个完全循环的模型。$\lambda_{reg}$ 的中间值将导致“近似”循环的模型。例如，为了模拟(4 ⊗ 6)循环模型的设置，我们选择 $k=4$, $L=24$，并使用一个较大的正则化强度，如 $\lambda_{reg} = 10$。为了公平比较，所有其他超参数与基线训练保持相同。我们尝试了除余弦相似度之外的其他选项，如 $l_2$ 范数，以使权重更接近，但发现余弦相似度更鲁棒和有效。

余弦相似度验证。首先，我们通过在训练结束时测量连续k层块之间的余弦相似度来检查正则化是否达到了预期的效果。我们确实发现，对于所有参数组，余弦相似度都在0.98或更高（见图5）。

归纳偏置验证。为了确认正则化器的归纳偏置，我们在图2中可视化了$\lambda_{reg} = 10$和基线模型的对数困惑度与下游任务的等性能图。虽然对于闭卷问答任务，这些图是相似的，但对于开卷问答和推理问题，显示出了强烈的归纳偏置。至关重要的是，正则化模型在不损害困惑度的情况下，在推理方面表现良好（见表4）。

5 循环模型的理论分析

分析目标。在本节中，我们提出理论结果来理解前面几节的现象——为什么参数少的循环模型能在推理问题上与计算量相当的非循环基线模型相媲美？虽然一个完整的理论是具有挑战性的，因为“推理”是一个非常宽泛的概念，但我们的目标是为循环模型的表达能力提供一些直觉和形式化。首先，我们证明循环Transformer可以有效地解决群组合问题（加法问题的推广）。然后，我们展示一个非常普遍的结果，即一个具有很少不同层的非循环模型可以被一个模型大小有小幅增加的循环模型模拟。这个结果接着被用来通过一个单层循环Transformer解决p-跳问题。我们对群组合和p-跳问题的构建在深度上几乎是最优的，并且在参数方面比非循环模型效率高得多。

5.1 预备知识与符号

Transformer架构定义。我们首先定义标准的Transformer架构。在整篇论文中，我们将嵌入维度固定为 $d \in \mathbb{N}^+$，词汇表为 $V$，最大序列长度为 $n_{\text{max}}$。我们将用 $\text{id}$ 表示恒等映射。这里，我们描述高层符号。详细符号请参考附录B.1。我们用FF和MHA分别表示前馈层和注意力层，$\theta_{\text{MHA}}$ 和 $\theta_{\text{FF}}$ 分别表示这些层中的参数。

定义5.1 (Transformer块)。给定层数 $L \in \mathbb{N}^+$ 和参数 $\theta_{TB} = (\theta^{(l)}{MHA}, \theta^{(l)}){l=0}^{L-1}$，对于任何 $n \in \mathbb{N}^+$，L层Transformer块 $TB^d)^n$ 定义为：}}: (\mathbb{R}^d)^n \to (\mathbb{R

$$\mathrm{TB}_{\theta_{\mathrm{TB}}} \triangleq\left(\mathrm{id}+\mathrm{FF}_{\theta_{\mathrm{FF}}^{(L-1)}}\right) \circ\left(\mathrm{id}+\mathrm{MHA}_{\theta_{\mathrm{MHA}}^{(L-1)}}\right) \circ \cdots \circ\left(\mathrm{id}+\mathrm{FF}_{\theta_{\mathrm{FF}}^{(0)}}\right) \circ\left(\mathrm{id}+\mathrm{MHA}_{\theta_{\mathrm{MHA}}^{(0)}}\right),$$

Transformer模型定义。我们还用EMBED和OUTPUT分别表示输入嵌入层和输出softmax层。精确定义请参考附录B.1。最后，我们定义了将词元序列映射到词元分布的整个Transformer模型：$p_{\theta}: \cup_{n \le n_{\text{max}}} V^n \to \Delta^{|V|-1}$。

$$p_{\theta} \triangleq \text{OUTPUT}_{\theta_{\text{OUTPUT}}} \circ \text{TB}_{\theta_{\text{TB}}} \circ \text{EMBED}_{\theta_{\text{TE}}, \theta_{\text{PE}}}$$

其中 $\theta = (\theta_{TB}, \theta_{TE}, \theta_{PE}, \theta_{\text{OUTPUT}})$ 表示所有的Transformer参数。特别地，我们使用 $TF_{\theta}(v_1, \dots, v_n) \triangleq \arg\max_{v \in V} p_{\theta}(v|v_1, \dots, v_n)$ 来表示Transformer模型的确定性版本。我们现在定义一个循环Transformer模型，它也包含了非循环模型。

定义 5.2 ((L ⊗ T) 循环Transformer)。给定循环次数 $T \in \mathbb{N}^+$，参数 $\theta = (\theta_{TB}, \theta_{TE}, \theta_{PE}, \theta_{\text{OUTPUT}})$，其中 $\theta_{TB} = (\theta^{(l)}{MHA}, \theta^{(l)}})^{L-1{l=0}$，我们定义一个 (L ⊗ T) 循环Transformer为 $p} \triangleq \text{OUTPUT{\theta}}} \circ (TB_{\theta_{TB}})^T \circ \text{EMBED{\theta$。}, \theta_{PE}

5.2 群组合问题

问题与结果。我们考虑组合n个群元素的问题，并证明一个标准的1层Transformer循环O(log(n))次可以解决此问题。这是模加法问题的推广，并有悠久的历史（见附录B.2.2）。最近，Liu等人【36, Bingbin Liu et al., Transformers learn shortcuts to automata, arXiv 2022】表明，深度为 $\log_2 n$ 的Transformer可以计算n个群元素的组合，无论该群是否可解。然而，Liu等人【36】的构造中每一层都使用不同的注意力参数。在这里，我们提供一个更节省参数的构造，通过循环一个单层Transformer log(n)次来解决这个问题。

定理 5.1。对于任何有限群G和每个 $n \in \mathbb{N}^+$，存在一个常数精度的循环Transformer $TF_{\theta,T}$，它使用一个1层Transformer块计算n个G中元素的组合，循环次数 $T = \lceil \log_2 n \rceil$，词汇表为 $G \cup {#}$，嵌入大小 $d = 3 (\lceil \log_2 |G| \rceil + \lceil \log_2 n + 1 \rceil)$，MLP中的隐藏维度 $d_{FF} = |G|^2 + 6 \lceil \log_2 |G| \rceil$，隐藏注意力维度 $d_{ATTN} = \lceil \log_2 n \rceil$，以及2个注意力头。更具体地说，对于任何 $g_1, \dots, g_n \in G$，$TF_{\theta} (#, g_1, \dots, g_n) = g_1 \circ \dots \circ g_n$。

结果意义。上述结果与Liu等人【36, Bingbin Liu et al., Transformers learn shortcuts to automata, arXiv 2022】为非循环模型展示的深度上界相匹配，并且非常接近那里展示的超常数下界。因此，循环模型可以用已知的最佳深度解决这个问题。

循环模型模拟非循环模型。我们的第二个理论结果表明，一个具有重复层的非循环Transformer可以被一个参数更少、深度相同的循环Transformer模拟。

定理 5.2。对于任何具有L层、d嵌入维度、MLP隐藏维度 $d_{FF}$、H个注意力头（隐藏维度为 $d_{ATTN}$）、最多R个不同Transformer层且激活值有界的Transformer $p_{\theta}$，存在一个循环Transformer $p_{\theta',L}$，它在相同的词汇表V外加一个虚拟词元#上工作，循环一个1层Transformer块L次，嵌入维度为 $d + R + 2$，MLP隐藏维度为 $Rh_{FF} + O(L)$，RH个注意力头（隐藏维度为 $d_{ATTN}$），使得对于任何字符串 $v_1, \dots, v_n \in V$，$p_{\theta}(v_1, \dots, v_n) = p_{\theta',L}(#, v_1, \dots, v_n)$。

应用。我们也可以使用上述定理将模拟群组合的 $O(\log^2 n)$ 深度Transformer转换为1层循环Transformer，尽管它在参数效率上不如前一节的结果。完整证明请参阅附录B.2.1。

p-跳理论。第2.1节中关于p-跳归纳的实验惊人地显示，一个小模型多次循环可以非常有效地解决它。我们为此发现建立了理论基础。更具体地说，我们证明一个常数层Transformer加上log(p)次循环足以解决p-跳归纳问题。这一结果实际上与Sanford等人【54, Clayton Sanford et al., Transformers, parallel computation, and logarithmic depth, arXiv 2024b】证明的非循环模型所需层数的下限相匹配。该结果由定理5.2和Sanford等人【54】的非循环模型构造（在定理B.12中重述）得出。

推论 5.3。p-跳问题（定义A.1）可以通过循环一个1层Transformer $\lfloor\log_2 p\rfloor + 2$ 次来解决，该Transformer具有 $O(\log n)$ 位的精度，嵌入大小 $d = d_{FF} = d_{ATTN} = O(1)$，MLP和注意力的隐藏维度，以及最多3个注意力头。

5.4 循环模型可以模拟CoT推理

理论模拟。在3.4节中，我们讨论了循环模型如何被视为在每次迭代中生成多个隐性思维。以下定理表明，可以使用具有L次循环的循环Transformer来模拟另一个大小相似的Transformer的L步CoT。

定理 5.4 (循环Transformer模拟CoT)。对于任何L层非循环Transformer $TF_{\theta}$，具有固定输入长度n和CoT步数m，存在一个循环Transformer $TF_{\theta'}$，其层数为 $L + O(1)$，嵌入维度增加 $\Omega(\log(n + m))$，注意力头数量增加常数个，使得对于任何输入 $v = (v_i){i=1}^n$，非循环Transformer经过m步CoT后的输出 $TF^m(v, #^m)$ 相同。}(v)$，与循环Transformer在输入x后连接m个虚拟词元并进行m次循环的输出 $TF_{\theta',m

证明思路。以下我们简述定理5.4背后的高层思想；完整证明见附录B.3。
* （除一以外全屏蔽） 我们在每个位置维护一个计数器t（引理B.10），记录当前的循环次数，对于第i个位置，我们只在 $i-n \ge t$ 时“更新”嵌入，否则将嵌入重置为循环Transformer层的输入。这可以通过MLP实现。因此，与CoT类似，前 $n + i - 1$ 个嵌入在第i次循环中不会改变。
* （移位一） 我们可以使用注意力来获取当前循环在上一位置的输出，并将其用作下一循环在当前位置的输入。（引理B.9）
* （词元解码） 我们证明可以使用带有ReLU激活的MLP来模拟CoT的编码和解码过程。（引理B.7）

A4 实验环境

合成推理任务
* n元加法：
* 数据集：操作数 $n \in {2, 4, 8, 16, 32}$ 的混合，每个操作数为3位随机数。
* 模型架构：标准Transformer，输入维度256，8个注意力头，前馈层隐藏维度1024。
* 硬件/软件：使用Adafactor优化器【59, Noam Shazeer and Mitchell Stern, Adafactor: Adaptive learning rates with sublinear memory cost, ICML 2018】，批量大小1024，训练200k步。

• p-跳归纳：

  • 数据集：字母表大小为4，序列长度为256，p在16到32之间变化。训练集4M，测试/验证集各262k。
  • 模型架构：Transformer，模型维度128，隐藏维度512，8个注意力头，使用旋转位置编码。
  • 硬件/软件：使用Adafactor优化器，批量大小256，训练200k步。

• i-GSM (合成数学题)：

  • 数据集：深度为4的有向无环图生成的符号化数学问题，运算为模7。训练集约4M，测试集约50k。
  • 模型架构：与p-跳归纳任务相同。
  • 硬件/软件：与p-跳归纳任务相同。

语言建模

• 数据集：The Pile数据集【17, Leo Gao et al., The pile: An 800gb dataset of diverse text for language modeling, arXiv 2020】，使用250B个词元进行训练。

• 模型架构：

  • 基线模型：一个1.5B参数的仅解码器Transformer模型，包含24层，模型维度2048，隐藏维度5120，32个注意力头。
  • 变体模型：循环模型和较浅的基线模型仅改变层数，其他超参数保持不变。

• 硬件/软件：

  • 训练目标：因果语言建模。
  • 配置：批量大小512，序列长度1280，使用余弦学习率调度，训练400k步，峰值学习率为0.01。
  • 代码实现：基于GPT2风格的模型。

A4 实验结果

合成推理任务（第2节）

• 实验内容：在n元加法、p-跳归纳和i-GSM任务上，比较了循环模型 (k ⊗ L) 与等参数 (k ⊗ 1) 和等计算量 (kL ⊗ 1) 的非循环模型。

• 实验结果：

  • 在n元加法和p-跳归纳任务中，(k ⊗ 12/k) 和 (k ⊗ 6/k) 的循环模型性能远超等参数的 (k ⊗ 1) 模型，并且几乎与参数量多得多的 (12 ⊗ 1) 和 (6 ⊗ 1) 等计算量基线持平（见表1）。
  • 在i-GSM任务中，k 层模型循环 L 次的性能同样可以达到或超过 kL 层非循环模型的性能（见表2）。

• 分析结论：这些结果有力地支持了核心论点1，即这些推理任务主要依赖于计算深度，而非大量独立的参数。通过循环增加深度是一种极其参数高效的方法。

语言建模（第3节）
* 实验内容：在1B参数规模下，预训练并评估循环语言模型，与等参数和等计算量的非循环模型进行比较。评估指标包括困惑度和四大类下游任务（闭卷QA、开卷QA、数学应用题、推理基元）。
* 实验结果：
* 困惑度与记忆任务：循环模型 (k ⊗ 24/k) 的困惑度低于等参数的 (k ⊗ 1) 模型，但高于等计算量的 (24 ⊗ 1) 基线。在衡量记忆的闭卷QA任务上，其性能提升幅度（% Gap）与困惑度的改善幅度相似（约34%-52%）（见表3）。
* 推理任务：在开卷QA和数学应用题上，循环模型的性能提升幅度（% Gap）显著更高（60%-100%+）。特别是在推理基元任务上，循环模型的性能全面超越了参数量多几倍的等计算量基线（% Gap > 100%）（见表3）。
* 归纳偏置：通过困惑度-下游性能图（isoplots）发现，在相同的困惑度水平上，循环模型在推理任务上的性能显著优于基线模型，证明了其具有促进推理的归纳偏置（见图2）。
* 扩展行为：模型性能随着循环次数的增加而稳定提升，尤其是在推理任务上，其扩展效率（α_loop/α_base）与通过增加新参数来加深模型相当，甚至在推理基元上更优（见图3）。

• 分析结论：循环模型在语言建模中表现出强烈的推理归纳偏置。尽管在困惑度和记忆任务上不占优势，但它们能以极高的参数效率提升模型的推理能力，这种能力还会随着循环次数的增加而扩展，类似于思维链（CoT）的机制。

循环启发式正则化（第4节）
* 实验内容：将一种鼓励层间权重相似的正则化方法应用于标准的24层1B模型，并评估其性能。
* 实验结果：应用正则化后，模型在数学应用题和推理基元等任务上取得了显著的性能提升，同时验证困惑度几乎没有变化（见表4）。训练后的模型，不同块之间的权重余弦相似度非常高（见图5），表明其结构确实趋向于循环模型。
* 分析结论：该正则化方法成功地将循环模型的推理归纳偏置迁移到了标准模型中，实现了在不牺牲模型通用能力（如困惑度）的前提下，定向增强其推理能力。

A5 结论

结论总结。本文开创了“循环模型用于推理”这一新研究方向。研究发现，循环模型不仅能用极少的参数解决多种推理问题，还表现出一种归纳偏置，即在语言模型中能不成比例地提升推理性能而非记忆能力。关于循环模型表达能力的理论结果初步揭示了其深度最优性的原因。

局限性与未来工作。尽管本文在一系列推理问题上测试了循环模型，但一个自然的问题是，这些结果是否能推广到其他形式的推理（如多模态和常识推理）。对推理问题本身进行简洁的形式化也是一个有趣的未来方向。此外，在相同困惑度下推理性能提升的归纳偏置非常引人入胜，值得进一步探索。我们发现循环模型的扩展行为非常迷人，其与隐性思维和CoT推理的联系为这种行为提供了初步解释。我们希望这能启发未来在利用循环模型实现更高效的推理时扩展方面进行探索，以辅助更深层次的推理。

A6 附录

A 实验

A.1 简单推理设置细节

n元加法。所有实验均在标准Transformer架构上运行，输入维度为256，8个头，前馈层隐藏维度为1024。我们使用Adafactor【59, Noam Shazeer and Mitchell Stern, Adafactor: Adaptive learning rates with sublinear memory cost, ICML 2018】进行训练，采用线性预热和余弦衰减的学习率调度。所有运行的批大小为1024，学习率为0.005，运行200k步。这对应于2亿个样本，与可能的总样本数（> 10^320）相比微不足道。因此，记忆答案不是问题。由于训练有些嘈杂，对于每种设置，我们运行3个不同的随机种子，并选择平均准确率最高的运行。我们选择最大值而不是平均值，因为我们关心这些模型的表达能力。

p-跳归纳。我们正式定义p-跳问题如下：

定义A.1 (p-跳, Sanford等人【54, Clayton Sanford et al., Transformers, parallel computation, and logarithmic depth, arXiv 2024b】)。对于一个有限字母表 $\Sigma$，定义映射 $hopp_p: \Sigma^n \to \Sigma \cup \{\perp\}$ 为 $hopp_p(v) = v_{\text{find}_p(v,n)}$ 如果 $\text{find}_p(v, n) \ne 0$ 否则为 $\perp$，其中

$$\begin{aligned} \begin{aligned} \operatorname{find}_1(\boldsymbol{v}, i) & =\max \left(\{0\} \cup\left\{j \leq i, v_{j-1}=v_i\right\}\right) \\ \operatorname{find}_p(\boldsymbol{v}, i) & =\operatorname{find}_1\left(\boldsymbol{v}, \operatorname{find}_{p-1}(\boldsymbol{v}, i)\right) \text { for } p \geq 2 . \end{aligned} \end{aligned}$$

p-跳归纳实验设置。对于p-跳问题，我们随机抽样实例，同时强制输入序列中始终存在p-跳。我们通过首先随机选择p-跳序列，然后将它们与填充标记一起打乱到一个序列中，再用剩余的字符填充。打乱后，我们抽样剩余的字符来替换填充标记，同时尊重p-跳的顺序。我们的训练集包含4M个样本，测试集和验证集各有262k个样本。对于所有在此数据集上训练的模型，模型维度为128，隐藏维度为512，并使用了8个注意力头。也使用了旋转位置编码。我们使用Adafactor训练200,000步，批大小为256，基础学习率为$10^{-3}$，并使用线性预热和余弦衰减的学习率调度。

i-GSM。我们在此更详细地描述i-GSM数据集是如何生成的。我们从一个深度为4的实体层次结构开始，从中构建一个随机抽样的结构图，其中有向边连接第i层的实体到第i+1层的实体。结构图中的每条边定义一个实例参数，它是一个整数（例如，市中心和停车场之间的边表示市中心的停车场数量）。然后，我们构建一个随机抽样的数学依赖图，它是所有实例参数上的一个有向无环图（DAG），通过将每个参数与其他参数关联起来。最后，我们选择依赖图中的一个节点进行查询，目标是计算该节点对某个质数P取模的数值。关于结构图和依赖图抽样过程的更多细节，我们请读者参考Ye等人【74, Tian Ye et al., Physics of language models: Part 2.1, grade-school math and the hidden reasoning process, arXiv 2024】。与Ye等人【74】相比，我们做了3个简化：我们将问题用符号语言表述，没有英文句子结构（见表2）；我们不允许问题中出现抽象参数；我们进行模7的算术运算，而不是模23。我们的训练数据集包含约400万个样本，我们在约5万个样本上进行测试。鉴于可能存在的唯一解法模板数量要大得多，问题模板空间中的训练-测试重叠概率很高。对于所有在此数据集上训练的模型，模型维度为128，隐藏维度为512，并使用了8个注意力头。也使用了旋转位置编码。我们使用Adafactor训练200,000步，批大小为256，基础学习率为$10^{-3}$，并使用线性预热和余弦衰减的学习率调度。

A.2 语言建模设置

训练设置。我们在Pile数据上使用因果语言建模目标进行训练。数据集经过预处理和缓存，以确保所有模型都在完全相同的词元上以相同的顺序进行训练。对于所有实验，我们使用512的批大小和1280的序列长度。我们使用一个在400k步内衰减的余弦学习率调度，峰值学习率为0.01，这是根据基线模型调整的。基础模型是一个1.5B参数的仅解码器Transformer模型，有24层，模型维度为2048，隐藏维度为5120，32个头。对于较浅的基线和循环模型，我们只改变层数，并保持所有其他超参数相同。

A.3 各任务组的结果

详细结果。在第3节中，我们讨论了闭卷问答、数学应用题等各种任务组的结果。为了完整性，这里我们展示了每个单独任务的结果。闭卷问答的详细结果在表5中，开卷问答在表6中，数学应用题在表7中，推理基元在表8中。这些表既包括表3中的循环模型，也包括表4中用正则化训练的模型。

B 理论结果

B.1 详细符号说明

定义B.1 (嵌入层)。给定一个有限词汇表 $V$，嵌入维度 $d \in \mathbb{N}^+$，词元嵌入参数 $\theta_{TE} \in \mathbb{R}^{d \times |V|}$ 和位置嵌入参数 $\theta_{PE} \in \mathbb{R}^{d \times n_{\text{max}}}$，我们定义嵌入层为一个序列到序列的映射，记为 $\text{EMBED}_{\theta_{TE}, \theta_{PE}}: V^n \to (\mathbb{R}^d)^n$ 对于任何 $1 \le n \le n_{\text{max}}$，其中

$$\mathsf{EMBED}_{\theta_{\mathrm{TE}}, \theta_{\mathrm{PE}}}\left(v_{1}, \ldots, v_{n}\right)=\left(\theta_{\mathrm{TE}}\left(v_{1}\right)+\theta_{\mathrm{PE}}(1), \ldots, \theta_{\mathrm{TE}}\left(v_{n}\right)+\theta_{\mathrm{PE}}(n)\right) .$$

多头自注意力机制。给定注意力参数 $\theta_{\text{ATTN}} \{W_Q, W_K, W_V, W_O\}$，其中每个 $W^m_Q, W^m_K, W^m_V, W^m_O \in \mathbb{R}^{d_{\text{ATTN}} \times d}$，我们在算法1中为仅解码器Transformer定义了带因果掩码的自注意力层。我们还定义了一个多头注意力层，作为一组具有非共享参数 $\theta_{\text{MHA}} = \{\theta^{(h)}_{\text{ATTN}}\}_{h=1}^H$ 的自注意力层的集合，其输出是每个头的输出之和。即，$MHA_{\theta_{MHA}} = \sum_{h=1}^H \text{ATTN}_{\theta^{(h)}_{\text{ATTN}}}$。

算法 1 因果自注意力, ATTN
输入: 参数 θ_ATTN = (W_Q, W_K, W_V, W_O), 输入嵌入 x_1, ..., x_n) ∈ R^d n.
输出: 输出嵌入 x' = (x'_1, ..., x'_n) ≜ ATTN_θ_ATTN (x_1, ..., x_n).
1: q_i ≜ W_Q*x_i, k_i ≜ W_K*x_i, v_i ≜ W_V*x_i, ∀i ∈ [n]
2: s_i ≜ softmax(⟨q_i, k_1⟩, ..., ⟨q_i, k_i⟩) || (0, ..., 0).
3: h'_i ≜ W_O^T * Σ_{j=1 to n} (s_i)_j * v_j .

前馈网络。给定全连接前馈网络层的参数 $\theta_{FF} = (W_1, b_1, W_2, b_2) \in \mathbb{R}^{d_{FF} \times d} \times \mathbb{R}^{d_{FF}} \times \mathbb{R}^{d \times d_{FF}} \times \mathbb{R}^{d_{FF}}$，我们定义前馈层 $FF_{\theta_{FF}}: \mathbb{R}^d \to \mathbb{R}^d$ 为 $FF_{\theta_{FF}}(h) \triangleq W_2 \text{relu}(W_1h + b_1) + b_2$。

定义 B.2 (输出层)。给定参数 $\theta_{\text{OUTPUT}} \in \mathbb{R}^{d \times |V|}$，我们将输出层表示为 $\text{OUTPUT}_{\theta_{\text{OUTPUT}}}: (\mathbb{R}^d)^n \to \Delta^{|V|-1}$，其中

$$\mathrm{OUTPUT}_{\theta_{\mathrm{OUTPUT}}}\left(x_1, \ldots, x_n\right) \triangleq \operatorname{softmax}\left(x_n^{\top} \theta_{\mathrm{OUTPUT}}\right)$$

完整Transformer模型。最后，我们定义整个Transformer模型 $p_{\theta}: \cup_{n \le n_{\text{max}}} V^n \to \Delta^{|V|-1}$ 为

$$p_{\theta} \triangleq \mathrm{OUTPUT}_{\theta_{\mathrm{OUTPUT}}} \circ \mathrm{TB}_{\theta_{\mathrm{TB}}} \circ \mathrm{EMBED}_{\theta_{\mathrm{TE}}, \theta_{\mathrm{PE}}}$$

对于任何 $\theta = (\theta_{TB}, \theta_{TE}, \theta_{PE}, \theta_{\text{OUTPUT}})$。为方便起见，我们也将 $[p_{\theta}(v_1, \dots, v_n)](v)$ 写为 $p_{\theta}(v | v_1, \dots, v_n)$。特别地，我们使用 $TF_{\theta}(v_1, \dots, v_n) \triangleq \arg\max_{v \in V} p_{\theta}(v|v_1, \dots, v_n)$ 来表示Transformer模型的确定性版本。

有限精度建模。在本文中，我们假设Transformer是有限精度的。具体来说，我们遵循Li等人【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】的设置，并使用简写 $\mathcal{F}_s \triangleq \{c \cdot k \cdot 2^{-s} | c \in \{-1, 1\}, 0 \le k \le 2^{2s} - 1, k \in \mathbb{N}\}$ 来表示常数精度s的定点数，并用舍入操作 $[\cdot]_s: \mathbb{R} \to \mathcal{F}_s$ 来表示校正舍入，即从 $\mathbb{R}$ 到 $\mathcal{F}_s$ 中最接近的可表示数的映射。（我们通过选择绝对值较小的数来打破平局）。我们假设(1) Transformer的所有参数都在 $\mathcal{F}_s$ 中，并且(2) 在Transformer的前向传播中，每次二元运算后都会执行校正舍入。我们将请读者参考Li等人【33】以获取关于这种有限精度建模的详细讨论，并仅列出本文将使用的重要符号和引理。

符号与引理。我们用 $1_s$ 表示长度为s的全一向量。类似地，我们定义 $\langle \cdot, \cdot \rangle_s, \times_s$ 和 $\text{softmax}_s$。我们回顾，对于任何 $s \in \mathbb{N}$ 和整数 $0 \le x \le 2s-1$，我们用 $\text{bin}_s(x) \in \{0, 1\}^s$ 表示整数x的常规二进制编码，使用s个二进制位，即 $x = \sum_{i=1}^s 2^i(\text{bin}_s(x))_i$，用 $\text{sbin}_s(x) \in \{-1, 1\}^s$ 表示有符号二进制编码，即 $2\text{bin}_s(x) - (1, \dots, 1)$。最后我们定义 $B_s = \max \mathcal{F}_s = 2^s - 2^{-s}$。

引理 B.1。 [引理 E.1, 【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】] 对于任何 $s \in \mathbb{N}^+$，有 $[\exp(-B_s)]_s = 0$。

引理 B.2。 [引理 E.2, 【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】] 对于任何 $s \in \mathbb{N}^+$，有 $[\exp(B_s)]_s = B_s$。

引理 B.3。 [引理 E.5, 【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】] 无限扇入的AND, OR (相应地, MAJORITY) : $\{0, 1\}^n \to \{0, 1\}$ 可以由一个2层前馈ReLU网络模拟，该网络具有常数（相应地, $\log n$）位的精度，常数隐藏维度，以及额外的n个值为1的常数输入。

引理 B.3 数学表达。数学上，令 $\mathcal{FF}[s(n)]$ 是函数 $C: \{0, 1\}^n \to \{0, 1\}$ 的集合，这些函数可以是一个具有至多 $s(n)$ 位精度和常数隐藏维度的两层前馈ReLU网络 $FF_{\theta}: \{0, 1\}^{2n} \to \{0, 1\}$，$FF_{\theta}(x') = W_2 \times_s \text{relu}([W_1 \times_s x' + b_1]_s)$，其中 $\theta = (W_2, W_1, b_1)$，使得对于任何 $x \in \{0, 1\}^n$，

$$\mathrm{FF}_\theta(x_1, 1, x_2, 1, \dots, x_n, 1) = C(x).$$

我们有无限扇入的 AND, OR $\in \mathcal{FF}[1]$ 和 MAJORITY $\in \mathcal{FF}[\log n]$。

向量交错。给定两个相同长度s的向量x, y，我们用 $x \frown y$ 表示它们的交错，即 $(x \frown y)_{2i-1} = x_i, (x \frown y)_{2i} = y_i$ 对于所有 $i \in [e]$。

引理 B.4。 [引理 E.3, 【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】] 对于任何 $s \in \mathbb{N}^+$，令 $q_i = \text{sbin}_s(i) \frown 1_s$ 和 $k_i = B_s \cdot (\text{sbin}_s(i) \frown (-1_s))$ 对于所有 $i \in [2^s-1]$，则有 $-\exp(\langle q_i, k_j \rangle_s)_s = 1[i=j]$ 对于所有 $i, j \in [2^s-1]$。

B.2 证明

B.2.1 循环模型可以模拟非循环模型

定理5.2的证明。我们首先引入一些更多的符号。我们将对任何固定的序列 $v = (v_1, \dots, v_n)$ 证明该定理。我们使用 $x^{(l)} = (x^{(l)}_i)_{i=1}^n$ 来表示 $p_{\theta}$ 在第 $l$ 层的中间嵌入。更具体地，我们定义

$$ \boldsymbol{x}^l = (\text{id} + \text{FF}_{\theta_{\text{FF}}^{(l-1)}}) \circ (\text{id} + \text{MHA}_{\theta_{\text{MHA}}^{(l-1)}}) \circ \cdots (\text{id} + \text{FF}_{\theta_{\text{FF}}^{(0)}}) \circ (\text{id} + \text{MHA}_{\theta_{\text{MHA}}^{(0)}}) \circ \text{EMBED}_{\theta_{\text{EMBED}}}(\boldsymbol{v}) $$

我们也使用 $x^{(l+0.5)} = (x^{(l+0.5)}_i)_{i=1}^n$ 来表示 $p_{\theta}$ 在第 $l$ 层经过注意力层后的中间嵌入。

$$x^{l+0.5} = (\text{id} + \text{MHA}_{\theta_{\text{MHA}}^{(l-1)}})(x^l).$$

循环Transformer的中间嵌入。类似地，对于构造的循环Transformer $p_{\theta,T}$，我们用 $v' = (\#, v_1, \dots, v_n)$ 表示其输入。为简单起见，我们约定#在位置0。如果#从位置1开始，证明仍然成立，因为我们只需将后面的词元移动1个位置。我们定义 $x'^{(l)} = (x'^{(l)}_0, x'^{(l)}_1, \dots, x'^{(l)}_n)$ 为 $p_{\theta}$ 在第 $l$ 层的中间嵌入，以及 $x'^{(l+0.5)} = (x'^{(l+0.5)}_0, x'^{(l+0.5)}_1, \dots, x'^{(l+0.5)}_n)$ 为 $p$ 在第 $l$ 层的中间嵌入。

构造属性。下面我们首先陈述我们构造将满足的关键属性，这些属性蕴含了定理的正确性，然后我们陈述我们对 $p_{\theta',T}$ 的构造，并证明这些属性确实被满足：
* $x_i'^{(l)} = (x_i^{(l)}, 1_R - e_{r(l)}, l, 1[i=0])$
* $x_i'^{(l+0.5)} = (x_i^{(l+0.5)}, 1_R - e_{r(l)}, l, 1[i=0])$
* $x_0^{(l)} = x_0^{(l+0.5)} = 0$

构造细节。要获得最后一个坐标l，即深度计数器，我们只需在MLP中增加1个隐藏维度。

MLP构造。接下来我们展示如何使用两层MLP，其中有L+2个神经元，来得到从 $\ell$ 到 $1_R - e_{r(l)}$ 的映射。令 $(\theta_{FF}^{(i)}, \theta_{MHA}^{(i)})_{i=1}^R$ 为 $\theta$ 中R个不同层的参数。我们假设在第 $l$ 层，使用的是 $r(l)$ 的参数。这是因为 $e_{r(l)} = \sum_{i=1}^L e_{r(i)}0.5 * ([l-i+1]^+ - 2[l-i]^+ + [l-i-1]^+)$。

停用MLP神经元。现在我们解释如何停用不需要的MLP神经元。换句话说，我们对 $\theta'_{FF}$ 的构造本质上是在FF的隐藏维度上拼接 $\theta_{FF}^{(i)}$ for $i \in [r]$，并附加控制条件，使得在第 $l$ 层，$FF_{\theta'_{FF}}((x_i^{(l)}, 1_R - e_{r(l)}, l), 1[i=0]) = \sum_{i=1}^R FF_{\theta_{FF}^{(i)}}(x_i^{(l)}) 1[r(l)=i, i \ne 0]$。这个控制可以通过从 $1 - e_{r(l)} + 1[i=0]$ 中减去一个大于隐藏层中最大预激活值的常数来实现。

停用注意力。最后我们解释如何停用不需要的注意力。我们将只使用注意力来更新嵌入的第一部分，即 $x_i^{(l+0.5)}$。这里一个关键步骤是我们将#的词元嵌入设为0。我们按如下方式构造键和查询：
1. $W'^{(r')}_Q(x'^{(l)}_i) = (W^{(r')}_Q x_i^{(l)}, 1 - 1[r'=r(l)])$ for $r' \in [R]$ and $i = 0, \dots, n$
2. $W'^{(r')}_K(x'^{(l)}_i) = (W^{(r')}_K x_i^{(l)}, -B1[i=0])$ for $r' \in [R]$ and $i = 0, \dots, n$，其中B是某个大于注意力中最大先前内积的常数，$\max_{l \in [L], i, j} \langle W_K x_i^{(l)}, W_Q x_i^{(l)} \rangle$。
3. $W'^{(r')}_O W'^{(r')}_V(x'^{(l)}_i) = (W^{(r')}_O W^{(r')}_V x_i^{(l)}, 0, 0, 0)$。

构造有效性。这种构造是有效的，因为只有“期望的”注意力头 $r=r(l)$ 会被激活并像非循环情况一样运作，否则该注意力头中的所有位置将完全关注到位置0并返回一个零向量。（我们可以选择足够大的B，使得注意力分数计算出的分布是位置0处的delta分布），这会产生一个零向量作为其值。这完成了证明。□

B.2.2 群组合

背景。自动机理论中的里程碑式成果，克罗恩-罗兹分解定理（Krohn & Rhodes, 1965），表明所有具有可解变换群的半自动机（包括可解群的组合问题）都可以被置换-重置自动机的级联模拟，而后者可以被TC0电路模拟。Liu等人【36, Bingbin Liu et al., Transformers learn shortcuts to automata, arXiv 2022】进一步表明，这种具有可解变换群的自动机可以被常数深度Transformer连续模拟。然而，也有研究表明（Barrington, 1986），不可解群的组合问题是

算法 2 群组合
输入：群元素 g0, g1, ..., gn ∈ G, 其中 g0 = e。
输出：g0 ◦ g1 ◦ ... gn。
1: g_i^(0) = g_i, ∀0 ≤ i ≤ n。
2: for l = 1 → ⌈log2 n⌉ do
3: a_i^(l) = g_([2i-n-1]+)^(l-1), b_i^(l) = g_([2i-n]+)^(l-1) ∀0 ≤ i ≤ n。
4: g_i^(l) = a_i^(l) ◦ b_i^(l), ∀0 ≤ i ≤ n。
5: end for
6: return g_n^(⌈log2 n⌉)。

NC1完备性与Transformer。NC1完备的，例如5个元素上的置换群S5的组合。在NC1 $\neq$ TC0的普遍硬度假设下，常数深度Transformer无法通过单次前向传播解决S5的组合问题【40, William Merrill and Ashish Sabharwal, The parallelism tradeoff: Limitations of log-precision transformers, TACL 2023b】【36, Bingbin Liu et al., Transformers learn shortcuts to automata, arXiv 2022】【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】。但通过CoT，非常浅的Transformer（深度等于一或二）可以模拟任何群的组合问题【33, Zhiyuan Li et al., Chain of thought empowers transformers to solve inherently serial problems, ICLR 2024】【39, William Merrill and Ashish Sabharwal, The expresssive power of transformers with chain of thought, arXiv 2023a】。

定理5.1的证明。我们将#的词元嵌入设置为与G的单位元e相同。在下面的证明中，我们直接将#视为e。我们将按照算法2来构造模拟群组合的Transformer，该算法给出了构造的高层思想。算法2的正确性源于群组合的结合律。更具体地，我们可以通过归纳法验证，对于所有的 $l = 0, \dots, \lceil \log_2 n \rceil$，$g_0^l \circ g_1^l \circ \dots \circ g_n^l$ 都是相同的，并且在最后一轮，即当 $l = \lceil \log_2 n \rceil$ 时，对于所有 $i < n$，都有 $g_i^{(l)} = e$。

Transformer构造细节。下面我们展示如何构造一个给定大小的Transformer来模拟上述算法2。我们将每个 $g \in G$ 嵌入为一个不同的向量 $\mathbf{g} \in \{-1, 1\}^{\lceil \log_2 |G| \rceil}$，每个位置 $0 \le i \le n$ 嵌入为其在 $\mathbf{i} \in \{-1, 1\}^{\lceil \log_2 n+1 \rceil}$ 中的二进制表示，这是 $\text{sbin}_s(i)$ 的简写，其中 $s = \lceil \log_2 n + 1 \rceil$。我们将它们拼接起来得到 $\{x_i^{(0)}\}_{i=0}^n$，即 $x_i^{(0)} = (\mathbf{g}_i, \mathbf{i}, [\mathbf{2i-n-1}]^+, [\mathbf{2i-n-1}]^+, 0^{\lceil \log_2 |G| \rceil}, 0^{\lceil \log_2 |G| \rceil})$。为方便起见，我们将省略末尾的0（在本文的其他证明中也是如此），并将其写为 $x_i^{(0)} = (\mathbf{g}_i, \mathbf{i}, [\mathbf{2i-n-1}]^+, [\mathbf{2i-n-1}]^+)$。下面我们展示可以构造一个参数为 $(\theta_{MHA}, \theta_{FF})$ 的1层Transformer块，满足：
1. $h_{MHA_{\theta_{MHA}}}((\mathbf{g}_i, \mathbf{i}, [\mathbf{2i-n-1}]^+, [\mathbf{2i-n-1}]^+))_{i=0}^n)_k = (0^{\lceil \log_2 |G| \rceil + 3 \lceil \log_2 n+1 \rceil}, \mathbf{g}_{[\mathbf{2k-n-1}]^+}, \mathbf{g}_{[\mathbf{2k-n}]^+})$ 对于所有 $g_0=e, g_i \in G \forall i \in [n], k=0, \dots, n$；
2. $FF_{\theta_{FF}}(\mathbf{g}, \mathbf{i}, \mathbf{j}, \mathbf{k}, \mathbf{g}', \mathbf{g}'') = (\mathbf{g}' \circ \mathbf{g}'' - \mathbf{g}, 0^{3 \lceil \log_2 n+1 \rceil}, -\mathbf{g}', -\mathbf{g}'')$，对于所有 $i,j,k=0,\dots,n, g,g',g'' \in G$。

MHA和FFN的实现。第一个声明成立是因为我们可以使用两个注意力头分别检索 $\mathbf{g}_{[\mathbf{2k-n-1}]^+}$ 和 $\mathbf{g}_{[\mathbf{2k-n}]^+}$，其中两者都使用 $\mathbf{k}$ 作为键，并分别使用 $-[\mathbf{2k-n-1}]^+$ 和 $-[\mathbf{2k-n}]^+$ 作为查询。这是可能的，因为所有需要的信息都已在 $x_i$ 中。我们进一步将注意力温度设置得足够低，以便注意力返回的概率是键等于负查询（四舍五入后）位置的独热分布。

FFN的实现。现在我们转向关于MLP的第二个声明。我们将使用 $|G|^2$ 个带有ReLU激活和偏置的神经元来模拟 $\mathbf{g}'$ 和 $\mathbf{g}''$ 的乘积。我们可以用 $(h, h')$ 对每个神经元进行索引，其中 $h, h' \in G$，并设置其输入权重 $[W_1]_{(h,h'),:} = (\mathbf{h}, \mathbf{h}')$，偏置 $(b_1)_{(h,h')} = -2\lceil \log_2 |G| \rceil + 1$，这确保了只有当 $\mathbf{g}'=\mathbf{h}, \mathbf{g}''=\mathbf{h}'$ 时，神经元 $(h, h')$ 的激活才会为1，否则为0。然后将神经元 $(h, h')$ 的输出权重设置为 $\mathbf{h} \circ \mathbf{h}'$，并将第二层的偏置设置为0，就完成了模拟群组合的构造。最后，我们使用剩余的 $6\lceil \log_2 |G| \rceil$ 来模拟剩余 $3\lceil \log_2 |G| \rceil$ 嵌入维度的负恒等映射 $x \to -x$。证明完毕。□

CoT推理递归定义。在本节中，我们建立了循环模型和CoT推理之间的联系。我们首先将CoT推理的递归定义如下：

$$\mathsf{TF}_{\theta}^{i}(v_{1}, \ldots, v_{n}) \triangleq \mathsf{TF}_{\theta}^{i-1}(v_{1}, \ldots, v_{n}, \mathsf{TF}_{\theta}(v_{1}, \ldots, v_{n})),$$

对于满足 $i+n \le n_{max}-1$ 的 $i, n \in \mathbb{N}^+$，以及基本情况 $TF^1_{\theta}(v_1, \dots, v_n) \triangleq TF_{\theta}(v_1, \dots, v_n)$。对于所有 $0 \le i \le n_{max}-n-1$，经过i步CoT的输出是 $v_{n+i+1} = TF^{i+1}_{\theta}(v_1, \dots, v_n) = TF_{\theta}(v_1, \dots, v_n, v_{n+1}, \dots, v_{n+i})$。

定理 B.5。我们首先给出正式的陈述。

定理 B.5 (循环transformer模拟CoT)。对于任何L层非循环transformer $TF_{\theta}$，存在一个循环transformer $TF_{\theta'}$，其层数为 $L+O(1)$，嵌入、MLP和注意力层的维度增加常数个，注意力头的数量也增加常数个，使得对于任何输入 $v = (v_i)_{i=1}^n$ 和整数m，非循环transformer经过m步CoT后的输出 $TF^m_{\theta}(v)$，与循环transformer在输入x后连接m个虚拟词元并进行m次循环的输出 $TF_{\theta',m}(v, \#^m)$ 相同。

辅助引理。以下是一些有助于证明定理B.5至少与CoT一样强大的辅助引理。

引理 B.6 (使用MLP模拟arg max)。对于每个 $d \in \mathbb{N}$ 和精度 $s \in \mathbb{N}^+$，存在一个3层网络，具有relu激活和 $d^2$ 宽度 $f$ 以及s精度，使得对于任何 $x \in \mathcal{F}_s^d$，如果存在 $k \in [d]$，使得 $x_k > \max_{j \ne k, j \in [d]} x_j$，则 $f(x) = e_k$。

引理 B.6的证明。对于每个 $i \in [n]$，定义 $g_i = 2^s \cdot \text{relu}(2^{-s} - \sum_{j \ne i} \text{relu}(x_j - x_i))$。我们声称，如果存在 $k \in [d]$，使得 $x_k - \max_{j \ne k, j \in [d]} x_j \ge 2^{-s}$，那么对于所有 $i \in [d]$，$g_i = 1$ 当且仅当 $i = k$。首先 $g_k = 2^s \cdot \text{relu}(2^{-s}) = 1$。其次，对于 $i \ne k$，显然有 $\sum_{j \ne i} \text{relu}(x_j - x_i) \ge 2^{-s}$，因此 $g_i \le 0$。这个构造显然可以通过一个具有s精度的3层relu网络实现。□

引理 B.7 (使用MLP模拟解码和嵌入)。给定任何s精度的 $\theta_{TE} \in \mathbb{R}^{d \times \Sigma}$ 和 $\theta_{\text{OUTPUT}}$，存在一个5层网络 $f: \mathbb{R}^d \to \mathbb{R}^d$，具有relu激活和 $\max(|\Sigma|^2)$ 宽度以及s精度，使得对于所有s精度的 $x \in \mathbb{R}^d$，其对于 $v \triangleq \arg\max_{o \in \Sigma} (x^\top \theta_{\text{OUTPUT}})(o)$ 有唯一的arg max，则有

$$f(x)=\theta_{\mathsf{TE}}(v).$$

引理 B.7的证明。这是引理B.6的一个简单应用。

引理 B.8 (控制门)。一个具有s精度的2层relu网络可以实现 $F : \mathcal{F}_s \times \mathcal{F}_s \times \{0, 1\}$，$F(x, y, M) = Mx + (1-M)y$。

引理 B.8的证明。注意到 $F(x, y, M) = \text{relu}(x - 2s \cdot (1-M)) - \text{relu}(-x - 2s \cdot (1-M)) + \text{relu}(y - 2s \cdot M) - \text{relu}(-y - 2s \cdot M)$。证明完成。□

定义 B.3 (带掩码的Transformer块)。给定层数 $L \in \mathbb{N}$，参数 $\theta_{TB} = (\theta^{(l)}_{MHA}, \theta^{(l)}_{FF})_{l=0}^{L-1}$，以及掩码函数 $M: \mathbb{N} \to \{0, 1\}$，我们定义L层带掩码的Transformer块 $TB_{\theta, M}: (\mathbb{R}^d)^n \to (\mathbb{R}^d)^n$ 对于任何 $n \in \mathbb{N}^+$ 如下：

定义 B.4 (带掩码的循环Transformer)。给定循环次数 $T \in \mathbb{N}^+$，参数 $\theta = (\theta_{TB}, \theta_{TE}, \theta_{PE}, \theta_{\text{OUTPUT}})$，以及掩码函数 $\{M_t\}_{t=1}^T$，其中 $\theta_{TB} = (\theta_{MHA}^{(l)}, \theta_{FF}^{(l)})_{l=0}^{L-1}$，我们定义带掩码的循环Transformer为 $p_{\theta,T,M} \triangleq \text{OUTPUT}_{\theta_{\text{OUTPUT}}} \circ TB_{\theta_{TB}, M_T} \circ \dots \circ TB_{\theta_{TB}, M_1} \circ \text{EMBED}$，以及相应的确定性版本为 $TF_{\theta,T,M}(v_1, \dots, v_n) \triangleq \arg\max_{v \in V} p_{\theta,T,M}(v|v_1, \dots, v_n)$。

定义 B.5 (移位层)。我们定义移位层 SHIFT: $(\mathbb{R}^d)^n \to (\mathbb{R}^d)^n$ 如下，对于任何 $d, n \in \mathbb{N}^+$ 和 $x_1, \dots, x_n \in \mathbb{R}^d$：

$$\text{SHIFT}(x_{1}, x_{2}, x_{3}, \dots, x_{n}) = (x_{1}, x_{1}, x_{2}, x_{3}, \dots, x_{n-1}).$$

引理 B.9。对于输入序列长度最大为某个整数n，SHIFT可以通过一个注意力层实现，方法是将每个嵌入 $x_i$ 与 $(\text{sbin}_s(i), \text{sbin}_s(f(i)))$ 拼接，其中 $s = \lceil \log_2 n + 1 \rceil$。

引理 B.9的证明。这等价于证明我们可以构造一个注意力头，在每个位置i计算 $x_{f(i)}$。

实现细节。为此，我们只需调用引理B.6，并用它来设置键和查询，这样位置i就会关注到位置f(i)。我们将位置i的值设为 $x_i$。证明完成。□

引理 B.10。对于任何正整数 $s > 0$，存在一个常数深度的MLP F，每层有 $O(s)$ 个隐藏神经元，参数在 $\mathcal{F}_s$ 中，使得对于任何输入 $\text{bin}(x) \in \{-1, +1\}^s$，其中 $0 \le x \le 2^s - 1$，有

$$F(x) = \text{bin}(x+1).$$

引理 B.10的证明。根据引理B.3，只需证明我们可以用 $O(s)$ 宽、常数深度的布尔电路（带无限扇入的AND、OR、NOT门）来模拟 $\text{bin}(x) \to \text{bin}(x)+1$。这可以立即看出，因为

引理 B.11。对于任何正整数 $s > 0$，存在一个常数深度的MLP F，每层有 $O(s)$ 个隐藏神经元，参数在 $\mathcal{F}_s$ 中，使得对于任何输入 $(\text{bin}(x), \text{bin}(y)) \in \{-1, +1\}^s$，其中 $0 \le x, y \le 2^s - 1$，有

引理 B.11的证明。根据引理B.3，只需证明我们可以用 $O(s)$ 宽、常数深度的布尔电路（带无限扇入的AND、OR、NOT门）来模拟 $\text{bin}(x), \text{bin}(y) \to 1[x>y]$。这可以立即看出，因为

定理 B.5的证明。我们考虑掩码 $M_t(i) = 1[i-t \ge n]$，我们称之为CoT掩码。设 $s = \lfloor\log(n+m)+1\rfloor$，为方便起见，我们用 $\mathbf{i}$ 表示 $\text{sbin}_s(i)$，其中 $1 \le i \le n+m$。我们新构造的transformer的嵌入大小比要模拟的目标transformer大一个加性常数3s。我们用 $(x_i^{(t)}, p_i^{(t)})$ 表示第t次循环后位置i的新嵌入，其中 $x_i^{(t)}$ 将是待模拟transformer的原始嵌入，而 $p_i^{(t)}$ 的维度为3s，仅依赖于i和t。特别地，我们可以证明，我们可以设置 $p_i^{(t)} \triangleq (\mathbf{i}, [\mathbf{i-2}]^+ + 1, \mathbf{n+t})$ 来保存位置和循环次数的信息——$p_i^{(0)}$ 来自位置嵌入，更新是由于引理B.10。MLP和注意力（键、查询、值）的隐藏维度大小也将增加 $O(s)$。

证明步骤。证明包含两个步骤：
1. 证明存在一个带有CoT掩码的Transformer，通过循环其自身的L+O(1)层块m次，来模拟具有m步CoT和L层的目标Transformer。
2. 证明上述带有CoT掩码的循环Transformer可以被一个没有掩码、层数增加常数个的标准循环Transformer模拟。

第一步证明。对于第一个声明，我们从与带CoT的Transformer相同的参数 $\theta$ 开始，构建一个新的循环模型，其参数为 $\theta'$，每个Transformer块中层数增加常数个，每个注意力层中头数最多增加常数个。首先，我们可以增加常数个层，使用MLP通过引理B.7来模拟解码-编码过程。接下来，我们可以在每个块中增加一个Transformer层，并使用注意力层通过引理B.9来模拟移位层，因为我们有额外的位置嵌入 $p_i^{(t)}$。特别地，我们在t次循环后在位置n+t得到的嵌入 $x_{n+t}^{(t)}$，现在模拟了CoT Transformer的n+t的词元嵌入。根据我们定义CoT掩码M的方式，对于每个 $t \ge -n+1$，嵌入 $\hat{x}_{n+t}^{(t')}$ 对于所有 $t' \ge \max(t, 0)$ 都将保持不变。在第t次循环中，唯一重要的嵌入更新发生在第n+t个位置，因为在更早的位置没有更新发生，而在更晚的位置n+t'（对于某个t'>t）的更新最终将在未来的循环t'中被某个与当前循环t的值无关的值覆盖。最后，我们知道定义B.4中的嵌入 $x_i^{(T)}$ 与CoT Transformer中的完全相等，最终输出也一样。

第二步证明。对于第二个声明，因为CoT掩码可以由一个 $O(\log(n+m))$ 宽、常数深度的MLP计算（引理B.11），再结合引理B.8，我们知道，通过将每个Transformer块的层数以及嵌入和隐藏维度增加常数，就足以用一个没有掩码的Transformer来模拟带掩码的Transformer。□

定理 B.12 (定理4.2的重述, Sanford等人【54, Clayton Sanford et al., Transformers, parallel computation, and logarithmic depth, arXiv 2024b】)。p-跳问题（定义A.1）可以由一个 $\lfloor \log_2 p \rfloor + 2$ 层的非循环transformer解决，该transformer具有 $\log n$ 位的精度，最多3个不同的层，嵌入大小 $d = d_{FF} = d_{ATTN} = O(1)$，MLP和注意力的隐藏维度，1个注意力头。