门控循环Transformer：以循环调制获得表达性深度 Gated Recurrent Transformers: Expressive Depth through Recurrent Modulation
用GRU式门控让单个共享层块在每次循环中特化，以约37%参数量匹配稠密基线质量
前置知识
Transformer 残差流与参数-深度耦合
标准 Transformer 把 $L$ 个结构相同的块依次堆叠在残差流上：$h^{(\ell)} = h^{(\ell-1)} + \mathrm{Block}_\ell(h^{(\ell-1)})$，每个块包含 pre-norm 多头自注意力和逐 token 的 MLP。由于每层都持有一套独立权重，总参数量按 $\Theta(L \cdot d^2)$ 随深度线性增长（$d$ 为隐藏维度），想加深网络就必须成倍增加存储。
本文要解决的核心问题正是这种深度与参数的刚性耦合；GRT 的目标是在不增加唯一权重的前提下获得深层表达能力，必须先理解普通 Transformer 如何堆叠层、参数量如何随深度增长。
GRU 门控机制
门控循环单元（GRU）用 sigmoid 产生的更新门 $g \in [0,1]^d$ 按元素控制旧状态与新计算的候选状态之间的混合：$h_t = g \odot h_{t-1} + (1-g) \odot \tilde{h}_t$。门接近 1 时状态原样保留，接近 0 时被新内容覆盖，这既缓解梯度消失，也让网络学会哪些维度该记住、哪些该更新。LSTM 遗忘门正偏置初始化（如 +1 或 +2）是经典实践。
GRT 的核心创新就是把 GRU 的门控从时间维（序列步之间）搬到深度维（循环步之间），门控的直觉是读懂其更新公式和初始化设计的钥匙。
权重共享与循环深度
让多层 Transformer 复用同一组权重并迭代多次，可以用固定参数量换取任意有效深度。代表工作包括 ALBERT（全层共享）、Universal Transformer（加自适应停机）、Meta 的循环深度 Transformer（prelude–共享核–coda 布局）、MoR（token 级递归路由）、RRT（用 LoRA 松弛权重绑定）和 Ouro（监督每次循环迭代）。
GRT 建立在循环深度框架之上，实验中与 MoR、RRT、Ouro、heavy-tail depth sampling 四个循环基线逐一对比，不了解这一谱系就无法理解实验设计的意图。
isoFLOPS 与 isoPARAMS 两种评估制度
isoFLOPS 指在训练/推理计算量匹配的前提下比较方法：循环模型参数更少但把层迭代多次，衡量'用参数换效率'的能力；isoPARAMS 指参数量匹配的前提下比较：循环带来更高的每步 FLOPs，衡量'用计算换质量'的能力。Kaplan 等人的 scaling law 研究曾发现循环模型在 isoPARAMS 下更好、在 isoFLOPS 下更差。
论文所有主实验（Table 1）都同时在这两种制度下报告，不理解这两个词就无法判断结果对比的公平性，也无法理解'统一评估'这一贡献的含义。
KV 缓存与自回归解码
自回归生成时，每层会缓存历史 token 的键/值张量（KV cache）以避免重复计算，缓存大小随层数、批大小与序列长度线性增长。循环深度模型每做一次循环就产生一份额外的 K/V，朴素实现会把缓存膨胀 $R$ 倍，因此需要压缩策略（只留最后一步、只留第一步、跨步平均）。
GRT 宣称的显存优势（峰值显存 −59%）很大程度取决于 KV 缓存如何处理，文中 Table 3、Table 4 和 Figure 7 专门分析了这个问题，是读懂效率部分的前提。
研究动机
标准 Transformer 中深度与参数刚性耦合：每加一层就引入一整套新权重，参数量按 $\Theta(L\cdot d^2)$ 增长，例如 GPT-2 Medium 需 24 层、354M 参数，GPT-2 Large 达 36 层、774M 参数，在固定硬件预算下想加深网络很快撞上显存墙。权重共享看似解药——同一组块迭代 $R$ 次即可用 9 个唯一块的存储执行 24 次块计算——但存在更深的架构性矛盾：第一次循环后的隐状态与第八次循环后的隐状态在表示上截然不同，静态共享强迫同一变换作用于完全不同的表示阶段，会坍缩让深度有价值的功能多样性。论文的消融量化了这一点：在 GPT-2 Small 设置下，朴素循环（无任何辅助机制）的验证损失为 3.264，比同训练配方下的 12 层稠密基线 3.157 差 0.107 nats。Kaplan 等人的 scaling law 研究也观察到循环模型在匹配参数量时更好、在匹配计算量时更差，说明此前的权重共享方法没有真正解决共享核的输入特化问题。
本文的目标是本文的目标是设计一个循环深度 Transformer，让单个共享核心在每次循环中'表现不同'，从而在两种制度下同时取得强结果：isoFLOPS 制度下用远少于稠密基线的唯一参数（约 36–37%）达到接近甚至超过的质量，并在所有规模×预算组合中领先现有循环方法；isoPARAMS 制度下在参数量匹配时通过加深循环稳定超越稠密基线。同时希望不引入任何辅助损失或路由机制，就能自然支持推理时提前退出（early exit），让单个检查点在浅层快速推理与深层高质量推理之间连续调节。工程上还要求循环引入的额外开销足够小（门控与投影仅每步 $10d^2$ FLOPs），且解码时的 KV 缓存膨胀可通过简单策略控制。
与已有工作不同的是，已有工作把循环深度当作两种割裂的工具：MoR、RRT 等用它做效率工具（参数更少匹配质量），Ouro 等用它做性能工具（参数相同提升精度），缺乏统一处理；而且这些方法共享权重时都没有让共享核针对每一步的表示状态做输入特化——RRT 的 LoRA 增量在训练时确定、对所有输入相同，买到的多样性被提前固定。GRT 的独特切入是从 RNN/GRU 借来门控思想但作用在深度维而非时间维：一个逐元素门以当前隐状态、固定的 prelude 表示和每步重采样的随机噪声为条件，构造出'每次循环都不同的、锚定原始输入的'新输入，让一个权重张量表现为多个特化层。这也区别于 ACT、Mixture-of-Depths、MoR 等需要专用停机信号或 token 路由机制的自适应计算方法。
核心方法
GRT 采用 prelude–共享核心–coda 三段架构，记作 $n_{pre}+n_{rec}\times R+n_{coda}$。前 $n_{pre}$ 个块对 token 嵌入 $h^{(0)}$ 只执行一次，产生固定的上下文条件信号 $h^{(pre)}$；中间 $n_{rec}$ 个块构成共享核心，以完全相同的权重迭代 $R$ 次，每步读入上一步状态、写回更新后的状态；最后 $n_{coda}$ 个块把精炼后的 $h^{(R)}$ 投影为最终隐藏态并送入 LM head。唯一参数量为 $\Theta((n_{pre}+n_{rec}+n_{coda})\cdot d^2)$，与 $R$ 无关。以 2+5×4+2 为例：每次前向执行 $2+5\cdot4+2=24$ 次块计算，但只存储 9 个块的权重，比同 FLOPs 的 24 层 GPT-2 Medium 少 2.6 倍唯一块。每 token 前向成本为 $\text{FLOPs}=(n_{pre}+n_{rec}R+n_{coda})(24d^2+4Sd)+R\cdot 10d^2$，其中 $24d^2+4Sd$ 是标准块成本，$10d^2$ 是循环投影 $W_{proj}$ 与门控 MLP 的每步开销。直觉上：prelude 像固定的输入编码器，共享核心像 GRU 单元在深度方向上反复更新一个'思考状态'。
核心创新是带随机性的逐元素循环门控。每步先做循环投影 $\tilde{h}^{(r)} = W_{proj}[\,h^{(r-1)}+\epsilon_x,\ h^{(pre)}\,]$，把当前状态与固定的 prelude 输出拼接后投影，其中 $\epsilon_x\sim\mathcal{N}(0,0.1^2)$ 是每步重采样的状态噪声，防止模型学到跨步的脆弱精确匹配；共享块处理得候选输出 $o^{(r)}$。门由两层 MLP（隐宽 $d$、SiLU 激活）计算：$g^{(r)} = \sigma(f_g([\mathrm{LN}(h^{(r-1)}), \mathrm{LN}(h^{(pre)})])/\tau + \epsilon_g)$，温度 $\tau=1$，门噪声 $\epsilon_g\sim\mathcal{N}(0,0.1^2)$ 防止门在训练中坍缩为常数。更新规则 $h^{(r)} = g^{(r)}\odot h^{(r-1)} + (1-g^{(r)})\odot o^{(r)}$：$g\to 1$ 时残差流原样通过，$g\to 0$ 时块输出完全接管。关键在初始化——$f_g$ 第二层偏置初始化为 +4，使初始 $g^{(r)}\approx 0.98$，训练从近似恒等映射出发、随训练逐步学会选择性覆写。与 Gemma3n 等模型的 PLE 需要 $L$ 个独立嵌入层提供逐层上下文不同，GRT 用单一 $h^{(pre)}$ 在每步提供原始输入的恒定视图，再以 $h^{(r-1)}$ 和随机性逐元素调制。
方法步骤详情
完整流程分五步。第一步，token 序列 $X$ 经嵌入得 $h^{(0)}\in\mathbb{R}^{T\times d}$，依次通过 $n_{pre}$ 个 prelude 块得到 $h^{(pre)}$，它在所有循环步中保持不变，充当原始输入的锚。第二步，进入共享核心：第一步以 prelude 输出加状态噪声 $\epsilon_x$ 初始化，之后每步计算循环投影 $\tilde{h}^{(r)} = W_{proj}[h^{(r-1)}+\epsilon_x,\ h^{(pre)}]$，把 $2d$ 维拼接投影回 $d$ 维。第三步，$n_{rec}$ 个权重相同的共享块处理 $\tilde{h}^{(r)}$ 得候选输出 $o^{(r)}$。第四步，按 $g^{(r)}=\sigma(f_g([\mathrm{LN}(h^{(r-1)}),\mathrm{LN}(h^{(pre)})])/\tau+\epsilon_g)$ 算出 $[0,1]^{T\times d}$ 的门，按 $h^{(r)}=g^{(r)}\odot h^{(r-1)}+(1-g^{(r)})\odot o^{(r)}$ 混合写回，循环 $R$ 次。第五步，$h^{(R)}$ 经 $n_{coda}$ 个 coda 块得最终隐藏态，LM head 输出 logits。训练时每个优化步从 $r\sim\mathrm{Uniform}\{1,\dots,R\}$ 采样循环深度：既隐式训练了每个退出点（实现免费 early exit），又起随机深度正则化作用。优化用 AdamW（$\beta_1=0.9,\beta_2=0.95$，权重衰减 0.1），2000 步预热后余弦衰减 $6\times10^{-4}\to 6\times10^{-5}$，梯度裁剪 1.0，bf16，20,000 步、每步约 491,520 token、共约 9.8B token。
技术新颖性
新颖性体现在四个层面。其一，门控维度的转换：GRU 的门在时间步之间控制状态覆写，GRT 把同一原理搬到深度循环之间，且门的条件里包含固定的 prelude 表示——共享块的权重固定但行为动态化，靠'每次换新输入'而非'每次同一变换'实现功能多样性。其二，随机化设计：状态噪声 $\epsilon_x$ 在每步重采样而非仅在初始化时注入，配合门噪声 $\epsilon_g$，消融显示仅状态噪声一项就贡献 0.018 nats。其三，无辅助损失的自适应计算：均匀深度采样让所有中间状态都被训练去预测最终目标，使 early exit 成为涌现能力（一半循环步保留约 92% 精度），而 ACT、Mixture-of-Depths、MoR 都需要专用停机信号或 token 路由。其四，经济性：门与投影的每步开销仅 $10d^2$，相对每块 $24d^2$ 的标准成本很小；偏置 +4 初始化沿用 LSTM 遗忘门正偏置的经典实践（Jozefowicz 等发现 +1 或 +2 即可）但首次系统用于深度维。与 RRT 的本质区别：LoRA 增量对所有输入相同、多样性在训练时就固定，而 GRT 的门以'即将被更新的状态'为条件，能在均匀共享坍缩表示之处保持状态可分。
实验结果
isoFLOPS 制度（Table 1 上半）：GRT 在全部 9 个规模×预算格中领先 MoR 和 heavy-tail depth sampling。小规模（1+1×10+1，仅 3 层、35M 参数）验证损失 3.14，优于 124M 参数的 GPT-2 Small 基线 3.15，三个种子为 3.145±0.004；同时超过 MoR（3.30）、heavy-tail（3.23）、Ouro（3.19），与 RRT（3.14）持平。中规模 2.89、大规模 2.77，落后稠密基线（2.84、2.71）0.05/0.06 nats，但领先所有循环对手：RRT 在中/大规模落后 0.06/0.08，MoR 与 heavy-tail 落后 0.08–0.16。isoPARAMS 制度（Table 1 下半）：参数匹配时 3.04/2.76/2.65 全面超过稠密基线 3.15/2.84/2.71，代价是推理 FLOPs 大幅上升（大规模 109.0G 对 21.1G）。下游评测（Table 2，大规模零样本）：isoFLOPS 版九项基准平均 42.08 对 GPT-2 Large 的 42.05（+0.03），参数仅 288M（−63%）；isoPARAMS 版平均 44.15（+2.10），8/9 任务领先，LAMBADA Standard +8.29、LAMBADA OpenAI +5.22、HellaSwag +4.31。Early exit（Figure 3）：只用一半循环步保留约 92% 精度，同推理 FLOPs 下 GRT 浅退出优于稠密早层退出。效率（Table 3/4）：torch.compile 下大规模生成延迟 3.67 ms/token（基线 3.33，+10%），峰值显存 639MB 对 1570MB（−59%）；KV 跨步平均在 B=32 把解码内存降到 0.39×（1.44 GiB），HellaSwag 反而略升（33.90 对 33.65）。消融（Figure 5）：门控是最大单一贡献者（−0.048 nats），完整模型 3.141 是唯一低于稠密基线 3.157 的配置。机制分析：早期循环写密集（门≈0）、后期复制密集（门≈1）；困难 token 的跨步改进量约为简单 token 的 10 倍。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 验证损失（小规模 isoFLOPS，1+1×10+1） | validation loss (nats) | 3.14（GRT，3 层，35M 参数） | GPT-2 Small 12L：3.15（124M）；MoR 3.30 / heavy-tail 3.23 / Ouro 3.19 / RRT 3.14 | 低于稠密基线且参数仅为其约 28%，并列最强循环基线 RRT，领先其他循环对手 0.05–0.16 nats |
| 验证损失（中规模 isoFLOPS，2+5×4+2） | validation loss (nats) | 2.89（127M 参数，9 层） | GPT-2 Medium 24L：2.84（354M）；最强循环基线 Ouro 2.93、RRT 2.95 | 领先所有循环对手至少 0.04 nats，落后稠密基线 0.05 nats，参数为基线 35.9% |
| 验证损失（大规模 isoFLOPS，1+5×6+5） | validation loss (nats) | 2.77（293M 参数，11 层） | GPT-2 Large 36L：2.71（774M）；最强循环基线 Ouro 2.82、RRT 2.85 | 领先循环对手 0.05–0.14 nats，落后稠密基线 0.06 nats，参数为基线 37.9% |
| 验证损失（中规模 isoPARAMS，2+20×4+2） | validation loss (nats) | 2.76（357M 参数） | GPT-2 Medium 24L：2.84（354M 参数） | −0.08 nats（匹配参数量，推理 FLOPs 26.4G 对 7.35G） |
| 验证损失（大规模 isoPARAMS，3+30×6+3） | validation loss (nats) | 2.65（779M 参数） | GPT-2 Large 36L：2.71；MoR 2.69 / heavy-tail 2.70 | −0.06 nats 且为该制度最优，同时优于两个循环基线 |
| 九项基准平均准确率（大规模，零样本） | average accuracy (%) | isoFLOPS：42.08（288M，−63% 参数）；isoPARAMS：44.15 | GPT-2 Large：42.05 | isoFLOPS +0.03 点；isoPARAMS +2.10 点，9 项任务中 8 项领先 |
| LAMBADA（Standard） | zero-shot accuracy (%) | 38.02（isoPARAMS） | GPT-2 Large：29.73 | +8.29 点，Table 2 中最大单项增益 |
| 生成延迟与峰值显存（大规模，batch 4，prompt 1024） | ms/token 与 MB | GRT Large：compiled 3.67 ms/tok，峰值 639MB（293M 参数） | GPT-2 Large：compiled 3.33 ms/tok，峰值 1570MB（774M 参数） | 编译后延迟仅 +10%，峰值显存 −59%；eager 模式延迟 +23% |
| Early exit（只用一半循环步） | retained accuracy | 约 92%（无任何辅助损失，单检查点） | 多退出模型通常需辅助损失训练；稠密模型同 FLOPs 早层退出质量更差 | 涌现能力：深度采样隐式训练了所有退出点 |
| 端到端解码内存（中规模 isoFLOPS，bf16，T=1024，B=32，平均 KV） | 相对 GPT-2 Medium 的内存比 | 0.39×（1.44 GiB），HellaSwag 33.90 | GPT-2 Medium：1.00×（3.66 GiB），HellaSwag 34.45；GRT 全量 R× 缓存：0.91× | 内存降 61%，且平均 KV 的 HellaSwag 反而高于全量缓存的 33.65 |
局限与改进
作者承认三点局限：其一，循环深度 $R$ 在推理时固定，没有逐 token 停机机制，所有 token 支付相同的循环计算；其二，门偏置（+4）与噪声幅度（$\sigma_x=\sigma_g=0.1$）是针对 GPT-2 家族调出的，换到其他规模或架构可能需要重调，且偏置只在 {−2, 0, +2, +4} 内做了全程训练扫描、极差仅 0.019 nats，覆盖有限；其三，最优共享比例（prelude/共享/coda 的划分）随规模变化，缺乏系统研究。我自己的观察还有：实验上限仅 774M 参数、9.8B token 训练，与当代数十亿至千亿参数级 LLM 的差距使结论外推存疑；中规模 isoPARAMS 下 heavy-tail Poisson（2.74）反而略优于 GRT（2.76），说明门控优势并非处处成立；isoFLOPS 下中/大规模仍落后稠密基线 0.05–0.06 nats，'省参数不亏质量'目前只在小规模或数据加倍的前提下成立；isoPARAMS 的质量增益以推理 FLOPs 暴涨为代价（大规模 109.0G 对 21.1G，约 5.2 倍）；KV 跨步平均虽省显存但改变解码语义，对长上下文、大 batch 的稳定性未验证；eager 模式下门控带来 +23% 延迟，说明方法对 kernel 融合等编译优化依赖较强。
独立分析的弱点
第一，统一深度、无逐 token 停机：知识型 token 常在早期就稳定（Figure 4 中'1984 作者'提示在第 2 步 Orwell 已达 83%），却被强制跑满 $R$ 步。改进方向：把门值的累积消耗作为停机信号（类似 ACT 的 halting score），门本身已在逐步度量'还需覆写多少'，天然适合改造。第二，KV 缓存 $R$ 倍膨胀是真实的服务成本：batch=32 时朴素策略只省到 0.91×；跨步平均虽到 0.39×，但本质是把 $R$ 步 K/V 压成一份，长序列下的信息损失未量化。改进方向：按步加权的缓存蒸馏或低秩跨步共享。第三，门控逐元素操作的 kernel 开销：eager 模式 +23% 延迟，约一半差距来自 kernel 启动，可把投影、门计算与混合融合为单个 CUDA kernel。第四，规模与数据域单一：全部实验基于 GPT-2 词表（50,257）、序列长 1024、约 9.8B token，数据集只描述为'多样化文本'，应先在 1–7B 参数、数百 B token 上复现 scaling 趋势。第五，中规模 isoPARAMS 输给 heavy-tail Poisson（2.74 对 2.76），提示门控收益与共享布局存在交互，应联合搜索 $n_{pre}/n_{rec}/n_{coda}$ 比例与 $R$。第六，门条件依赖 $h^{(pre)}$ 这条固定视图，若 prelude 欠训练或输入分布漂移（如超长文档后半段），锚可能失效，值得补充稳健性消融。
未来方向
作者明确提出两个方向：动态逐 token 停机（让不同 token 消耗不同循环步数，门控信号可直接复用为停机依据），以及从稠密教师模型向循环学生做知识蒸馏（缓解循环模型在低数据期的劣势）。基于本文成果还可延伸：其一，系统化共享比例研究——作者已指出最优 sharing fraction 随规模变化，可把 $(n_{pre}, n_{rec}, n_{coda}, R)$ 作为超参数纳入 scaling law 外推；其二，把门控与 Mixture-of-Recursions 的 token 级递归路由结合，实现'深度维门控 + token 维路由'的双粒度自适应计算；其三，机制研究延伸——论文已发现早期步写密集、后期步复制密集，且困难 token 的跨步改进量约 10 倍于简单 token，可在更大模型上验证这些规律并发展成可解释性诊断工具；其四，与长上下文技术交叉：循环状态天然适合做记忆压缩，可探索用 $h^{(r)}$ 状态序列部分替代 KV 缓存或做检索式记忆；其五，推理型任务验证——Figure 4 显示代码补全提示在更深循环下持续修正（'range' 置信度 21%→36%），数学/代码基准上的系统评测很有价值；其六，用稠密模型中间层监督循环步（层间对齐蒸馏），可能加速门控的分化。
复现评估
复现条件较好。代码已开源（github.com/Amr-Hegazy1/gated-recurrent-transformer），基于 nanoGPT 搭建，训练配方完整公开：AdamW 超参（$\beta_1=0.9,\beta_2=0.95$、权重衰减 0.1）、2000 步线性预热、余弦学习率 $6\times10^{-4}\to 6\times10^{-5}$、梯度裁剪 1.0、bf16、20,000 步、每步约 491,520 token（batch 8 × 梯度累积 60 × 序列 1024），合计约 9.8B token；状态噪声 $\sigma_x=0.1$、门噪声 $\sigma_g=0.1$、门温度 $\tau=1.0$、偏置初始化 +4 等关键细节均有数值。最大模型约 774M 参数、109G FLOPs/前向，中大规模需多卡（如 8×A100 数天量级），小规模 124M 级别单机可跑，属于学术可复现范围。需要注意的坑：四个循环基线（MoR、heavy-tail、RRT、Ouro）均为作者从零复现，可能存在实现偏差；评测用 lm-eval-harness 零样本协议，明确可对齐；种子敏感性只在小规模报告了 3 个种子（±0.004–0.056），其他规模为单次运行；数据集只写'diverse dataset of text'，未给出具体构成、规模与去重方式，这是严格对齐数字的最大不确定性来源。
论文图表
