Gated Recurrent Transformers: 通过循环调制实现表达深度

TL;DR Gated Recurrent Transformer (GRT) 通过门控循环复用共享核心，实现自适应深度专业化，在 isoFLOPS/isoPARAMS 下以显著更少参数和内存匹配或超越深层 Transformer。

问题

问题背景

在 Transformer 语言模型规模扩展中，表达力 与 内存效率 之间始终存在张力。业界一方面追求更深层模型带来的功能分化能力，另一方面受限于实际部署时的参数存储与推理内存开销。

现有方法局限

传统做法存在两难：

• 每层独立权重：能保持从输入接地到抽象细化的功能专用化，但内存占用线性增长，尤其在大规模模型上难以承受。
• 标准深度共享（如 weight tying）：强制所有层执行相同变换，导致表征多样性坍缩，模型质量下降。
• 现有改进如 MoR（mixture-of-recurrence）或 heavy-tail depth sampling 虽允许部分深度复用，但缺乏自适应的逐 token 调制机制，在 isoFLOPS 约束下仍明显落后于稠密模型。

为什么这个问题难且重要

核心挑战在于如何在参数共享的前提下保留功能分化。递归深度方法必须让同一个共享核心在不同 recurrence step 产生不同行为，同时避免梯度不稳定和训练开销过大。业界对此高度关注，因为大模型推理时的内存和延迟直接决定部署可行性，而参数效率是降低硬件门槛的关键。

行业类比

类似在边缘设备上部署 LLM：我们希望在有限内存与计算预算内，通过动态路由或条件计算获得接近大模型的效果，避免为每个功能模块单独存储一套参数。

核心洞察

• GRT 的核心创新在于用轻量投影和 elementwise update gate 调制共享层的循环更新，并且每步重采样噪声，使同一个层在不同递归步骤中根据 hidden state 和固定 prelude 输出动态改变行为，从而在参数大幅减少的情况下保持功能多样性，避免了标准 depth-sharing 带来的表示坍缩。
• 与 MoR 和 heavy-tail depth sampling 相比，GRT 在 isoFLOPS 和 isoPARAMS 约束下均展现出更优的精度－参数权衡，因为它通过门控机制实现自适应深度重用，而非依赖固定共享或随机丢弃，让模型能够根据 token 难度和递归步数自主决定更新幅度，更接近深层模型的功能特化。

方法

输入与整体流程

输入 token 序列经 embedding 后，首先通过固定深度的 prelude 块（若干标准 Transformer 层）编码，得到锚点表示。随后进入仅含一层或少数层的 共享核心 块，该核心被迭代 R 次，每次迭代复用同一组参数。

关键模块：门控循环更新

每次迭代中，当前隐藏状态、prelude 输出以及每步重新采样的噪声被输入一个轻量投影，产生元素级 更新门（update gate）。该门控以逐元素方式调制循环更新，决定保留多少历史状态、注入多少新信息。这种设计使模型能针对不同 token 或不同迭代步骤动态调整处理强度，实现功能专业化。

训练与输出

训练时采用 循环深度采样（recurrence depth sampling），每次前向随机选择循环次数 R（或从分布中采样），既提升对不同深度的适应性，又降低计算开销。最终，共享核心的输出经固定深度 coda 块处理，得到下一个 token 的 logits，用于语言建模损失。

与同类方法的差异

不同于标准深度共享（如 ALBERT）或均匀重复的 MoR / heavy-tail 深度采样，GRT 通过输入条件门控和每步噪声注入实现自适应深度复用，在参数共享下保持表征多样性。

实验

实验设计

实验在 isoFLOPS 与 isoPARAMS 两个约束下展开。

• isoFLOPS：控制总浮点运算量，对比 GRT 与 12 层 GPT-2 Small、MoR、heavy-tail depth sampling 等基线，覆盖九个预算单元。
• isoPARAMS：固定参数量与数据预算，比较循环深度版本与非循环版本的验证损失。
• 额外分析：KV cache 共享、wall-clock 延迟、解码内存、机制消融与门控行为。

关键发现

• 在 isoFLOPS 下，3 层 GRT 可匹配 12 层 GPT-2 Small 的精度，且训练与推理 FLOPs 相近；在九组预算对比中领先于 MoR 与 heavy-tail depth sampling。
• 中等与大规模预算下，GRT 在标准 token 预算接近稠密模型质量；中等预算 token 翻倍后超越稠密模型。
• isoPARAMS 下，更深循环得到验证损失 2.76，优于非循环基线的 2.84。
• 大规模场景：参数减少 63%，峰值解码内存减少 59%，编译生成延迟仅增加 10%。

对比解读

GRT 与 权重绑定 / 深度共享 方法不同：标准深度共享强制每层相同变换，导致表征多样性坍塌。GRT 通过 elementwise gate 与 resampled noise 对共享核心进行动态调制，使同一个层在不同循环步骤中呈现功能专门化。相比 MoR（混合循环）与 heavy-tail 采样，GRT 在等算力下全面占优，表明自适应门控循环是比随机深度或均匀复用更有效的参数-质量折衷策略。

行业影响

落地场景

Gated Recurrent Transformer(GRT) 通过共享核心层迭代替代深层堆叠,显著降低参数量与解码内存,尤其适合资源受限场景与大规模部署。在电商智能客服中,GRT 可作为对话系统骨干,以更少参数维持多轮理解与生成质量,同时降低单次推理的 GPU 显存需求,支持更高并发。在内容平台个性化搜索与推荐中,可用于用户查询编码与候选内容排序,在相同硬件预算下服务更多用户请求,或在不牺牲效果的前提下减少模型副本数。

商业价值

核心收益来自降本与体验提升两条线。GRT 在 isoFLOPS 条件下匹配更深度模型准确率,在 isoPARAMS 条件下以 2.76 验证损失优于非循环对照的 2.84,说明相同计算量下可提升模型质量。论文显示,大规模使用时可减少 63% 参数 与 59% 峰值解码内存,仅增加约 10% 编译延迟。对云服务商而言,这意味着同等硬件可承载更多模型实例或更大 batch,直接降低单位 token 成本;对边缘设备商,则能用更小存储实现接近大模型效果,提升产品竞争力。

集成接口

GRT 遵循标准 Transformer 接口,易于融入现有 PyTorch / TensorFlow 技术栈。模型架构由固定的 prelude 与 coda 加共享核心循环构成,训练时采用循环深度采样策略,无需特殊优化器或推理框架改动。工程团队可复用现有数据管线与训练集群,仅需将骨干网络替换为 GRT,并在配置中设置循环次数 R。在已有模型微调场景中,GRT 的共享参数结构也便于做模型压缩与蒸馏后的二次优化。其 开源实现 提供了细粒度超参数,可适配不同规模与硬件约束。

局限

• **推理延迟与部署成本** 论文明确指出在 isoPARAMS 条件下，更深的循环会导致编译生成延迟增加约 10%。由于核心模块被循环迭代 R 次，计算图中存在动态循环结构，可能阻碍某些静态编译优化或算子融合，从而在实际部署时需要在吞吐量和内存节省之间做出权衡。对于对延迟敏感的在线服务，这一成本可能抵消部分参数效率收益。
• **实验规模与泛化性** 现有验证主要集中在 GPT-2 Small 级别的模型（约 1.2 亿参数）以及 isoFLOPS / isoPARAMS 的预算设定下，尚未在数十亿参数或更大规模的语言模型上验证。门控机制、状态噪声和循环深度采样在大模型上的稳定性仍未可知，而且不同深度采样策略的训练动态可能发生变化。因此其结论能否平滑扩展到产业级模型规模存在不确定性。
• **超参数敏感性** 方法引入了更新门、轻量投影、每步重采样的噪声等多个组件，并依赖 gate temperature、noise magnitude 等超参数来控制门控行为和噪声强度。论文中的消融实验显示这些超参数对最终性能有显著影响，这增加了训练调参的复杂度和实验成本。与标准 Transformer 相比，需要更细致的超参数搜索才能获得理想结果，降低了方法的开箱即用性。