一句话总结把上一个 token 的中间层表示注入下一个 token 的相同区域让推理中间态跨时间步存续测试效果稳定超过标准 Transformer新增推理开销不超过 8%但训练算力需要扩增 2–4 倍论文标题T²MLR: Transformer with Temporal Middle-Layer Recurrence论文地址https://arxiv.org/abs/2607.15178作者背景普林斯顿大学代码地址https://github.com/princeton-pli/T2MLR一、动机1.1 自回归低效性标准 decoder-only Transformer 逐 token 生成文本。为了产出每一个词模型要把输入送过几十层网络每一层都会算出一个几千维的隐状态向量——里面装着模型对当前局面的理解和算到一半的推理。但到了出字环节这些高维隐状态被投影回词表坍缩成一个离散词下一步的前向计算只能从这个词的 embedding 出发重新开始这带来两个结构性问题信息瓶颈几千维的连续表示每一步都要先被压缩成一个离散的词然后下一步再把它重新转换回向量送回模型。这一路上只塞得进一个词的信息量中间形成的推理状态存活时间极短跨层信息无交互尽管注意力机制能挑选重要上下文 token实现过往信息的再加工但标准 Transformer 架构下各层 token 只能访问到相同 layer 下的缓存上下文 token 在更深层的表示虽然还在显存里但计算时读不到1.2 隐空间循环方法潜空间推理以 COCONUT 为代表让模型暂不吐字把最后一层的隐状态直接喂回输入端在连续空间里默想几轮再开口。这个思路确实绕开了单个词的瓶颈但递归信号的路径是 “最后一层 → 输入 embedding” —— 一直贴着网络的两端打转。而大量可解释性研究BERT 探针、知识定位、层堆叠分析等指向同一个结论早层处理词法句法晚层负责投影回词表抽象推理主要发生在中间层。递归信息被架构限制在网络外壳进不了真正执行抽象计算的区段循环 Transformer在同一个 token 上把整叠层原地多跑几遍等效于临时加深网络。推理能力确实提升但每个 token 的推理计算量随循环次数线性增长循环两次就是 2 倍、三次就是 3 倍对延迟敏感的线上部署难以接受从这些线路可看出大家都忽视了一个研究盲区都在争论要不要递归、递归几次却很少问 “递归应该安在哪一层”二、解决方案2.1 核心机制T²MLR 的改动可概括为选定中部偏浅的层 l_start 和中部偏深的层 l_end维护一个维度大小恒定的递归缓存 r第 t 个 token 经过第 l_end 层后用该层表示更新 r_t第 t1 个 token 传播到第 l_start 时把 r_t 与当前表示做门控融合再继续往上算其余层的计算与标准 Transformer 完全一致这样上一步在中间层算出的推理状态不需要压过词表、也不需绕到输入端而是直接送进下一步的中间层即抽象推理最活跃的区段。实验表明这样的确能获得更低的预训练困惑度与更好的下游表现2.2 门控融合递归缓存不能直接灌进当前的计算流否则会破坏模型原有的表示。融合模块给两股信息各配了两级阀门全局门tanh处理 r_t控制整条递归链路的强度局部门sigmoid处理当前输入决定其在各维度上被放行多少全局门初始化为零即训练起步时递归通路完全关闭模型行为与标准 Transformer 相同以保证训练早期的稳定性中间缓存更新时直接让新旧隐层向量相加再过 RMSNorm 即可缓存大小始终恒定让模型的参数增量不足 0.1%2.3 近似训练所有递归架构都有一个绕不开的死结串行的逐 token 计算。在 Transformer 中搞隐空间递归会让原本的并行优势荡然无存。以 COCONUT 为例在隐空间中打转推理不仅很慢还没有监督信号可用因为没有映射回语言空间这加大了预训练难度T²MLR 的解法是雅可比不动点迭代训练时先假设递归缓存不存在把整句话并行跑一遍取各位置第 l_end 层的输出在时间上平移一位当作所有位置缓存的初值再把它融合进第 l_start 层、对整句话重刷一遍中间块得到刷新后的缓存每刷一轮正确的递归信息只往后传一格若要严格还原整条递归链需迭代 length 次序列长度但这显然开销太大。停在 16 次拿到的是一个近似解换来全程序列并行。此外仅仅在中间层做递归输出层正常吐字并与监督文本对比而产生稠密学习信号对照实验表明雅可比近似训练与精确逐 token 递归的验证损失差距在 0.005 以内。在线推理时则完全回归到自回归解码流程只多一次融合模块的计算而不需再做迭代实测生成开销不超过约 8% 且随生成长度和模型规模增大而降低雅可比迭代想象一排人依次报数每个人要报的数取决于左边那位报了什么。老实的办法是从最左边一个个问过去但这样有多少人就得问多少次。雅可比的办法是让所有人先同时瞎猜一个数喊出来然后每个人只看左边邻居上一轮喊的数来更新自己大家再一起喊一遍。这样每轮所有人同时开口保持并行并且能逐步变准确。T²MLR 只做16 轮在保持并行、开销可控的前提下实现了递归结果近似三、实验结果3.1 状态跟踪与上下文检索能力在上自然语言评测之前作者先用一个合成任务把架构的归纳偏置隔离出来。S5-Retrieval 要求模型同时做两件事持续追踪一个 S5 置换群状态状态跟踪并按当前状态去上下文中的随机字典里查出对应编码上下文检索经典的 S5 状态跟踪桌上摆 5 个编号杯子总共有 5! 种摆法“S5” 指的就是这 120 种排列集合。输入是一长串洗牌指令每条指令描述一种重排方式如 “1、3 号互换、2 号挪到最后”。模型要做的是从初始摆法出发每执行一条指令就报一次杯子当前的摆法难点在于第 i 步的摆法是前 i 条指令按顺序执行的结果洗牌不满足交换律错一步则后面全错。相关理论也证明了没有能一口气算出来的捷径 —— 想精确解计算深度必须随步数增长。所以它是学界公认的 “状态跟踪” 试金石考的就是能不能维护一个随时间演化的内部状态S5-Retrieval 在此基础上加了一层检索。开题前先发一个暗号本120 种摆法各配一个 4 位编码比如摆法 32514 对应 F4Nd模型每报一次当前摆法还需要查出这个摆法的暗号这个任务的巧妙之处在于它精确打击两类架构各自的短板层数固定的 Transformer 存在推理深度上限难以拟合深层次函数LSTM 类模型受益于循环机制序列越长推理越深但中间状态的容量有限限制了上下文检索能力。实验结果也正是如此4 层 Transformer 和 4 层 LSTM 各自失败参数量匹配的 T²MLR 在训练长度内接近满分超出训练长度后 token 级准确率也保持非平凡水平上图分别展示了三种模型的测试结果。第一行子图统计 “序列级全对率”要求答案一个 token 不差才得分第二行是 “token 级平均准确率”按位置给部分分用来衡量错到什么程度。横轴表示状态数量蓝色竖线是 4 层 Transformer 理论上能精确解出的最大状态数2⁴16橙色竖线是训练时见过的最大状态数32两个 baseline 的序列级全对率几乎为 0按 token 打分也只能撑住最前面几步Transformer 实际在远没到蓝线N16时就已经崩了明显低于理论上界T²MLR 在训练长度内N≤32序列级全对率接近满分超出训练长度后全对率逐渐回落但 token 级准确率仍维持在较高水平递归通路提供了随时间增长的电路深度补状态跟踪短板注意力骨架保留了对全上下文的直接访问补检索短板这正是 “attention 递归” 混合架构想要的兼得3.2 预训练预训练主实验以 SmolLM2-135M 为基线在 FineWeb-Edu 的 10B token 官方子集上训练一个 epoch用 lm-eval-harness 测 7 个零样本基准ARC-C/E、HellaSwag、OBQA、PIQA、SciQ、Winogrande与同类隐空间计算方法相比pause-token 的推理开销随暂停数平方增长looped 随循环数线性翻倍而 T²MLR 推理几乎不加价。还有一个更优意思的现象middle-looped×3 也是在中间层多算却连基线都没超过因为它只是把中间块原地重复空间上加深与把中间态送往下一个 token时间上递归是两种不同的信息流递归位置消融实验效果如下。可见最好的配置是只递归中间 6 层20%而非所有此前潜空间推理工作默认的全层递归其次递归块也不能太窄D2 时收益几乎消失3.3 下游任务把预训练 checkpoint 在五个多跳推理/数学任务上微调T²MLR 均超过了参数与数据匹配的基线且中间层递归变体D22/14/6一致强于全层递归D30在 zero-shot 和下游任务基准上将参数规模提高至 1B、预训练参数量扩充至 50B 均能明显提高效果表明架构具备一定的规模化趋势前景3.4 旧模型改装尽管近年来学术界各类预训练架构新方案层出不穷但大都难以运用至实际场景因为这些新架构一般都另起炉灶从头训练这意味着只有预训练组才玩儿得起对实践者最有吸引力的问题是这架构必须从头预训练吗对此作者拿现成的 SmolLM2-1.7B-Instruct 模型插入递归融合通路覆盖第 5 到 28 层占全网络 75%在 OpenMathReasoning 上微调一个 epoch效果有显著提升全局门零初始化在这里派上大用场改装瞬间模型行为与原模型完全一致递归通路随微调逐渐上线不会砸坏预训练好的能力四、局限最大局限性在于T²MLR 需要更大的训练开销约是普通 Transformer 的 2–4 倍如果把训练所需的算力对齐让基线训练至 2.24 个 epoch后者便能实现反超。这也是不少新架构所面临的问题 —— 一顿操作下来是否真比简单地多训几轮、多洗一些数据更划算
