七年前,《Attention is all you need》这篇论文横空出世,transformer架构从此碘伏了整个深度学习领域。如今,各家大模型几乎都以transformer为基础,但它内部到底是怎么运作的——依然像一个黑盒子,没人能完全说清楚。去年,transformer论文作者之一Llion Jones创立了Sakana AI,而最近这家公司发表了一篇很有意思的研究,标题是《Transformer Layers as Painters》。他们把预训练transformer中的信息流,类比成画家在流水线上作画的过程。

这个类比很直观:画布(输入)被依次递给一群画家,有的画家擅长画鸟,有的擅长画轮子。每个画家从下一个画家那里拿到画布,然后决定要不要添上几笔,或者干脆跳过(利用残差连接),再传给上一个画家。当然,这不是严格的理论,只是一种思考工具。顺着这个思路,研究团队提出了一连串问题,对预训练的大模型做了一系列实验——注意,没有对模型做任何微调。实验对象是仅解码器(Llama)和仅编码器(BERT)两类模型,测试了标准transformer执行策略的各种变化。
各层是否都在使用相同的表征空间?
为了回答这个问题,研究者先看了一个朴素现象:如果跳过某层,或者把相邻两层的顺序互换,模型会崩溃吗?拿Llama2-7B来说,第6层通常等着接收第5层的输出,如果突然给它第4层的输出,它会不会“直接罢工”?
实验结果(图2)显示,除了第一层和最后几层,中间层对这种跳层或交换表现得相当皮实。这说明中间层共享一套表征空间,而“外围层”(第一层和最后几层)用的则是另一套。为了进一步验证,他们又测量了不同层隐藏状态之间的平均余弦相似度(图3),发现所有中间层的一致性很高。于是第一个问题的答案出来了:是的,中间层共享一个共同的表征空间。
所有层都是必要的吗?
既然中间层共享空间,那是不是意味着有些层是多余的?研究团队设计了“跳过层”实验:把第N层的输出直接喂给第N+M层的输入,中间跳掉M−1层(图1a)。目的是看看第N+M层能不能理解来自第N层的激活——毕竟它原本只接收第N+M−1层的输出。
结果(图4)显示,无论是在Llama2-7B还是BERT-Large上,性能都只是适度下降,并没有崩溃。所以第二个问题的答案也很明确:不,至少可以删除一些中间层而不会出现灾难性故障。
中间层都执行相同的功能吗?
既然共享表征空间,会不会意味着这些层做的是同样的事?研究者把中间层的权重替换成某一层(比如中心层)的权重,然后让模型在被替换的层上循环T−2N+1次(T是总层数,Llama2-7B有32层,BERT-Large有24层)。结果(图5)非常糟糕:随着被替换的层数增加,模型在基准测试上的得分迅速下滑,甚至比后面其他实验中的任何方法都差(图11也证实了这一点)。
所以结论是:中间层虽然共享空间,但各自执行不同的功能,让它们共享权重根本行不通。
层的顺序重要吗?
既然中间层功能不同,那它们的排列顺序有多大意义?研究者做了两组实验:第一组,把中间层按训练时的相反顺序运行;第二组,随机打乱顺序(取10个随机种子的平均值)。结果(图6、图7)显示,两种情况下模型性能都出现了下降,但下降比较平缓。而且更有意思的是,随机打乱的效果竟然比完全反过来好。原因可能是随机打乱在某些程度上保留了层之间“大小关系”(比如层i在层j之后,i>j),而完全反过来彻底打破了这种关系。这说明顺序虽然重要,但并非不可动摇,层本身的功能比顺序更关键。
这些层可以并行运行吗?
如果把中间层全部并行起来,把它们的平均结果送给最后的几层,会怎样?实验(图8)显示,在大部分基准测试中性能平缓下降,只有面对GSM8K里的数学应用题时表现明显变差。也就是说,并行运行在绝大多数情况下可行,唯独对付需要顺序逻辑的数学问题有点乏力。
对于某些任务来说,顺序是否比其他因素更重要?
纵观所有“改造”过的模型,只要碰到抽象推理(ARC)或数学推理(GSM8K)这类基准,性能往往掉得最快。研究者认为,这是因为逐步推理任务对层级顺序特别敏感,远高于那些靠语义理解就能搞定的常识任务。推理需要同时把握结构和含义,不像拼贴画那样可以随意安排各部分。这就像画一幅由各种元素组成的拼贴画,顺序无所谓;但若是画一张精确的建筑透视图,每一笔的顺序就至关重要了。结论是:数学和推理任务对层顺序的依赖性更强,而纯语义任务就没那么讲究。
循环有助于层之间并行吗?
回到画画的比喻:画家在画一幅画时,不会一开始就把所有东西画满,而是先画出车身,再根据车身去画车轮。如果transformer的层也能先得到正确的“草图”,是不是就能更高效地贡献后面的笔画?实验设计是:把并行层的平均输出重新反馈回同一层,进行固定次数的迭代。图9展示了循环3次的结果:效果显著优于单次并行。特别是当起始层N设定为15(Llama2-7B)或11(BERT)时,只有一层被影响,三次循环的效果等于把那层重复了三次,性能跟完整模型一样好。进一步实验(图10)显示,除了极端情况(几乎并行所有层),最佳迭代次数与并行化层数大致成线性关系。结论是:最佳迭代次数与并行化层数成正比。
如何调整层,对模型性能的影响最小?
最后,研究者把所有实验放在一起比了比(图11)。中间重复——用同一层的副本替换其他层——表现最差,很快掉到随机基线水平。相反,随机化层顺序和循环并行的影响最小。整体上,这些实验中的性能下降都比较平缓,但研究者自己也不完全清楚为什么层能扛住这么多扰动,这个问题只能留到未来的研究中继续探索了。
更多细节请参见原论文:https://arxiv.org/pdf/2407.09298v1
