用同样的权重，为什么第八次和第一次能算出完全不同的东西

2017年，一群做BERT的研究者发现了一件挺尴尬的事。

他们把ALBERT里所有的transformer层都换成同一套权重，参数量确实降下来了，可模型的表现也跟着往下掉。原因很简单：第一层要处理的是刚从词表里查出来的原始向量,充满噪声,需要做的是最基础的语义消歧；到了第十二层,模型手里的东西已经是经过反复提炼的抽象表示,需要做的是逻辑推理和答案生成。这是两种完全不同的工作,用同一把扳手去拧所有型号的螺丝,总有几个拧不紧。

这就是这篇论文要解决的核心矛盾:参数和深度,能不能解耦?

先说清楚这事儿难在哪。

标准的transformer架构里,每加一层就要加一套全新的权重。GPT-2 Large有36层,这36层各自拥有独立的参数,互不共享。这样做的好处是每一层都可以专门学会自己该干的事,坏处是参数量随层数线性暴涨,内存跟着遭殃。想让模型更"深"（更能思考）,就得接受它更"重"（更占显存）。

于是有人想了个取巧的办法:干脆让所有层共享同一套权重,反复循环使用。这就是"循环深度"（recurrent depth）的思路,ALBERT和后来的Universal Transformer都是这条路子。问题是,这样一来,输入到共享层的东西在第一次循环和第八次循环时,状态天差地别,可处理它们的权重却是同一套,没法区分对待。结果就是,要么表现打折扣,要么模型学不到真正有深度的东西。

这篇论文的作者提出了一个思路:如果权重必须共享,那能不能让"输入"变得不一样,从而让同一套权重在不同阶段表现出不同的行为?

这就是Gated Recurrent Transformer（简称GRT，门控循环Transformer）的核心想法。它不是在层与层之间共享权重后就撒手不管,而是引入了一套类似GRU的门控机制,让模型在每次循环时都能"重新认出"自己走到了哪一步,从而做出不同的处理。

三段式结构：固定的入口和出口,可循环的中间层

GRT把整个模型切成三块。

最前面是"前奏"（prelude）,由固定数量的普通transformer层组成,只运行一次,负责把原始输入编码成一个稳定的、带上下文的表示。中间是"共享核心"（shared core）,这部分的权重被反复循环使用R次,是整个架构省参数的关键所在。最后是"尾声"（coda）,同样固定不循环,负责把经过反复打磨的表示,转换成最终能拿去预测下一个词的输出。

论文用一种简写记法描述这种结构,比如"1+5×6+5"就表示1层前奏、5层共享核心循环6次、5层尾声。

这里有个细节值得琢磨:为什么不是所有层都参与循环,偏偏要单独留出前奏和尾声?

论文里提到一个佐证。有研究者（Koishekenov等人）分析过预训练好的大模型内部,发现前面几层主要在做编码（把文字变成语义）,中间层主要在做推理,最后几层主要在做解码（把语义变回具体的词）。这跟人写文章有点像:开头需要先把题目和材料理解透（这是编码）,中间反复琢磨论证逻辑该怎么组织（这是推理,也是最耗时间、最需要反复推敲的部分）,最后落笔成文、遣词造句（这是解码）。如果把"理解题目"和"反复推敲逻辑"用同一套流程混在一起循环,那就相当于每次循环都要重新读一遍题目,效率必然低,而且容易把已经想清楚的东西又想乱。GRT把编码和解码单独摘出来固定住,只让"反复推敲"这部分循环,正好对应了这个分工。

如果不做这个切分会怎样？

论文的对比实验给出了答案:把整个层堆栈全部拿去循环共享,损失会比精心设计的三段式结构高出至少0.1个nats（自然对数损失单位,数值越低表示模型预测得越准）。这不是一个可以忽略的差距,足以让模型在下游任务上明显逊色。

门控机制:让同一套权重,在不同时刻表现不同

三段式结构解决了"哪些层该共享",接下来要解决的是"共享的那部分,怎么才能不僵化"。

这正是GRT最有意思的部分。

在每一次循环里,模型会做三件事。第一步,把当前的隐藏状态（也就是上一轮循环处理完的结果）和"前奏"输出的那份固定表示拼在一起,再加上一点随机噪声,经过一个投影层,变成这一轮真正要喂给共享核心的输入。第二步,共享核心（也就是那几层被反复使用的transformer层）对这个输入做一次处理,产生一个"提案"（proposal）,也就是它认为这一轮应该往隐藏状态里写入的新内容。第三步是关键:一个可学习的门控,逐元素地决定这个提案里,有多少内容真正被写进隐藏状态,有多少内容原封不动地保留。

门控这一步是什么意思?

打个比方。你有一块黑板,每隔一段时间就有一位老师进来,在黑板上做一次批注。这块黑板的空间是有限的,老师不可能每次都把黑板擦干净重写,那样之前的思考痕迹全没了,谁也不知道上一轮想到了什么。更现实的做法是:这位老师每次进来,先看一眼黑板上已经写了什么（这是隐藏状态）,再瞄一眼最初记在小纸条上的原始问题（这是前奏输出,始终固定不变,相当于题目本身）,然后决定这一次要擦掉哪些内容、保留哪些内容、新增哪些内容。早期几轮,黑板还很空,老师会大刀阔斧地写很多东西;写到后面,轮廓已经清楚了,老师就只做小修小补,大部分内容原样保留。

这个门控具体是怎么算出来的？

它是一个小型的两层神经网络,输入是归一化后的当前隐藏状态和归一化后的前奏输出,拼在一起,经过网络算出一个0到1之间的数值,这个数值就是"保留多少旧内容"的比例。数值越接近1,表示这一轮几乎不改动,原样保留;数值越接近0,表示这一轮几乎全部采纳新提案。

门控（gate）：一种逐元素控制信息流的机制,决定新计算结果和旧状态之间该以什么比例混合,常见于GRU、LSTM等循环神经网络。

值得一提的是,这个门在训练刚开始时被故意设置成"几乎全保留"的状态,数值接近0.98。也就是说,一开始模型几乎不改动残差流里的内容,循环形同虚设。随着训练推进,模型才慢慢学会在哪些位置、哪些维度上应该主动改写。这种设计参考了LSTM研究中的一个经验:如果一开始就让"遗忘门"倾向于遗忘,模型很容易在训练早期就迷失方向,不如先让信息尽量流动起来,再逐步学会取舍。

如果没有这个门控机制,只是简单地把上一轮的输出原样喂给下一轮呢？

论文的消融实验给出了明确答案。仅仅加入循环这一项操作,不做任何区分处理,验证损失比不循环的密集基线模型（也就是标准的、每层权重都独立的GPT-2）还要高出0.107个nats。换句话说,如果只是机械地反复套用同一套权重,模型的表现反而比老老实实用12层不同权重的传统做法更差。只有把前奏尾声结构、状态噪声、前奏重新注入、以及这个逐元素门控全部叠加起来,验证损失才降到3.141,低于密集基线的3.157。这五项改进各自的贡献分别是负0.035、负0.018、负0.022、负0.048个nats,门控这一项贡献最大。

除了门控之外,论文还往投影这一步和门控计算里各加了一份高斯噪声。往投影输入里加噪声,是为了防止模型学出一种"死记硬背"式的捷径,也就是不管输入是什么,每一轮都套用完全相同的固定模式;往门控计算里加噪声,是为了防止门控数值早早收敛成一个不再变化的常数,失去区分不同时刻的能力。

跟别人比,谁更聪明

论文没有闭门造车,而是拿GRT去跟好几种同类方法正面对撞,而且是在同一套代码、同一批数据、同样的训练步数下从头训练的,保证比较公平。

对比的对象里,有一种叫MoR（Mixture-of-Recursions,混合循环深度）,它给不同的词分配不同的循环步数,需要一个专门的路由器来做决策;还有一种延续Geiping等人思路的"重尾泊松深度采样"方法,循环次数是随机采样的,但没有门控这类精细的调节手段;还有RRT（Relaxed Recursive Transformer,松弛循环transformer）,它给每一次循环加一个独立的LoRA适配器（一种轻量级的参数微调模块）,相当于给共享权重打了几个小补丁,让每次循环稍微有点差异;还有Ouro,它对每一次循环都单独计算损失并加以监督。

在"计算量相同,参数量减少"（论文称为isoFLOPS）的对比条件下,GRT用GPT-2 Small规模的配置（3层结构,循环10次,只存了35M参数）,达到3.14的验证损失,和12层、124M参数的密集基线打平,同时参数量少了64%。到了Medium规模,GRT的127M参数拿到2.89,比密集基线的2.84略高一点,但比MoR的3.02、重尾泊松的2.97、Ouro的2.93、RRT的2.95都要低,是四个循环深度方法里表现最好的。Large规模也是类似的格局,GRT拿到2.77,依然领先其他三个循环方法,只是还没追上774M参数的密集基线的2.71。

反过来,在"参数量相同,允许计算量增加"（isoPARAMS）的对比条件下,GRT表现得更抢眼。同样的参数预算下,让模型多循环几轮,GRT在Medium规模拿到2.76,比不循环的对照组2.84还要低,在Large规模是2.65对2.71。这说明如果你愿意为同样大小的模型多花一些推理时的计算量,循环深度确实能换来实实在在的质量提升。

论文还在九个真实的下游任务（比如常识问答、阅读理解、句子补全）上做了零样本测试。isoFLOPS版本的GRT用只有37%的参数量,平均分42.08,几乎打平密集基线的42.05;isoPARAMS版本的GRT在九个任务里的八个上都超过了密集基线,平均分高出2.10分,其中LAMBADA标准版任务的提升达到8.29分。

早退出:一个意外获得的免费功能

训练过程中还藏着一个巧妙的设计。每次训练,模型循环的步数不是固定的,而是从1到最大值R之间随机抽取一个数,这就是所谓的"深度采样"（depth sampling）。

深度采样：训练时不固定每次前向传播要循环多少步,而是随机选取一个步数,使得模型的每一个中间循环状态都被训练成"能直接输出合理预测"的状态。

这样做带来一个意外的好处:训练完成之后,模型天然具备"提前退出"的能力。你想要快一点、省一点计算量,就让它只循环两三轮就直接出结果;你想要更准确的答案,就让它循环满六轮。整个过程不需要额外训练任何专门的退出判断模块,也不需要给每个退出点单独设计损失函数。

论文里给出一个具体的对比:同样的推理计算量下,GRT提前退出后的损失,比标准密集模型在对应层数提前退出后的损失要低得多。图里的曲线显示,GRT从第0步（只经过前奏）到第6步,损失从6以上迅速降到2点多;而密集模型如果在对应的浅层就掐断,损失还停留在7到10之间的高位。这意味着如果你只有一半的计算预算,GRT依然能保留大约92%的准确率,而传统密集模型在同样浅的深度上远远达不到这个水平。

如果不做深度采样,只用固定的循环次数训练呢？

论文对此也做了检验,发现深度采样不仅带来了免费的早退出能力,还起到了一种"随机深度正则化"的作用,让最终的验证损失比固定深度训练更低。这相当于一种意外之喜:本来是为了解决推理灵活性问题设计的采样策略,顺带还改善了训练效果本身。

打开黑箱:模型循环时到底在想什么

论文花了相当大的篇幅去拆解GRT内部到底发生了什么,这部分内容读起来像是一场侦探式的调查。

第一个发现是,损失下降得非常集中。只经过前奏,模型的损失是5.29;循环一次之后骤降到3.77;两次之后降到3.14。前两步就贡献了从前奏到最终输出总损失下降量的77%。剩下四步只是在做精修,把损失从3.14磨到最终的2.68。这跟前面提到的门控行为完全吻合:早期步骤门比较开,大幅改写;后期步骤门逐渐关闭,只做微调。

第二个发现更有意思。研究者做了两组极端对照:强行把门控设为0（意味着每次循环都完全采纳新提案,不保留任何旧状态）,损失飙升到12.03,比正常模型差得离谱;强行把门控设为1（意味着每次循环都完全不改动,原样复制）,损失是5.26,正好等于只经过前奏、完全没有循环的效果。这两个极端结果像两道护栏,夹出了训练好的门控（2.68）到底做对了什么:它既不是简单粗暴地全部重写,也不是形同虚设地全盘照搬,而是学会了在具体哪些位置该写、哪些位置该留的精细平衡。

第三个发现关于表征几何。研究者用一种叫CKA（Centered Kernel Alignment,中心核对齐,一种衡量两组神经网络内部表征相似程度的数学工具）的方法,比较了循环的第1步到第6步之间的表征相似度,同时也拿这些表征去跟一个36层的标准GPT-2 Large做对比。结果发现,前奏和第1步的表征几乎完全一致,但从第2步开始,表征发生了一次明显的"相变",第2到第6步彼此之间高度相似,形成一个紧密的簇。跟标准GPT-2 Large对比时,前奏和第1步的表征最接近GPT-2 Large的浅层（第2层左右）,而第2步到第6步则全部稳定地对齐到GPT-2 Large的第11层附近,再也没有推进到更深的层。这说明GRT的循环核心,本质上是在反复"复用"标准模型中间层那部分的功能,而不是模拟出一整套深浅不一的层级体系。

第四个发现关于难易分工。研究者按照BPE分词后的token出现频率,把词分成标点符号类、常见词、中等频率词、生僻词四类,追踪它们各自在六轮循环中损失下降的情况。结果显示,标点符号这类简单的token,第一轮循环就能解决掉49%的初始损失,累计六轮下来能恢复67%;而生僻词第一轮只能解决23%到28%,累计六轮也只恢复42%。这说明同一次前向传播里,模型其实同时在干两件不同的事:对简单的、大概率已经猜对的词做"补全式"的加速确认,对生僻的、不确定的词做"逐步收窄"式的持续修正。而这一切分工,是靠逐元素的门控自动完成的,没有任何显式的路由机制去指挥"这个词简单,那个词难"。

如果你觉得这背后没什么道理,可以换个角度想。一个熟练的校对员看一篇稿子,遇到常见的、明显没问题的词,扫一眼就过去了;遇到拿不准的专业术语或者生僻字,会反复琢磨,查资料,来回修改好几遍。校对员没有专门的"难度分类系统"去提前给每个词打标签,这种差异化投入完全是根据实际读到的内容自然产生的。GRT的门控机制起到了类似的作用,只不过它是在向量空间里,用一个数值大小来决定该在哪个词上多花几轮计算。

内存和速度的账,怎么算

省参数听起来很美好,但落到实际部署上,还有一笔账要算清楚:解码时的KV缓存（key-value cache,transformer在生成文本时用来存储历史token的键值向量,避免重复计算的机制）。

因为GRT要循环R次,如果每次循环都各自保留一份KV缓存,内存开销会直接乘以R倍,把省下来的参数优势吃掉一大半。

论文测试了三种压缩策略:只保留最后一步的缓存、只保留第一步的缓存、把所有步骤的缓存做平均。结果显示,平均策略效果最好,不仅内存开销显著降低,在HellaSwag（一个常识推理评测任务）上的准确率反而比不压缩的完整缓存还高一点,33.90对33.65。这大概率是因为平均本身起到了一种轻微的正则化效果,削弱了个别步骤里可能存在的噪声。

具体到数字,在Medium规模模型上,批大小为1时,GRT用完整的R倍缓存,总内存（包括权重和缓存）已经是密集模型的0.55倍;批大小为32时,缓存开销开始主导,完整缓存策略的优势收窄到0.91倍;但用平均策略,批大小32时依然能压到0.39倍。

延迟方面,在编译优化（torch.compile）开启的情况下,GRT Large相比GPT-2 Large,生成延迟只多了10%,但参数量少了62%,峰值解码内存少了59%。这笔账整体上是划算的:用一点点速度上的代价,换来大幅的内存节省。

三个反过来印证设计合理性的细节

论文里还有几个细节值得单独拎出来说。

第一,论文测试了如果把"锚点"（也就是喂给每次循环的那份固定前奏输出）换成别的东西会怎样。换成全零向量,损失暴涨到8.08;换成原始的词嵌入,损失是3.73;换成"上一步的隐藏状态"（也就是让锚点跟着循环漂移,而不是固定不变）,损失是3.38。跟训练好的模型（用固定前奏作锚点)的2.68相比,固定锚点这一项设计单独贡献了0.70个nats的改善。这说明给每一轮循环提供一个稳定不变的参照系,价值巨大,哪怕门控本身依然正常工作。

第二,论文分析了投影层的权重结构,发现处理"当前状态"的那一半权重和处理"前奏锚点"的那一半权重,它们的行向量平均余弦相似度是负的0.189。也就是说,这个投影层不是在简单地把两路信息做平均,而是在计算两者之间的差异,相当于问"跟最初相比,现在改变了什么"。这也解释了为什么固定前奏这么重要:它提供了计算"变化量"所需要的参照基准。

第三,论文追踪了门控里"当前状态"和"前奏锚点"两路信号各自贡献了多少方差。第一步时两者接近对半开(59%对41%),到第六步时变成74%对26%,当前状态的话语权越来越大。追根究底,这跟处理当前状态那部分权重的最大奇异值(9.87)比处理锚点那部分权重的最大奇异值(6.85)高出44%有关,这是写在权重里的结构性偏好,而不是训练过程中偶然形成的。

写在后面

读完这篇论文,我印象最深的其实不是那些漂亮的benchmark数字,而是那个"Gate=0对Gate=1"的极端对照实验。12.03对5.26,两个数字像两道墙,把训练出来的2.68死死夹在中间。这种做法很老实,不像很多论文只给你看最好的结果,它把两个失败的极端也摆出来了,让你自己判断这个门控到底值多少钱。

另一个让我反复琢磨的细节是CKA分析里那个"相变"现象。前奏和循环第一步几乎完全重合,但第二步开始突然跳到一个新的簇里,而且这个簇一直稳定到第六步,始终对齐在标准GPT-2 Large的第11层附近。这意味着这个共享核心,其实并没有学会去扮演"浅层做编码、深层做解码"这种完整的层级分工,它更像是反复在扮演一个固定的中间角色,只是扮演的次数和火候不一样。这跟我们直觉里"循环等于模拟深度"的想象不完全一致,循环出来的不是一个纵深的层级体系,而是同一个中间态被反复精修的过程。

论文里还有一个没有深入展开但值得追问的地方:循环深度R目前是训练前就定死的超参数,推理时没法根据每个token的难度动态调整循环次数。论文自己也承认这是局限,提到未来可以引入类似ACT（Adaptive Computation Time,自适应计算时间）那种逐token的停止机制。如果真能做到,不同难度的词用不同的循环次数,而不是整个序列统一循环R次,这大概会是下一步更有意思的方向。

一个门控,决定了权重被重复使用了多少次之后,是继续做同一件事,还是已经悄悄换了一份工作。这听起来有点像我们每天早上醒来,用的还是同一个大脑,但清醒的第一分钟和忙碌了一整天之后的第八个小时,这个大脑处理信息的方式,真的还是一回事吗？

Q&A

Q1：Gated Recurrent Transformer是什么？

A：一种通过给共享权重加入门控机制来实现"循环深度"的语言模型架构,能让同一套transformer层在反复循环使用时,针对不同阶段做出不同的处理,从而在大幅减少参数量的同时保持接近传统深层模型的表现。

Q2：GRT和普通的权重共享模型（比如ALBERT）有什么区别？

A：普通权重共享模型对每次循环都用完全相同的处理方式,容易导致表现下降;GRT额外加入了逐元素门控、固定前奏锚点和噪声注入,让共享的那部分权重在不同循环步骤里表现出不同的行为,弥补了单纯权重共享带来的表达能力损失。

Q3：GRT训练完之后能不能提前退出、节省计算量？

A：可以。因为训练时循环步数是随机采样的,模型的每一个中间步骤都被训练成能直接输出合理预测,所以推理时只跑一半的循环步数,依然能保留大约92%的准确率,不需要额外训练专门的退出模块。