先说一件挺反直觉的事。

过去几年,只要有人做"循环Transformer"(就是把一部分神经网络层反复执行好几遍,而不是简单地堆更多层),几乎都会拿它和一个参数量相同的普通模型做对比,然后欢呼:"看,循环模型用更少的参数达到了差不多甚至更好的效果!"这个结论听起来很美好,但它藏着一个陷阱:循环意味着同样的层要多算几遍,这些多算的部分是要花算力和时间的。用现实类比一下,这就像两个人比赛谁做饭做得好吃,一个人只用了三种食材,另一个人也用了三种食材,但后一个人在锅里多炒了两遍。你说食材数量一样公平,可是多炒的那两遍时间、燃气费可没算进去。

这篇来自清华大学、字节跳动 Seed、M-A-P 和 TokenWave.AI 团队的论文,就是冲着这个不公平的比较去的。他们提出了一个问题:如果把训练花的算力、模型存下来的参数总量、还有推理时占用的显存(KV缓存)这三样东西都严格对齐,循环Transformer还能不能赢?


【资料图】

答案是能赢,而且赢得还不少。他们把这套经过反复打磨的方案取名叫 SMELT(Sparse MoE Transformer, middle layers Loop Twice,意思是"稀疏混合专家Transformer,中间层循环两次")。在最大的实验规模上,SMELT能用少花6.8%到18.0%的训练算力,达到和普通模型一样的效果。

这篇文章就带你一步步搞清楚,这个"少花算力多办事"的效果是怎么来的。

三个必须同时锁死的预算

要理解这篇论文的价值,得先明白它解决的是一个"公平秤"的问题。

之前的研究里,比较循环模型和普通模型的方式五花八门。有的固定住模型存储的参数量,然后不断增加循环执行的次数,这样效果自然会涨,因为你偷偷多花了计算量而不自知;有的干脆拿一个小循环模型去对比一个大好几倍的普通模型,强调"参数效率",可这同样没说清楚,效果好到底是因为循环这个设计本身厉害,还是仅仅因为多花了算力。

MoE(Mixture-of-Experts,混合专家):一种神经网络设计,模型里放着一大堆"专家"(其实就是一组一组的参数),但对每个输入的词,只挑其中一小部分专家来实际计算,这样存储的参数总量可以很大,但真正参与计算的只是一部分,从而把"存了多少"和"算了多少"这两件事解耦开。

这篇论文的作者认为,真正干净的比较,必须同时锁死三个东西。

第一是每个词处理时花费的浮点运算次数,这决定了训练和推理的真实成本;第二是模型的总参数量,这大致决定了模型能装下多少知识;第三是KV缓存的大小,这决定了模型能同时处理多长的文本。

FLOPs(浮点运算次数):衡量计算量的单位,可以理解成"电费账单",越大说明这次计算越费资源。

KV缓存:Transformer在生成文字时,需要把之前算过的一些中间结果(键和值)存起来重复利用,这部分占用的显存空间就叫KV缓存,它的大小直接限制了模型一次能处理多长的对话或文档。

之所以这三个都要卡住,是因为它们分别对应了实际部署时的三种代价:算力账单、模型下载和存储的体积、以及服务器上能同时开多长上下文的能力。只要有一个没对齐,那接下来所有关于"谁更强"的结论,都可能只是"谁花的钱更多"这句话的另一种说法。

那这三个预算怎么同时对齐?这就要说到MoE架构给这场比较提供的一个巧妙杠杆了。

用MoE给循环"买单":窄一点,专家多一点

普通的循环模型,如果把一部分层多算一遍,效果一定是每个词处理的浮点运算数量涨了。要把这部分涨出来的算力"还回去",最直接的办法是把模型其他地方缩窄一些,比如把每层神经元的隐藏维度调小。

问题来了:把隐藏维度调小,会连带把模型的总参数量也压下去,这样一来,参数量这个预算又对不齐了。

这时候MoE的好处就体现出来了。因为MoE把"存了多少参数"和"算了多少浮点运算"这两件事本来就分开了,你完全可以在缩窄隐藏维度、降低单次计算量的同时,多塞一些"专家"进去,把参数总量重新补回来。因为每次真正干活的专家数量是固定的(论文里固定为每次挑8个专家),专家池子的整体大小可以随便扩容,而不会增加单次计算的开销。

这就好比一家餐厅想推出一道"复刻两遍工序"的招牌菜(对应循环两次),厨师多花的这道额外工序的时间成本,要靠别的地方省出来。餐厅的做法是把主菜的分量切薄一点(隐藏维度变窄),省下的时间用来多炒一遍,但同时又在后厨冰箱里囤了更多种类的配菜(专家数量增多),保证顾客点单时依然有丰富的选择,餐厅冰箱的"总库存"(总参数量)没有变少。如果不这么做,要么这道菜做出来慢半拍、成本超支,要么配菜种类被迫减少,客人能选的花样变少了。KV缓存这块则单独通过调整注意力头的大小和分组比例(GQA比例)来补齐,确保额外执行的那几层不会让显存占用涨太多。

GQA(分组查询注意力):一种让多个注意力头共享同一组键值参数的技术,可以在几乎不损失效果的情况下大幅减少推理时需要缓存的数据量。

按照论文里给出的一个具体例子:一个2亿参数规模、循环中间6层执行两遍的模型,把隐藏维度从1280降到1056,同时把每层专家数从192个提升到288个。这样调整之后,每个词的计算量只比不循环的基准模型多了2.9%,总参数量只多了0.4%,KV缓存差距也控制在4%以内。三个预算基本被摁平了,接下来的比较才算干净。

搞清楚了怎么"付账",下一个问题自然是:这笔账花在哪最划算?也就是说,究竟该循环模型的哪一部分、循环几遍才最好?

三次实验,锁定一份配方

论文的作者没有直接拍脑袋决定循环方案,而是做了三组系统的消融实验(也就是控制变量、逐一验证的对比实验),一步步试出最优解。

第一个问题是:循环整个模型好,还是只循环中间那一段好?

他们试着让循环的层数范围从0(等于完全不循环,就是基准模型)一路涨到12层(相当于整个模型都循环),发现无论是在中等稀疏度还是较高稀疏度的设置下,验证集上的损失(也就是模型犯错的程度,数值越低越好)都在"循环中间50%的层数"这个点上达到最低。循环全部层数或者只循环很少几层,效果反而都不如只循环中间那一半。

这个结果其实挺符合直觉的,也呼应了一些更早的研究:Transformer的第一层和最后几层往往承担着比较特殊、专门化的功能(比如最开始几层负责基础的词法和语法信息处理,最后几层负责把信息整理成最终输出),这些"专才"层如果被拿去反复循环共享参数,反而会拖累它们各自发挥专长的能力。中间层则更像是通用的"信息加工车间",多加工几遍反而有益。

打个比方,一份文件从收发室到最后盖章出库,中间要经过很多道审核加工的环节。收发室拆信封、盖章出库这两道工序是专门的岗位,谁都替代不了,硬要把这两道工序也拿来反复循环执行,只会让流程变得又慢又乱。但中间那些审阅、修订、核对的环节,多让同一批人多看几遍、反复打磨,往往能把内容打磨得更扎实。如果把循环范围放大到全部工序,收发室和盖章员工也被迫反复折腾,专业岗位的效率反而被拖累。

第二个问题是:循环模型和普通模型相比,更适合"瘦高"还是"矮胖"的身材?

论文把模型的物理层数(也就是真实存储、不重复的层数)在9到18层之间扫了一遍,同时对普通模型和循环模型分别找各自的最优点。结果发现,普通模型在物理深度是12层时效果最好,而循环模型呢,它的最优配置反而是物理深度也是12层,但因为循环执行两遍,实际跑起来相当于18层的深度。也就是说,在同样"存起来的层数"下,循环模型更愿意让自己实际执行的深度变得更深一些,这个深度和宽度的最优比例,比普通模型要更偏"深"一点。

作者给出的一个可能解释是:循环模型的额外执行深度,并没有真的增加物理层数,也就没有增加需要独立训练的参数。共享的那部分层在训练时会从多个执行位置(第一次循环和第二次循环)都拿到梯度信号,尤其是第二次执行时,这一层到最终输出的路径更短,梯度传得更直接。这种"多重深度反馈"可能让额外的串行计算比堆叠一批全新参数的层更容易被优化好。

第三个问题最直接:循环两遍够不够,要不要循环三遍、四遍?

论文的答案是:两遍就是最佳点,三遍和四遍效果都会往下掉。原因也不复杂,因为受到算力预算的限制,循环次数越多,为了不超支就必须把模型的宽度削得更窄,模型练三遍或四遍时已经瘦得"营养不良"了,反而不如循环两遍时那个更均衡的胖瘦比例好用。

综合这三条经验,作者把最终方案定型,取名SMELT:循环中间一半的层,两遍,并且让循环部分执行起来比普通模型显得更"深"一点。这份配方随后被拿去做大规模的验证。

从消融到规模:把配方铺到54B参数

光在小规模的实验里验证一个想法靠谱是不够的,很多方法在小模型上好使,放大之后就失灵了。论文接下来把SMELT这份配方铺到了从1亿到16亿(活跃参数)共四个规模档位,每个档位又对应四种稀疏程度(也就是模型里专家总数不同带来的参数和计算比例差异),总共训练了32个正式模型,加上不同的训练步数分支,一共评估了192个数据点。最大的一档模型总参数量能达到540亿(这是非嵌入部分的参数,也就是不算词表嵌入层的那部分)。

在这么大的跨度里,SMELT在每一个规模、每一个稀疏度上,训练损失都比对应的普通模型要低。接下来他们做的一件更有意思的事情,是给SMELT和普通模型各自单独拟合一条"缩放定律"曲线。

缩放定律:描述模型效果(比如损失值)随着参数规模、训练数据量、计算量增大而变化的一种规律,通常呈现为幂律关系,可以用来预测更大规模模型的表现。

这份工作参照了业界知名的Chinchilla缩放定律,把损失拆解成一个不可消除的固定误差项、一个和模型容量相关的项、还有一个和训练数据量相关的项。作者把这个公式里代表容量的部分,换成了每个词消耗的计算量以及一个衡量稀疏程度的指标,这样一条公式就能同时描述稠密模型和各种稀疏程度的MoE模型。

结果显示,SMELT在两个关键指数上都比普通模型高:一个是描述模型容量增长时损失下降速度的指数,另一个是描述数据量增长时损失下降速度的指数。这两个指数综合起来,决定了SMELT的损失曲线随算力增长下降得更快。具体折算成能省下的算力,在训练预算为10的20次方浮点运算时,SMELT能省下6.8%到10.0%的算力;当预算涨到10的21次方时,能省下的比例进一步扩大到14.7%到18.0%。这说明这个优势不是一次性的小恩小惠,而是会随着规模增长持续放大的系统性优势。

指标数据|Baseline|SMELT

容量增长指数 a|0.3703|0.3892

数据增长指数 c|0.6594|0.7011

综合频率指数 γ|0.237|0.250

1e20 FLOPs处省算力|—|6.8%~10.0%

1e21 FLOPs处省算力|—|14.7%~18.0%

这些数字听起来有点抽象,换个说法可能更直观。假设你有两个厨师,每人做同样一道菜,随着厨房设备(相当于算力预算)越来越先进,一个厨师做菜的速度提升幅度是每年10%,另一个是每年12%。刚开始两人差距不大,可十年之后,累积下来的差距会变得非常明显。SMELT就是那个提升幅度更快的厨师,规模越大,这个差距就体现得越明显。

下游任务的意外惊喜:验证损失没算到的部分

单纯看训练时的损失值,只是这个故事的一半。因为损失是在验证集这样一个统计意义上算平均值,它并不能完全代表模型在具体任务上做得好不好。

论文接下来做了一件很扎实的事:他们先用96组普通模型的数据,拟合出一条"验证损失对应下游任务分数"的标准曲线,相当于给出了一个预期值,"如果你的验证损失是这么多,正常情况下你的下游分数应该是这么多"。然后再把SMELT的实际下游分数和这条曲线预测出来的预期值做对比。

结果发现,SMELT在几乎所有规模和几乎所有基准测试上,实际表现都超过了这条曲线预测的水平。也就是说,SMELT不仅仅是因为验证损失更低才带动了下游任务分数提高,它带来的下游收益比单纯损失下降能解释的还要多。这个多出来的部分,还会随着模型规模变大而进一步扩大。

举个具体的例子,在DCLM Completion这个指标(衡量模型生成正确答案的能力,数值越低越好)上,1亿参数规模时超出预期的幅度只有1.7个千分之一,等到规模涨到16亿时,这个超出幅度涨到了12.3个千分之一,涨了整整7倍多。

这就好比两个学生复习同一门课,模拟考分数一样,但正式考试时,那个平时多做了"举一反三"训练的学生,成绩往往会超出模拟考成绩推算出的预期水平,而且随着考试难度越来越大,这种超预期的幅度反而更明显,说明他掌握的不只是死记硬背的知识点,还有一些能迁移到新场景的能力。

再往深挖一层,这份超预期的收益并不是平均分布在所有类型的数据上的。论文按照Code(代码)、Finance(金融)、Math/STEM(数学与理科)、Knowledge(知识)、Web(网络文本)这五类数据分别计算了"计算效率增益"(也就是SMELT相比普通模型能省下的算力比例),发现Code这一类的增益最高,达到20.4%,Finance和Math/STEM紧随其后,都在16%到17%之间,Knowledge和Web则相对较低,只有约15%。

作者的解释是,这个排序恰好对应了数据本身内在结构的强弱:代码有严格的语法规则和长距离的依赖关系(比如一个变量在文件开头定义,几十行之后才被用到),而网络文本相对来说随意松散得多。这暗示着,循环带来的额外一轮"精加工",对那些需要严密逻辑、远距离信息关联的任务帮助最大。

除了数据类型,样本长度和上下文示例的数量也是两个值得关注的维度。论文发现,SMELT相对普通模型的损失改善幅度,在长文本(512到4096个词)上是短文本(32到256个词)上的1.52倍。而作为对照,单纯增加参数量或者单纯增加专家数量这两种普通的"变强"手段,并没有表现出这种偏向长文本的倾向,改善幅度在长短文本上基本持平。

在上下文学习(也就是给模型看几个示例,让它照葫芦画瓢完成新任务)这个维度上,SMELT和普通模型的差距在没有任何示例时只有0.9个百分点,一旦给出至少一个示例,差距立刻扩大到1.9个百分点,并且一直保持到给出8个示例的情况。有一个特别能说明问题的任务叫Dyck Languages(括号匹配任务,模型需要判断嵌套括号是否正确闭合),这个任务的特点是不给示例的话,模型几乎完全无法作答(正确率接近0%),唯一的办法就是去读懂提示里给出的示例答案。在给到32个示例的情况下,SMELT达到了29.8%的正确率,而普通模型只有26.4%。

这些现象叠加在一起,勾勒出一个相当清晰的图景:SMELT这种"精加工一遍"的设计,特别擅长处理那些需要模型回过头去检索、比对、复用远处信息的场景,无论这个远处信息是文档里的前文,还是提示词里给出的示范例子。

那这背后的机制到底是什么?模型内部第二次循环执行时,到底在干什么?

拆开看第二次循环:它在补强,而不是推翻

为了搞清楚第二次循环内部到底发生了什么,作者做了一系列非常细致的机制分析,拿真实的验证数据跑一遍模型,观察内部的路由选择、残差流更新和注意力模式变化。

第一个观察角度是专家路由。MoE模型每次处理一个词时,会挑选一小撮专家来实际计算(论文里是挑8个)。同一个词第一次经过某一层和第二次经过同一层时,挑出来的专家集合有多大程度重合?

结果显示,在稀疏程度较低(专家池子小)时,第二次挑选的专家几乎和第一次一模一样;随着稀疏程度提高、专家池子扩大十几倍之后,重合的专家数量降到2到3个左右,但即便如此,这个重合度依然远高于纯粹随机选择应该达到的水平。这说明路由机制不是完全随机地探索新专家,而是有意识地保留一部分核心专家,同时把剩下的名额分配给一些新的专家做补充判断。

第二个观察角度是残差流的更新幅度。Transformer每一层都会往一条"残差流"(可以理解成信息在层与层之间流动的主干道)里写入一点新的内容。第二次执行同一层时,写入的这部分内容,无论是绝对大小还是相对大小,普遍比第一次执行时更大。而且论文进一步计算了第一次写入和第二次写入这两个方向之间的夹角(用余弦相似度衡量),发现同一物理层在两次执行之间的写入方向,相似度明显高于随便两个不同层之间的相似度。

这意味着什么?意味着第二次循环并不是把第一次写入的内容推倒重来、另起炉灶,而是顺着第一次已经建立起来的方向继续加码、把信号放大。

这就像修改一份写好的文章草稿。第一遍写完之后,作者拿着草稿再读一遍,这一遍修改通常不是把整篇文章推翻重写,而是顺着原来的思路,把论据补充得更充分、语气强调得更明确,写下去的内容和第一遍的方向基本一致,只是力度更大了。如果第二遍修改和第一遍完全南辕北辙、方向随机漂移,那这份稿子只会越改越乱,根本谈不上"精加工"。

第三个观察角度落在了注意力机制内部。注意力机制里有三组关键的投影:Q(查询)、K(键)、V(值)。简单说,Q和K决定了"该往哪里看",V决定了"看到的内容具体是什么"。

论文发现,Q和K在两次循环之间的相似度非常高(甚至比它们共同的输入本身的相似度还要高),说明第一次循环已经把"该往哪个位置看"这件事确定下来了,第二次循环基本延续了这个判断。而V的相似度明显更低,说明第二次循环虽然看的位置差不多,但读出来的具体内容有了实质性的更新。用一个具体统计来说,两次循环选出的"最受关注的8个位置",重合比例达到56%到66%,而如果换成同一层里的不同注意力头做对照,重合比例只有28%到34%。

这就有点像一个人回过头重新审视一份材料,他记得上次已经找到了几个关键段落,这次翻阅时不需要重新地毯式搜索,直接翻到那几个熟悉的段落,但这一次读的时候会读得更仔细,理解出来的信息含量比上次更丰富。如果每次重读都要重新翻遍整份材料去确定该看哪里,那多读一遍的效率就会大打折扣,读到的内容也未必更深入。

第四个也是最直观的一个观察,是关于"注意力沉没"(attention sink)现象。

注意力沉没:指Transformer模型在处理文本时,无论内容是什么,总会习惯性地把大量注意力分配到序列最开头的那个标记(token)上,这个现象在很多标准Transformer模型里都能观察到,被认为和模型内部数值稳定性有关,但也意味着这部分注意力被"浪费"在了一个和内容本身无关的位置上。

论文发现,第二次循环执行时,这种沉没在开头标记上的注意力比例明显下降,被释放出来的注意力转而流向了那些真正和内容相关的位置。他们用括号匹配任务(Dyck Languages)做了一个很直观的案例展示:第一次循环时,模型的注意力大量堆积在句子开头的标记上,好像还没搞清楚该往哪看;第二次循环时,这部分注意力几乎清空,转移到了提示词里那些具体的示范答案上。平均下来,开头标记的注意力占比从0.60骤降到0.02,而示范答案的注意力占比则从0.24飙升到0.85。

更有意思的是,这个"注意力沉没随深度增加"的现象在普通模型里是逐渐加重的,层数越靠后,堆积在开头的注意力反而越多。但SMELT恰好反过来,在物理深度相同的位置上比较时,第二次循环的注意力沉没水平反而比第一次循环更低,跟普通模型的这个趋势正好拧着来。

这几条线索拼在一起,勾勒出一个相当连贯的故事:第二次循环并没有把第一次的判断推翻重来,它保留了"该往哪里看""该找哪些专家帮忙"这些基本判断,但在这个基础上,把浪费掉的注意力资源重新分配,把内容读得更透,把结论写得更笃定。这可能正是SMELT在长文本、结构化数据、上下文学习任务上表现出额外优势的内在原因。

写在后面

读完这篇论文,最让我意外的不是SMELT省了多少算力,而是那个关于注意力沉没的实验。之前接触到的注意力沉没相关研究,大多把它当成一个训练完成之后的静态属性来观察,好像这是模型定型之后就固定不变的一个特征。但这篇论文发现,在同一次前向计算内部,仅仅是重新执行了一遍同样的参数,这个所谓的"静态属性"就发生了系统性的变化。这说明注意力沉没这件事,可能没有想象中那么"顽固",它更像是模型在信息不充分时的一种默认保守策略,一旦给它机会重新审视一遍,它是有能力自己纠正的。

还有一个细节值得单独说一说:论文里提到,循环模型偏好更"瘦高"的形状,这个观察其实和"深度决定串行计算步数"的理论是呼应的。有些理论工作证明过,Transformer要完成某些天然需要多步串行推理的任务,仅靠增加宽度是不够的,必须要有足够的执行深度。这篇论文用实验数据从另一个角度印证了这一点:在参数预算固定的情况下,模型更愿意把预算换成"多算几步",而不是"每一步算得更宽"。这可能提示我们,衡量一个模型"够不够聪明",光看参数量这个单一数字是不够全面的,执行时真正走过的计算路径长度,可能同样重要,甚至更重要。

这篇论文也留下了没解决的问题。作者自己也承认,所有的架构搜索实验都是在较小规模上做的,更大规模下最优的循环层数范围会不会不一样,目前还不知道;论文研究的也只是最简单的"完全共享参数"这种循环形式,像带低秩适配器的松弛版本、按token动态决定循环深度这些更灵活的设计,跟这套预算对齐的方法结合起来会发生什么,也还是空白。

如果这套"用MoE给循环买单"的思路继续往前走,会不会有一天,我们训练大模型的方式不再是简单粗暴地堆参数堆层数,而是先想清楚"这个任务到底需要几步串行思考",再反过来设计模型该多深、该多宽?

Q&A

Q1:SMELT是什么?

A:SMELT是一种循环Transformer架构方案,全称是"Sparse MoE Transformer, middle layers Loop Twice",核心做法是把混合专家Transformer模型中间一半的层循环执行两遍,同时通过收窄隐藏维度、增加专家数量、调整注意力头设置,让每词计算量、总参数量、KV缓存这三项预算都和不循环的普通模型基本对齐。

Q2:SMELT相比普通模型能省多少训练算力?

A:在训练预算为10的20次方浮点运算时能省6.8%到10.0%的算力,预算涨到10的21次方时能省14.7%到18.0%,而且这个优势会随着模型规模和算力预算增长而持续扩大,最大验证规模达到540亿非嵌入参数。

Q3:为什么循环两遍比循环三遍四遍效果更好?

A:因为在总计算量预算固定的前提下,循环次数越多就必须把模型宽度削得更窄来腾出计算量,循环三遍四遍时模型会变得过于纤细,反而不如循环两遍时那个更均衡的深宽比例效果好,论文的消融实验证明两次循环是最佳点。

标签: 算力 实验 新模型 神经网络

免责声明:市场有风险,选择需谨慎!此文仅供参考,不作买卖依据。