T²MLR: Transformer with Temporal Middle-Layer Recurrence
Authors: Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora
Affiliation: Princeton Language and Intelligence, Department of Computer Science, Princeton University
arXiv ID: 2607.15178v1
Published: 2026-07-16
Pages: 34
Code URL: Not explicitly provided in the paper
PDF: 20260719_transformer_temporal_middle_layer_recurrence.pdf
1. 研究摘要 (Research Summary)
大型语言模型(Large Language Models, LLMs)虽然在数学、物理与科学推理任务中展现出惊人能力,但其底层 Transformer(Vaswani et al., 2017)的推理过程仍被一种根本性的信息瓶颈所束缚:自回归解码(autoregressive decoding)在每一步都把丰富的高维隐态表示压缩成词元空间(token space)中的稀疏 one-hot 向量,而下一个前向传播只能以这个离散表示作为唯一输入。于是,在前一个词元上经过多层 Transformer 计算得到的抽象中间状态,无法直接为当前词元的早期层所利用;模型若想保留任何“思考过程”,只能将其重新编码为可生成的文本,或者通过注意力机制间接地、费力地重建。这种“深度—时间壁垒”(depth-time barrier)正是限制多步隐式推理(latent reasoning)持续演化的核心障碍。
T2MLR(Transformer with Temporal Middle-Layer Recurrence)的出发点是:我们不必让模型把推理状态吐回词元空间,也不必将所有层反复循环以增加深度;更有效的方式是让抽象中间表示在解码步骤之间直接传递。近年来大量机制可解释性研究(Tenney et al., 2019; Geva et al., 2021; Meng et al., 2022; Saunshi et al., 2024; Atanas & Liu, 2025)表明,Transformer 的中间层(middle layers)是抽象推理最活跃的区域,早期层主要处理词汇与句法,后期层则倾向于把表示投影到输出词表。因此,将循环路径(recurrent pathway)嵌入中间层,而不是放在词嵌入层或最终层,是释放隐式推理能力的关键洞见。
本文的主要贡献可以概括为四个方面。第一,提出了 T2MLR 架构:一个轻量的循环路径将前一词元在较深层的表示,通过可学习的门控融合模块(gated fusion module)直接注入当前词元的较浅层。第二,设计了门控融合与循环缓存更新机制,使模型在保留标准自回归接口的同时,能够跨时间步持续演化抽象推理状态。第三,引入了基于 Jacobi 迭代的时间并行近似训练方案(temporal-parallel approximated training),在可扩展的序列并行训练框架下实现稠密教师强制(dense teacher forcing)预训练,这是连续隐式推理变体中首次实现此类训练。第四,展示了 T2MLR 可以从头预训练,也可以通过继续微调(continued finetuning)直接嫁接到已预训练的 Transformer 上,显著降低实际采用门槛。
实验结果支持了这些设计选择。在需要同时进行状态追踪(state tracking)与上下文检索(in-context retrieval)的 S5-Retrieval 合成任务上,T2MLR 显著优于单独的 LSTM 与 Transformer 基线。在 FineWeb-Edu 上的自然语言预训练中,T2MLR 在参数量与推理计算量匹配的条件下取得更低困惑度(perplexity),并在零样本 NLP 基准上持续进步;最引人注目的是,仅循环 20% 中间层(T2MLR(13,18))的配置往往优于全层循环。在 GSM-Aug、ProsQA-Hard、HotPotQA-Simple 与 Variable Assignment 等多跳推理任务上,中间层循环变体同样一致超越基线。将 T2MLR 嫁接到 SmolLM2-1.7B-Instruct 上,仅用一个 epoch 的 OpenMathReasoning 继续微调,就把 GSM8K 准确率从 35.78 提升到 39.88,MATH500 从 12.80 提升到 18.00。这些结果说明,有效的隐式推理不必像先前工作那样循环全部层,而可以通过精准定位中间层循环来更高效地涌现。
从更广阔的视角看,T2MLR 的意义在于重新框定了“循环在 Transformer 中应该放在哪里”这一基础问题。它表明,推理增益不仅来自“让模型循环”,更来自“在正确的深度让模型循环”。由于推理时额外开销不超过约 8%,且可 retrofit 到现有模型,T2MLR 为 latent reasoning 架构的实际落地提供了一条更可行的路径,也可能启发未来研究去探索更细粒度的循环位置、融合机制与跨层信息流。
2. 理论框架 (Theoretical Framework)
2.1 知识脉络与问题根源
理解 T2MLR 的理论背景,需要同时把握两条研究线索。第一条线索关注 Transformer 的推理瓶颈。自回归语言模型的核心机制是:生成第 t 个词元时,模型以之前生成的离散词元序列作为输入,逐层计算隐藏表示,最终预测下一个词。然而,当模型需要多步推理时,它必须把中间结论“外化”为文本,或者把高维隐态压缩进词嵌入。这种压缩造成了信息损失,使得连续的抽象状态难以跨时间步累积。Hao et al. (2024) 的 Coconut 和 Shen et al. (2025) 的 CODI 等连续隐式推理方法尝试把推理过程完全移到连续隐态空间,但它们通常作用于最终层或词嵌入层,把循环信息限制在输出词表附近。
第二条线索来自计算深度与循环Transformer的研究。Universal Transformers(Dehghani et al., 2019)和近年来的 looped transformers(Saunshi et al., 2025; Geiping et al., 2025; Zhu et al., 2025b)通过在同一层上多次循环来增加有效深度,确实提升了推理能力,但推理代价随循环次数线性增长。T2MLR 试图把这两条线索结合起来:既保留 Transformer 在上下文检索上的优势,又引入跨时间步的循环状态传播,同时把循环放在计算开销最小、抽象推理最丰富的位置——中间层。
2.2 核心概念:深度—时间壁垒、中间层抽象与循环状态
T2MLR 的理论基础建立在三个核心概念上。首先是“深度—时间壁垒”。标准 Transformer 中,浅层无法直接读取前一词元深层产生的表示;信息必须沿着时间轴(通过词元序列)和深度轴(通过层层前向传播)两次传递。每一次传递都伴随着信息损失或延迟。其次,“中间层抽象”指 Transformer 的中间层被认为负责更高层次的语义与推理操作,而非低层词汇模式或高层词表投影。第三,“循环状态”是一个跨时间步持续存在的隐变量,它承载了前一词元的推理结果,并能在当前词元的早期层被调用。
T2MLR 的洞察在于:如果循环状态被放置在最需要进行抽象合成的位置,它就能最大化地促进推理,而无需承担全层循环带来的推理代价。这与把循环信息放到词嵌入层的方法形成鲜明对比——词嵌入层距离抽象推理层还很远,信息必须经过早期层再次向上传播,难以被高效利用。
2.3 数学形式化
设标准 L 层仅解码 Transformer 的隐藏维度为 d,第 t 个词元的输入嵌入为 h_t^(0)。每一层 ℓ 可以看作一个映射 F_ℓ,它接收当前词元表示和过去词元的 KV 缓存,返回更新后的表示与缓存:
这里 h_t^(ℓ) ∈ ℝ^d 是第 t 个词元经过第 ℓ 层后的表示,KV 缓存存储了历史词元的键与值,保证自注意力的高效计算。这一公式刻画了标准 Transformer 的“逐层、逐时间步”计算范式:每个词元的表示只沿深度方向传播,而时间方向的信息只能通过 KV 缓存和词嵌入间接传递。
为了打破深度与时间的壁垒,T2MLR 引入两个层索引 ℓ_start ≤ ℓ_end 和一个可学习的融合模块 Φ: ℝ^d × ℝ^d → ℝ^d,同时维护一个固定大小的循环缓存 R_t ∈ ℝ^d。在第 t 步,模型把来自前一步的 R_{t-1} 与当前词元进入 ℓ_start 之前的表示 h_t^(ℓ_start−1) 进行融合:
这个改动意味着:当前词元在较浅层就能直接“读取”前一词元在较深层的抽象状态。R_{t-1} 扮演了跨时间步的“推理记忆”角色,而 Φ 则负责决定如何把这种记忆与当前输入整合。对于 ℓ_start 以外的层,计算与标准 Transformer 完全相同;在通过 ℓ_end 之后,模型基于 h_t^(ℓ_end) 与 R_{t-1} 更新循环缓存,为下一步做准备。
融合模块 Φ 的具体形式是一个门控残差结构:
其中 f_cur 与 f_rec 是可学习的线性层,分别对当前表示与循环缓存进行调制;W_rec 将循环缓存投影到与当前表示相同的空间;σ 是逐元素 sigmoid,提供输入相关的局部调制;γ_cur 与 γ_rec 是可学习的标量门控,在训练初期被初始化为零,使 Φ 近似恒等映射,从而保证优化稳定性。这个设计把“无条件残差连接”与“条件门控调制”结合在一起:模型始终保留标准 Transformer 的默认计算路径,仅在需要时通过循环分支注入额外信息。
循环缓存的更新采用 RMSNorm 包裹的残差形式:
这里把当前词元在 ℓ_end 层的深层表示与旧缓存相加并归一化,使得循环状态能够平滑地累积和演化。RMSNorm 的使用保持了数值稳定性,同时让模型学习在哪些维度上保留、更新或遗忘信息。
2.4 理论联系、假设与适用范围
这些数学组件之间存在清晰的逻辑依赖关系。方程 (2.2) 定义了跨时间步的信息注入位置,方程 (2.3) 决定了注入的方式,方程 (2.4) 规定了状态的演化规则。三者共同把“时间”维度显式引入深度维度,使模型在每个解码步骤都能利用前一步的抽象状态,而不破坏自回归接口。
T2MLR 的理论假设包括:第一,中间层确实承载了最具推理价值的抽象表示;第二,跨时间步的循环信息应直接注入较浅层,而非词嵌入层;第三,门控残差融合足以在保持训练稳定的同时,让模型学会何时利用循环记忆。这些假设在附录中得到了进一步支持:例如,未来词元预测探测(future-token prediction)显示,T2MLR 的中间层表示对下一个词元的预测能力优于标准 Transformer,表明循环路径确实促进了前瞻性表征学习。
然而,该理论框架也有边界。它主要面向短上下文内的多跳推理,而非长上下文记忆(这是 state-space models 与 recurrent memory 方法更关注的场景)。此外,它假设循环缓存的固定维度 d 足以编码所需的推理状态,对于极其复杂的多分支推理可能需要更大的状态空间。最后,时间并行训练方案本质上是一种近似,虽然实验显示其与精确循环训练几乎无异,但在极端长序列或强非线性动态下,近似误差仍可能放大。
3. 技术架构 (Technical Architecture)
3.1 系统概览与数据流
T2MLR 的技术系统可以被视为一个标准仅解码 Transformer 的“最小侵入式”扩展。它的主体仍是 Llama 风格的 Transformer 块堆叠,包含多头自注意力、前馈网络、RMSNorm 与残差连接。变化仅发生在中间层的某个局部区间 [ℓ_start, ℓ_end]:在该区间之前,模型插入一个循环融合模块 Φ;在该区间之后,模型提取一个循环缓存 R_t,供下一步使用。其余所有层保持不变,KV 缓存机制也照常工作。
数据流可以这样描述:当模型准备处理第 t 个词元时,输入嵌入首先经过早期层(1 到 ℓ_start−1),得到当前词元的浅层表示 h_t^(ℓ_start−1)。与此同时,系统从缓存中读取前一步的 R_{t-1}。融合模块 Φ 把这两个信息源结合,生成 h_t^(ℓ_start−1) 的“增强版本”,然后送入第 ℓ_start 层。此后,增强后的表示继续向上通过中间层 ℓ_start 到 ℓ_end,期间与历史词元的 KV 缓存进行标准自注意力交互。在 ℓ_end 层之后,模型用当前深层表示与旧缓存更新 R_t。最后,表示继续通过后期层(ℓ_end+1 到 L),输出词元分布。整个过程的标准自注意力、位置编码、输出头均不受影响。
这种数据流的关键在于,循环信息不是以文本形式回到输入端,而是以高维连续向量的形式直接“短路”到中间层的早期位置。可以把 R_t 想象成一种隐式的工作记忆:它不占据词元序列中的位置,也不受词嵌入矩阵的约束,却能在每个新词元进入抽象处理阶段时被立即调用。
3.2 组件交互:Transformer 块、融合模块与循环缓存
在标准 Transformer 中,各层之间是严格前馈的:第 ℓ 层的输出作为第 ℓ+1 层的输入。T2MLR 引入了两条新的交互路径。第一条是“深度短路”:从 ℓ_end 层到 ℓ_start 层,当前词元的深层表示可以影响下一个词元的浅层表示。第二条是“时间循环”:R_t 从当前步骤传递到下一步,形成一个跨时间步的循环依赖。这两条路径通过融合模块 Φ 相交:Φ 的输入一端来自当前词元的浅层流,另一端来自前一步的循环流;其输出则进入 Transformer 的主干。
门控融合模块本身也是一个小型网络。它首先把当前表示与循环缓存拼接起来,通过两个线性层 f_cur 与 f_rec 分别产生当前门与循环门。sigmoid 函数使这些门具有输入相关的局部调制能力,而 tanh(γ_cur) 与 tanh(γ_rec) 则提供全局、可学习的尺度。初始化时 γ 均为零,因此早期训练阶段 Φ 近似恒等映射,T2MLR 的行为几乎等同于标准 Transformer。随着训练进行,门控参数逐渐分化:实验观察到 γ_rec 趋向正值,γ_cur 趋向负值,意味着模型学会了把循环信息作为对残差路径的“加法修正”,而非与当前输入竞争。
循环缓存的更新机制同样经过精心设计。它采用残差相加而非直接替换,这意味着前一步的状态不会突然被完全覆盖,而是与当前步骤的深层信息平滑融合。RMSNorm 则防止状态范数随时间步累积而爆炸或消失。这种设计借鉴了现代 RNN 与状态空间模型中的门控/归一化思想,但把它嵌入到了 Transformer 的中间层残差流中。
3.3 关键创新:为何放在中间层
T2MLR 最重要的技术创新在于“循环位置”的选择。先前的连续隐式推理方法通常把循环信息注入输入嵌入层,这相当于让前一步的结论以“新词元”的形式重新进入模型。但输入嵌入层距离抽象推理层还很远,信息必须再次经过大量早期层向上传播,容易在传播过程中被稀释或扭曲。相反,T2MLR 把循环信息直接送到中间层的起点,使前一步的抽象状态能够立即参与当前词元的抽象推理。
另一个关键创新是“局部循环”优于“全层循环”。实验一致表明,仅循环网络的 20% 中间层(例如 ℓ_start=13, ℓ_end=18,D=6)往往比从第 1 层循环到第 30 层(D=30)表现更好。这说明全层循环可能让循环状态承担过多职责——它既要编码抽象推理,又要编码词汇、句法和输出投影信息,反而造成干扰。把循环限制在中间层,让状态专注于抽象推理,可能是更高效的设计。
3.4 实现细节与训练算法
在实现上,T2MLR 的超参数主要包括循环边界 (ℓ_start, ℓ_end)、Jacobi 前向深度 d_forward 与反向截断深度 d_backward。论文默认使用 d_forward=16 与 d_backward=4,在近似质量与训练效率之间取得平衡。循环深度 D = ℓ_end − ℓ_start + 1 决定了循环块的大小;D 越小,训练和推理开销越低,但状态表达能力也越受限。
推理时的算法非常简洁:维护一个缓存 R_t,每生成一个新词元,只需先计算早期层,然后调用一次融合模块,再计算中间层与后期层,最后更新 R_t。由于融合模块只涉及几次线性投影与逐元素操作,其额外计算量很小。实验测得,在 135M 到 1B 参数规模上,T2MLR 的逐词元推理时间仅比同参数量 Transformer 多约 4% 到 8%,且随着序列变长与模型变大,相对开销下降。
训练则更为复杂。因为循环依赖 R_t 依赖于前一步的 R_{t-1},标准 Transformer 的逐词元并行训练(teacher forcing)无法直接应用。T2MLR 采用 Jacobi 固定点迭代来并行近似所有位置上的 R。具体而言,先做一次标准前向传播,假设没有循环缓存,得到初始缓存 R^(1);然后用这个缓存与预先计算好的 h_t^(ℓ_start−1) 再次通过中间层,得到更精细的缓存 R^(2);重复 d_forward 次后,得到最终用于完整前向传播的 R^(d_forward)。为了控制梯度深度,反向传播只在最后 d_backward 次迭代上传播,类似于 RNN 中的截断反向传播通过时间(TBPTT)。
这一训练算法的价值在于:它在保留序列并行训练可扩展性的同时,允许模型学习跨时间步的循环依赖。与 exact sequential training 相比,Jacobi 近似在 135M、360M 与 1B 三个规模上的验证困惑度差异都在 0.005 以内,说明近似并未引入明显退化。这也使 T2MLR 成为首个能够在稠密教师强制、序列并行设置下完成预训练的连续隐式推理变体。
4. 实验评估 (Experimental Evaluation)
4.1 实验设计与数据集
作者的实验策略是:首先在一个合成任务上分离 T2MLR 的核心归纳偏置,然后在自然语言预训练中评估其语言建模能力,最后在多跳推理任务上测试其下游表现。这种由简到繁的设计让读者能够分别验证“状态追踪+检索”能力、一般语言建模能力以及实际推理能力。
S5-Retrieval 是一个定制的合成任务,它把 S5 状态追踪(Liu et al., 2023)与键值检索结合起来。输入首先是一个从 S5 群到字符串的随机映射字典,然后是一个动作序列;输出要求在每一步追踪累积群状态,并根据当前状态从字典中检索对应值。单独 LSTM 擅长状态追踪但不擅长检索,单独 Transformer 擅长检索但浅层网络难以进行长程状态追踪;T2MLR 被设计成能同时胜任两者。
自然语言预训练在 FineWeb-Edu 的 10B 词元子集上进行,基线采用 SmolLM2-135M(Allal et al., 2025),T2MLR 变体通过微调隐藏维度保持参数量大致匹配(约 136.4M)。零样本下游评估使用 lm-eval-harness,涵盖 ARC-Challenge、ARC-Easy、HellaSwag、OpenBookQA、PIQA、SciQ、Winogrande 等常见基准。推理任务包括 Variable Assignment(深度从 2 提升到 5)、ProsQA-Hard(节点与路径加长)、HotPotQA-Simple(由 GPT-4o-mini 生成思维链 SFT 数据)以及 GSM-Aug(符号与自然语言两种变体)。
4.2 主要结果与基线比较
在 S5-Retrieval 上,T2MLR 在序列精确匹配(sequence exact-match)和逐词元准确率上都显著优于 LSTM 与 Transformer 基线。Transformer 基线在训练分布内的长度尚可,但泛化到更长序列时迅速崩溃;LSTM 则因固定状态维度无法有效完成检索。T2MLR 则在分布内接近完美,并在分布外较长序列上保持非平凡的逐词元准确率。这直接证明了中间层循环能够同时支持两种互补能力。
自然语言预训练的结果见表 1。T2MLR 在多数基准上与基线持平或提升,平均表现最佳的是 T2MLR(13,18)(D=6),仅循环 20% 的层。全层循环变体 T2MLR(1,30) 的平均表现反而弱于中间层循环变体,验证了“局部循环优于全层循环”的核心假设。困惑度曲线也显示,除极端配置外,大多数 T2MLR 变体的验证损失低于基线。
表 1:135M 参数 T2MLR 变体在 FineWeb-Edu 10B 词元预训练后的零样本表现
| Model/Config | ARC-C | ARC-E | HellaSwag | OBQA | PIQA | SciQ | Winogrande | Average |
|---|---|---|---|---|---|---|---|---|
| Baseline | 24.74 | 44.28 | 29.81 | 30.20 | 61.53 | 60.80 | 48.46 | 42.83 |
| T2MLR(1,30); D=30 | 24.49 | 43.48 | 29.98 | 30.00 | 60.72 | 62.60 | 52.25 | 43.36 |
| T2MLR(5,26); D=22 | 23.98 | 46.13 | 30.75 | 29.80 | 60.77 | 62.80 | 51.85 | 43.73 |
| T2MLR(9,22); D=14 | 24.15 | 45.24 | 30.40 | 29.20 | 61.15 | 66.40 | 51.54 | 44.01 |
| T2MLR(13,18); D=6 | 24.23 | 45.50 | 29.95 | 31.20 | 61.70 | 64.20 | 52.17 | 44.14 |
| T2MLR(15,16); D=2 | 24.40 | 45.83 | 30.11 | 29.20 | 59.63 | 60.20 | 50.83 | 42.89 |
与 looped Transformer 和 pause-token 基线的比较进一步凸显了 T2MLR 的优势(表 2)。Pause-token×2、Full-looped×2 与 Middle-looped×3 都在推理时引入额外计算,而 T2MLR 在取得更好平均表现的同时,逐词元推理开销基本与同规模 Transformer 相当。这说明 T2MLR 把“额外计算”从推理阶段转移到了训练阶段,是一种更侧重推理效率的权衡。
表 2:与 looped 和 pause-token 基线的 135M 参数比较(10B FineWeb-Edu)
| Model/Config | ARC-C | ARC-E | HellaSwag | OBQA | PIQA | SciQ | Winogrande | Average |
|---|---|---|---|---|---|---|---|---|
| Transformer Baseline | 24.74 | 44.28 | 29.81 | 30.20 | 61.53 | 60.80 | 48.46 | 42.83 |
| T2MLR(9,22) | 24.15 | 45.24 | 30.40 | 29.20 | 61.15 | 66.40 | 51.54 | 44.01 |
| T2MLR(13,18) | 24.23 | 45.50 | 29.95 | 31.20 | 61.70 | 64.20 | 52.17 | 44.14 |
| Pause-token×2 | 24.74 | 44.57 | 29.51 | 31.60 | 60.23 | 61.90 | 50.59 | 43.31 |
| Full-looped×2 | 24.83 | 44.95 | 29.84 | 30.80 | 60.23 | 60.70 | 49.57 | 42.99 |
| Middle-looped×3 | 23.04 | 45.16 | 30.12 | 30.20 | 60.17 | 59.80 | 50.28 | 42.68 |
在下游推理任务上,T2MLR(9,22)、T2MLR(13,18) 等中间层循环变体一致优于基线,而全层循环变体的提升较小。在 Variable Assignment、ProsQA-Hard、HotPotQA-Simple、GSM-Aug-Symbolic 与 GSM-Aug-Natural-Language 五个任务上,T2MLR 都取得更高的 pass@1 准确率。尤其是在数学推理上,中间层循环带来的相对提升在更大规模与更多数据上进一步放大(表 9)。
4.3 消融、扩展与统计稳健性
作者进一步通过固定宽度消融验证了“位置本身”的重要性。在 135M 参数、10B 词元设置下,保持 D=6 或 D=14 不变,把循环块分别放在早期层、中间层和晚期层。结果见表 11:中间层位置始终最优,早期层与晚期层均明显落后。这排除了“任何循环都有帮助”的简单解释,支持了循环应放在抽象推理最活跃区域的论断。
表 11:固定循环宽度下的位置消融(135M,10B 词元)
| Model/Config | ARC-C | ARC-E | HellaSwag | OBQA | PIQA | SciQ | Winogrande | Average |
|---|---|---|---|---|---|---|---|---|
| Baseline | 24.74 | 44.28 | 29.81 | 30.20 | 61.53 | 60.80 | 48.46 | 42.83 |
| 6-Layer Recurrence: 13–18 (Middle) | 24.23 | 45.50 | 29.95 | 31.20 | 61.70 | 64.20 | 52.17 | 44.14 |
| 6-Layer Recurrence: 1–6 (Early) | 24.66 | 45.08 | 29.83 | 30.00 | 60.07 | 61.00 | 48.54 | 42.74 |
| 6-Layer Recurrence: 25–30 (Late) | 25.09 | 45.24 | 29.65 | 28.20 | 60.12 | 61.20 | 50.59 | 42.87 |
| 14-Layer Recurrence: 9–22 (Middle) | 24.15 | 45.24 | 30.40 | 29.20 | 61.15 | 66.40 | 51.54 | 44.01 |
| 14-Layer Recurrence: 1–14 (Early) | 23.38 | 44.70 | 28.59 | 29.00 | 59.25 | 59.50 | 51.07 | 42.21 |
| 14-Layer Recurrence: 17–30 (Late) | 23.98 | 44.02 | 29.97 | 31.40 | 60.28 | 61.90 | 50.43 | 43.14 |
在 361M 与 1B 参数规模上的扩展实验(表 9)显示,T2MLR 的零样本 NLP 增益保持稳定,而推理任务上的相对提升甚至更大(HotpotQA-Easy 提升约 14%–16%,GSM-Aug-NL 提升约 10%–13%)。将预训练数据扩展到 50B 词元(表 10)后,增益进一步扩大,尤其是在 361M 规模上。这些跨规模、跨数据量的结果增强了结论的稳健性。
不过,作者也诚实地进行了训练计算量匹配的比较:如果把 Transformer 基线也训练 2.24 个 epoch(对应 T2MLR(13,18) 的 2.24× 训练耗时),基线在某些零样本 NLP 任务上会超过 T2MLR。这说明 T2MLR 的优势主要体现在“固定参数、固定数据、固定推理计算”的约束下,而非训练计算最优。
5. 案例研究 (Case Studies)
5.1 S5-Retrieval:一个检验“双重能力”的合成案例
S5-Retrieval 是理解 T2MLR 工作机制最清晰的案例。输入由两部分组成:首先是一个完整的字典,把 S5 群中的每个元素 a 映射到一个长度为 K 的字符串 D(a);然后是一个动作序列 (a_1, a_2, ..., a_N)。目标输出则是交替的累积状态与检索结果:a_1, D(a_1), a_1∘a_2, D(a_1∘a_2), ..., Π_i a_i, D(Π_i a_i)。
这个任务精心设计了两个难点。第一,模型必须持续追踪一个 120 元素的置换群状态,每一步把新的群元素复合到当前状态上。这本质上是一个需要状态机的动态过程,RNN 类模型天然适合,但浅层 Transformer 需要至少 Ω(log N) 层才能通过并行关联扫描(associative scan)实现。第二,模型必须根据当前状态从字典中检索对应字符串,这依赖于对输入上下文的精确访问,而 RNN 的固定状态难以保留所有键值对。
单独 LSTM 可以追踪状态,但把状态映射到字典值的检索能力很弱;单独 4 层 Transformer 可以做检索,但状态追踪能力不足。T2MLR 在同样的 4 层 Llama 骨干上插入 ℓ_start=0、ℓ_end=4 的循环路径后,在训练分布内(N≤32)达到接近完美的序列精确匹配,在更长的测试分布(N 最大到 48)上仍保持非平凡的逐词元准确率。这个案例生动说明,中间层循环让模型把 Transformer 的上下文检索机制与跨时间步的状态演化机制融合在一个统一框架中,而不需要牺牲其中任何一种能力。
5.2 ProsQA-Hard 与 GSM-Aug:自然语言推理中的行为
ProsQA-Hard 的输入是一个由大量逻辑蕴含关系构成的有向无环图,问题要求判断某个起点节点能否到达两个候选终点节点之一。正确输出需要逐步列出一条路径。在这个任务中,模型必须维护一个不断扩展的推理链:当前已推导出的结论成为下一步推理的前提。T2MLR 的循环缓存 R_t 可以自然地把这些中间结论作为连续状态保存,避免它们被每个新词元重新计算或遗忘。实验显示,中间层循环变体在该任务上明显优于基线,尤其是在路径加长到平均 8 步的 Hard 设置下。
GSM-Aug 则展示了数学推理场景。以自然语言版本为例,问题描述为:“一家公司 600 名员工中,30% 获得晋升,10% 获得奖金。有多少人既未晋升也未获奖金?”正确解题步骤需要依次计算 600×30%=180、600×10%=60、180+60=240、600−240=360。这类问题要求模型把多个中间数量(180、60、240)保持在工作记忆中,直到最后一步使用。T2MLR 的循环状态能够把这些数量作为连续隐态保存,而不必将每个中间结果都外化为文本。这或许解释了为何 T2MLR 在 GSM-Aug 上的提升随规模与数据量增长而放大。
5.3 从案例中提取的洞察
这些案例揭示了几个共性。第一,T2MLR 的循环缓存最适合承载“需要跨多步保持并渐进更新”的抽象信息,例如群状态、推理链中的节点、算术中间量。第二,由于循环信息直接注入中间层,它不会被早期层的词汇处理或晚期层的词表投影过度干扰,从而保持其抽象性。第三,模型并非总是依赖循环——门控机制让它在不需要时退化为标准 Transformer,这种“默认标准、按需增强”的行为可能是训练稳定和泛化良好的原因。最后,这些案例也暗示了失败模式:当任务需要同时维护多个独立分支或极长程依赖时,固定维度的循环缓存可能成为瓶颈。
6. 综合价值与局限 (Synthesis — Value and Limitations)
6.1 理论意义与实践影响
T2MLR 的理论意义在于它把“循环位置”从一个工程细节提升为架构设计的核心问题。它表明,Transformer 中的循环信息应该放在抽象处理最活跃的中间层,而不是被迫退回到词嵌入层或均匀分布在所有层。这一洞见与近年来关于 Transformer 层功能分化的机制研究相互印证,也为未来设计更精细的跨层、跨时间信息…
Topics:
- "reasoning"
- "llm"
- "memory_mechanism"
- "recurrent_neural_networks"
- "reasoning_memory"
References: - "princeton_university"
- "chain_of_thought"
- "backpropagation_through_time"
- "gsm8k"
- "hotpotqa"