GPT-6 Astra背后的Loop Transformer：循环架构原理与工程落地解析

过去一段时间，OpenAI 的 Astra 引发了极高的关注。有说法称这一代模型采用了 Loop Transformer 架构——信息不再一层层往前推，而是在同一组 Transformer 层里反复循环。虽然这些消息未经官方证实，但 Loop Transformer 迅速成了圈内讨论的热点。

这套做法的核心其实就一句话：让同一份权重反复跑 R 遍，用更少的参数换更深的计算。

接下来展开讲讲它的原理、工程落地难点，以及 OpenAI 是否真的用上了这个架构。

一、Loop Transformer 的原理

Loop 按循环粒度可以分成四种：整模型循环（从 embedding 到输出整体重跑）、层块循环（共享一叠 Transformer 层反复调用）、逐层循环（每层各自循环后再接下一层），以及给循环步加条件信号的潜变量条件化。下面以层块循环来展开。

设输入序列经过 embedding 是 $H^{(0)} = E = (e_1, \dots, e_n)$。循环单元 $F_\theta$ 是 N 层共享权重的 Transformer block。整个循环过程为：

其中 $R$ 是循环步数。

对比标准 Transformer 的前向过程：

差别就在于：标准 Transformer 是“层叠”，N 层独立参数；Loop Transformer 是“循环”，同一组参数跑 R 趟。unique 参数和 effective depth 在循环结构里被解耦了，这正是整套设计的出发点。

同一份参数跑 R 遍，这种结构能在参数量显著更小的情况下逼近更深的标准模型。字节的 Ouro 2.6B 在多数 benchmark 上打到了 8B 标准 Transformer 的水平。机制不复杂：约束传播、距离更新、多步算子这类迭代型任务，本质上是同一个变换被反复施加，共享权重的 R 次循环正好同构于这个过程——参数只付一份，深度按 R 复用。代价是循环单元不再是 N 层独立 block，一份权重要同时承担全部 R 层的角色。

循环块每跑一遍，变换函数不变，变的是输入。第 1 遍的输入是 embedding，第 2 遍的输入已经带上了第 1 遍的结果——同一个函数叠 R 次，近似等效于 R 层普通网络，而参数只占一份。但这个近似并不精确：weight-tying 会限制它能表达的函数类，后文会展开这层限制。这个做法最早可以追溯到 2018 年的 Universal Transformer。UT 本来就是深度方向的循环——同一个 Transformer block 反复跑，配的是 ACT 按 token 自适应停机，循环本身不在序列位置上。

架构上基本都是 Prelude + Recurrent Core + Coda 三段式：

• Prelude（非循环）：开头几层独立参数，把 token embedding 准备好。
• Recurrent Core（循环核心）：中间几层权重共享，反复跑 R 次。
• Coda（非循环）：末尾几层独立参数，把循环后的 hidden state 映射回词表。

Prelude 把 token 投到隐藏空间，Coda 把隐藏空间解回词表，中间那段循环块在同一个抽象层上来回磨。中间那段就是模型反复“思考”的地方，前后两段是硬功夫。硬功夫让独立参数承担，反复打磨让共享权重承担。这种切分既给了足够的灵活性，又把权重共享的约束限制在最通用的中间段。

训练语料里代码和数学占了相当比重，模型学到的 latent 迭代也就主要作用在推理类 token 上——这类 token 的收益随循环步数上升持续增加，整体呈现凹函数（上升、饱和，部分任务还会倒退）的形态。语言建模类任务则早早饱和甚至倒退。纯语言任务几乎不涉及多步逻辑推演，循环块多出来的算力会把已经收敛的 hidden state 继续往前推，推过最优点。

所以，循环结构的战场在推理类任务，不在语言建模。

那么循环结构有没有“理论上限”呢？TMLT 论文提到 weight-tying 的逼近误差同时受 sequence / contextual / token 三种连续性约束，而标准 Transformer 只受其中一种。这是 weight-tying 的“数学代价”。简单来说，标准 Transformer 每一层可以独立变换 hidden state，循环结构要求每一遍循环都做“几乎相同”的变换，这种约束让循环结构在某些抽象函数类上无法逼近到任意精度。

为此，TMLT 采用 timestep encoding 来解决：给循环步 $r$ 取一个 sinusoidal encoding $e_r \in \mathbb{R}^d$，作用方式类似 positional encoding，但索引的是循环步而不是序列位置；再过一层 2 层 MLP，映射成 RMSNorm 的增益 $\alpha(r)$ 和残差缩放系数 $\gamma(r)$。这让循环结构从“严格固定的操作”变成“按循环步数微调的操作”，等效扩展了表征空间。

这个缩放因子让循环块的“行为”在不同循环步上可以微调，但权重本体不变——等价于“权重共享 + per-step 微调”。要说清这是一笔交换：$\alpha(r)$ 和 $\gamma(r)$ 本身就是一组 per-step 参数，严格意义上的完全共享被放松了，换来的是表征空间变大。从信息论角度看，per-step 微调还给循环块引入了 $\log_2(R)$ 比特的信息（每个 token 知道自己在循环的哪一步），这让它能把“第 1 遍特征提取”和“第 2 遍特征组合”区分开，是 TMLT 绕开 weight-tying bottleneck 的关键。

在 Sudoku 9×9 任务上，6 层标准 Transformer 准确率 0%，Loop Transformer + TMLT（32 圈）跑出 90.2%。差别在串行计算深度：Sudoku 的约束传播要反复上百次变换才能收敛，而 6 层标准 Transformer 的串行深度只有 6 次——不管宽度加多大，它在结构上就只有 6 步串行变换。循环 32 圈把串行深度拉到 6 × 32 = 192 次，才够覆盖这道题需要的过程，这是质变的来源。

Edit Distance（length 60）这条线拉得更开：loop=4 只有 26.6%，loop=100 到 88.3%。这一类实验才是 Loop Transformer 的卖点——不是 MMLU 上几个百分点的提升，而是在标准 Transformer 结构上做不了的任务里让它能做。

二、自适应退出门：让循环步数跟着问题复杂度走

前面默认循环步数 R 是固定的。但 R 固定意味着所有输入都跑同样多的循环，简单的“今天天气不错”和复杂的“证明费马大定理”消耗一样的算力。

字节 Ouro 给出了一个解法：entropy-regularized halting gate。

每个 token 在第 $r$ 步结束时把隐藏状态过一个 sigmoid，输出退出概率 $p_r \in [0, 1]$。训练目标为：

第一项让每一遍循环都参与语言模型损失；第二项是与均匀先验之间的 KL（等价于熵正则），作用是拓宽退出步数的覆盖、避免门控坍缩到某个固定深度——它管的是探索，不是难度自适应，真正的自适应要靠后面的监督训练。论文里所有评测都在 $T_{max} = 4$ 的循环深度下做，性能随 $T = 1 \to 4$ 上升、在训练深度处达峰，外推到 $T > 4$ 反而略降。

Ouro 的训练分五段：Pre-train I 用 3T tokens 跑 8 圈，Pre-train II 再跑 3T 但把循环降到 4 圈（这一步叫 upcycle），接着 CT Annealing 1.4T、LongCT 20B、Mid-Training 300B，合计 7.72T tokens。先 8 圈、再降到 4 圈重训这个安排，是 Ouro 最终把循环数定在 4 的原因之一。CT Annealing 本身是“换更高质量的数据 + 把学习率退火到 3e-5”的持续训练阶段，强化的是数学、代码与通用能力，不负责训门控。

门控训练分两阶段。Stage I 在预训练里跟语言模型联合优化，用熵正则把退出分布的覆盖铺开；Stage II 把 LM 完全冻结，只训门控。Stage II 的监督信号来自相邻两步的 loss 改进量 $I$：把 $I$ 过一层 sigmoid 得到理想继续概率 $w = \sigma(k \cdot (I - \tau))$，其中 $k = 50$、$\tau = 0.005$，再用这个软目标做加权交叉熵。改进量大于 $\tau$ 就倾向继续，远小于 $\tau$ 就倾向退出。Stage II 训完的门控在相同计算预算下能比未训练版本再拿 2-3 个点，MMLU 上平均 2.5 步到 66%，跑满 4 步 67.35%，差 1.35 个点，省约 37% 算力。

这套信号是按“每多跑一步值不值”算的，不是按“此刻该不该退出”算的——两者互为镜像：改进量还大就继续跑下一步，改进量掉到阈值以下就该退出。所以门控学到的行为是贪心地继续循环，直到 loss 改进量低于阈值再退出。它跟 RLHF 里的奖励塑形不是一回事：这里是给门控造了一个连续的软目标，本质更接近自蒸馏。

门控先验的选择也踩过坑：均匀先验在 776M 模型上一致优于 PonderNet 那套几何先验。几何分布把概率质量压在浅层，深层循环分不到梯度信号，训练效果反而差；均匀先验给每一步等量的优化压力，收敛更平滑。

这套退出门跟早年的 Adaptive Computation Time 有点像。ACT 在 RNN 时代就用一个标量 sigmoid 决定每步停不停，PonderNet 把它概率化、加了几何先验，Ouro 又把先验换成最简单的均匀分布——三代做下来，结论是越简单的先验越稳。

Huginn 走的是另一条路，不训退出门，而是训练时随机采样循环次数（log-normal-Poisson），让模型学会“在任意步数下都能给出合理答案”。推理时默认给一个 num_steps 就行；官方实现另外提供了一组免训练的启发式退出判据——熵差、latent 差分、masked 分布 KL、argmax 稳定性，四种都能在推理时现算退出时机。

Ouro 和 Huginn 的差别在于自适应信号从哪来。Ouro 是训练出来的门控，模型自己判断该想几步；Huginn 是推理时用启发式规则现算，想几步由部署方按预算决定。前者对端侧部署更友好（能省算力），后者对 API 成本控制更友好（延迟更可预测、更好做容量规划）。这两种取向后来分成了 Loop Transformer 的两条路。

三、工程变体：循环结构怎么落地

循环结构在工程上带来三类额外开销：延迟线性涨（推理侧，未做并行优化的朴素情形）、KV cache 线性涨（推理侧）、激活显存随 R 线性涨（训练侧）。循环块反复调用同一份参数，反向传播要在 R 个时间步上累积梯度，R 一大就容易出梯度尖峰。Huginn 只用 truncated BP 回传最后 8 步，主要是为了压住训练时的显存和计算量。

激活显存是另一半。循环结构的中间激活要按 R 份保存才能反传，同一份权重只存一份——显存里激活项随 R 线性涨，权重项几乎不动。标准 Transformer 训练时激活项也常常吃大量显存，循环结构只是把这种压力按 R 倍放大，R 一大就成主要瓶颈。变通办法是 loop 内部做 gradient checkpointing，把 R 步的激活拆成若干段重算，代价是多花一次前向的算力。

字节 Seed 的 PLT（Parallel Loop Transformer）用 Cross-Loop Parallelism（CLP）打破循环的顺序依赖。做法是按对角线调度：解码 token $i$ 时，把“token $i$ 的第 1 步 + token $i-1$ 的第 2 步 + token $i-2$ 的第 3 步”拼成一个 micro-batch 一次性 forward。同一批里每一项依赖的输入，都在前面的波前里算完了，所以彼此没有数据依赖，可以并行——这跟卷积里做 wavefront 并行是同一个思路。

要做到这一点，训练侧有个关键技巧：把循环间状态沿序列方向右移一位再叠加 embedding（$H^{(r-1)} \leftarrow \text{concat}(0, H^{(r-1)}[:-1])$），消除“同位置、不同循环”之间的直接依赖，让并行成立。这个右移看似简单，本质是把循环的串行依赖解耦成“按 token index 错位的并行”。

为什么多塞几路并行几乎不涨延迟？因为 LLM decoding 是 memory-bound——瓶颈在显存带宽而不在算力。算力本来就没被吃满，多出来的并行正好被利用上，墙钟时间几乎不变。

PLT 还用了两层 KV 优化：KV 共享（所有非首循环用首循环的 KV cache）+ G-SWA 门控滑动窗（非首循环在共享 KV 之上叠加 w=64 局部 attention）。

（PLT 论文 Table 2，680M 激活 Seed-MoE，batch size=4）：

朴素循环让精度涨 5 点，代价是延迟和 cache 各翻一倍。PLT-3 把精度推到 40.8，延迟只比 vanilla 多 0.1ms、cache 只多 3M——循环结构的延迟与 cache 成本基本回到 vanilla 同等水平，精度红利还留着。

整套设计对于想自己上 Loop Transformer 的工程团队来说，CLP 是必做项，否则延迟直接乘 R；KV 共享配合 G-SWA 是配套项，否则 cache 撑不住。

PLT 各结构的复杂度拆分是这样的：

PLT 的迁移实验也给了数据：把 PLT 应用到 1.7B 激活参数的 MoE 模型上（vs 2.5B 激活参数 baseline MoE），平均精度反超 0.5 点，延迟低 30% 左右，KV cache 是 baseline 的 2/3。

Rice + Apple + UCSC + CMU 的 LT2 走另一条路：循环结构 + 线性注意力或稀疏注意力。LT2-linear 把 quadratic softmax attention 换成线性注意力（GatedDeltaNet），从 $O(n^2)$ 降到 $O(n)$；LT2-sparse 换成稀疏注意力，每个 token 只看局部邻居。

论文里有个关键发现：循环跟线性/稀疏注意力有协同效应。线性注意力的记忆会随时间衰减，而循环的 R 次重访把同一份信息反复拉回来，衰减被抵消掉；稀疏注意力只有固定窗口，每一圈循环都在窗口之外补进一层新信息，等效感受野随圈数被撑开。

把三条线放在一起看，各自啃的问题不一样：

Loop Transformer 在端侧、移动端、consumer GPU 上的部署门槛比 MoE 低很多。MoE 需要 all-to-all 通信和复杂的 expert parallelism，循环结构不需要。

四、Loop Transformer 在更大的架构图里站哪

Loop Transformer 跟 MoE 是互补的。 MoE 扩展“知识容量”（不同专家处理不同类型知识），Loop Transformer 扩展“算力迭代次数”（同一份权重反复打磨）。Ouro 1.4B 对 1.5B 激活 MoE（4B 总参数）的对照：Ouro 在数学/推理任务上反超 4 个点，MoE 在事实知识（TriviaQA / NaturalQS）上反超 6 个点。这说明在同样的参数预算下，循环更适合推理、MoE 更适合存知识。

跟 SSM 的关系也可以互补。Mamba 这类状态空间模型把注意力换成线性递推，省的是序列长度的二次项——本质上是把“看多远”做便宜；循环结构省的是参数，把“想多深”做便宜。LT2 那类工作已经证明了两者可以叠加，就是把循环块里装线性注意力，长序列和深推理两个目标一起要。

Loop Transformer 和 o1/R1 事实上都在做 test-time scaling，区别只在反复的载体。o1/R1 让模型在“输出序列”维度反复，每一步都写成文字；Loop Transformer 让模型在“隐藏状态”空间里反复，整个过程不落成 token。

五、OpenAI 有没有用 Loop Transformer？

不太清楚。目前只有 The Information 一家的报道，加上 OpenAI 首席提过一句：前沿模型（含 Astra）的计算图深度介于 GPT-4 的一半到两倍之间。这句话本身推不出 recurrent depth——计算图深度可以由很多别的原因造成。不过这个架构这么火，OpenAI 选型没有理由忽略它。但 Astra 的能力提升里有多少能归因于这个架构，这就不好说了，我觉得没有大家想象的那么高。

六、写在最后

参数效率的范式正在分叉。过去几年“提升模型能力 = 增加参数”，于是有 MoE、有 scaling law；Loop Transformer 给出另一条路——让现有参数多跑几趟。这条分叉不会终结 scaling law，两者管的不是一回事：scaling law 描述算力投进去能换来多少能力，循环结构管的是在某个能力水平上用什么代价换到。它们是互补的，不是替代的。

循环结构的“理论上限”到底在哪？TMLT 给了一种逼近误差分析，Ouro / Huginn 展示了有限步数的实际收益，但“循环越多越好”还是“到某个值就饱和”，两个方向的证据都有。TMLT 那类算法任务属于前者：Sudoku 要 32 圈才出来，Edit Distance 到 100 圈还在涨。Huginn 和 Ouro 的通用语言任务属于后者：Huginn 的收益曲线是凹的，饱和点随任务变化；Ouro 在训练深度 4 圈处达峰，外推到 4 圈以上反而略降。天花板的高度跟任务类型绑定，不存在一个通用的最优圈数。

还有一个问题是“循环结构的归纳偏置”。为什么 weight-tying 在推理任务上有效，在语言建模上不一定？一种可能的解释是：算法推理的每一步都在更新同一个数据结构（最短路径的距离、约束传播的标记），hidden state 反复磨这个数据结构的过程跟算法本身的迭代同构——循环结构匹配问题结构。语言建模则不同，每一步预测下一个 token 更像“从分布里采样”，没有那个可以被反复精化的中间对象；反复迭代会把 hidden state 推离训练时见过的状态区域，一旦越界，后面的变换没有纠偏机制，误差被逐圈放大。这个差别可能指向 Loop Transformer 真正的适用边界——它不是“通用更强的架构”，而是“在推理结构化任务上有专门优势的架构”。

这条线上的空白和问题还挺多的，做架构研究的话有太多方向可以细挖了。权重共享与 per-step 微调之间的取舍、自适应退出门的先验设计、循环与线性/稀疏注意力的组合方式，每个点都值得深入做下去。对 LLM 架构演进感兴趣的读者，值得持续关注这条线的后续工作。