打乱/跳过Transformer层会怎样?最新研究揭开其信息流动机制

  西风发自凹非寺

  量子位公众号 QbitAI

  Transformer 中的信息流动机制,被最新研究揭开了:

  所有层都是必要的吗?中间层在做同样的事吗?层的顺序重要吗?

  如果跳过一些层,比如第 4 层输出接到第 6 层会怎样。随机打乱层的顺序,比如4-6-5-7 又会怎样。

  最近一项名为“Transformer Layers as Painters”的研究火了,由来自 AI 初创公司 Sakana AI、Emergence AI 的研究团队完成。

  他们从 Transformer 内部工作原理出发,经过一系列实验对以上问题得出了结论。团队表示深入理解这些原理不仅能提高现有模型利用效率,还能帮助改进架构开发新的变体。

  谷歌 DeepMind 研究员、ViT 作者 Lucas Beyer 看过后直接点了个赞:

很棒的总结!尽管一些实验在之前的研究中已经被展示过了,但我喜欢你添加的新细节,特别是强调了“推理”类任务比其他任务受影响更大!

  还有不少学者、工程师也表示强烈推荐。

敢打赌,其中一些见解最终将会用于改进 Transformer。

其中的实验再次证实了:复制层对创造性任务有帮助,但对推理任务通常无效;改变层的顺序行不通;剪枝在中间层效果最佳,但仍需要进行修复调整。

  所以,在这项研究中,研究团队都进行了哪些实验?回答了哪些问题?

  实验模型选择和基准

  先来看一下实验配置~

  实验在decoder-onlyencoder-only模型上进行。

  其中 decoder-only 模型选择的是Llama2,主要研究 32 层、70 亿参数的 Llama2-7B,扩展实验中也包含 13B(40 层)和 70B(80 层)模型。

  encoder-only 模型选择的是BERT,有 24 层、3.4 亿参数。

  研究人员使用这些模型的标准预训练 checkpoints。在所有实验中,模型都是冻结的,除 BERT 的评估中包含一个标准的微调步骤,其它情况未通过微调等方法修改模型参数。

  基准测试方面,Llama2 使用以下标准基准:ARC(科学考试问题)、HellaSwag(常识问题)、GSM8K(数学题),WinoGrande(常识推理)、LAMBADA(词汇预测)。其中 LAMBADA 用于测困惑度,与训练期间使用的原始 token 预测最接近。

  对于 Llama2 的性能评估,提供了基准测试的标准化中位数,将性能从 0 到1(模型最优性能)进行量化。

  对于 BERT,采用 GLUE 基准并遵循其评估指标,包括基准的未标准化平均分。注意,标准的 BERT 评估包括一个微调步骤,因此对模型进行了适应性调整。在附录中研究人员也展示了一个只有模型头部可以调整的评估结果。

  实验动机最初源于这样一个问题:

  是否可以将多个层以某种方式合并成一个可能更大的单一层?假设可能由于训练过程中使用了残差连接,神经网络的中间层可能使用了一个共同的表征空间。(对于标准的多层感知机来说不成立,它们之间没有促使共同表征或层间排列一致性的机制)

  如果层能共享一个表征空间,将对后续条件计算或向预训练 Transformer 模型动态添加新知识及下游应用产生重要影响。

  关于 Transformer 的 8 大问题

  层是否使用相同的表征空间?

  为确定不同层是否共享相同的表征空间,研究人员检验了 Transformer 对于跳过特定层或更改相邻层顺序的鲁棒性。

  例如,在 Llama2-7B 模型中将输出流从“第 4 层->第 5 层->第 6 层”的正常顺序,改为“第 4 层->第 6 层”,跳过第 5 层,会怎样?

  又或者将第 4 层的输出送到第 6 层,然后将第 6 层的输出送到第 5 层,再送到第 7 层,会怎样?

  如下图所示,实验发现除了最前面的和最末尾的几层,Llama2-7B 跳过或改变层序表现出很好的鲁棒性

  也就是说,中间层共享一个表征空间,中间层与“外层”(最前面的和最末尾的几层)具有独立的表征空间。

  为了进一步证实这一假设,研究人员测量了不同模型(Llama2-7B、Llama2-13B 和 BERT-Large)中不同层的隐藏状态激活之间的平均余弦相似度,并跨基准测试进行了比较。

  下图 3 展示了所有中间层之间的一致性。例如,底部第四层的激活与顶部第四层的激活高度相似。对于 40 层的 Llama2-13B,可以看到这些层按相似性可划分成4-5 个组:第 0 层,1-3 层,中间层,然后是最后一两层。

  这表明模型可能对“开始”、“中间”和“结束”层具有三个不同的表征空间。研究人员还发现,“开始层”的数量似乎随着模型总层数的增加而增加。

  此外,高余弦相似度可能证明有一个共享的表征空间,低相似度更能表明这些空间不是共享的。而上图 3 中 Llama2-7B 的数据与图 2 所示的性能结果高度一致,这进一步证明了:

  至少中间层的表征空间是共享的。

  所有层都是必要的吗?

  为了进一步验证中间层的表征空间真正共享,研究人员还进行了层跳过实验(实验中未进行任何微调)。

  具体来说,将第N层的输出直接传递为第N+M层的输入(M>1),从而“跳过”了M-1 层,如下图所示。

  原本第N+M层仅针对来自第N+M-1 层的输入进行训练,那么现在它能否理解第N层的激活?

  此类实验中,研究人员正常执行第一层和最后N-1 层,而跳过或修改第N+1 到第T-N 层(T是模型总层数)。

  如下图4,在多个基准测试中,Llama2-7B 和 BERT-Large 的性能均逐渐下降(图从左至右展示了跳过层数逐渐递增的变化)。这一结果揭示了:

  不是所有层都是必要的,至少省略部分中间层不会对整体性能造成严重影响。

  中间层是否都执行相同的功能?

  如果中间层共享一个共同的表征空间,这些层是否多余?

  为了解答该问题,研究人员重新进行了前面的“跳过”实验,但这次不是跳过中间层,而是用最中心层的权重替换了这些所有中间层的权重,如下图所示。

  实际上就是在最中心层上循环执行了T-2N+1 次,其中T是模型总层数(Llama2-7B 为 32 层,BERT-Large 为 24 层)。

  结果基准测试中,随着被替换的层数增加,模型性能迅速下降。而且性能下降速度比仅仅跳过某些层要严重得多,这种权重替换极具破坏性。

  因此,中间层各执行不同的功能并非多余,中间层之间共享权重会产生灾难性后果。

  层的顺序重要吗?

  上面实验表明中间层虽共享表征空间,却在该空间上执行不同操作。那么这些操作顺序重要吗?研究人员进行了两组实验。

  首先,将中间层按照与其训练顺序相反的顺序(逆序)执行。将第T-N 层的输出传递给第T-N-1 层,依此类推,直至第N层,然后将该层的输出传至最后的T-N 层。

  如下图:

  第二个实验,随机排列中间层顺序,并平均 10 个随机种子结果。

  结果如下图,两种情况模型都呈现出缓慢的性能下降

  这里剧透一下下面的一项实验结果,无论是逆序还是随机顺序,模型表现均优于直接跳过这些层,说明即使层在非训练顺序的输入上运行,依然能进行有效输出。

  因此,层顺序重要吗?结论是:

  层顺序调整对性能有一定影响,随机顺序和逆序都表现出一定的性能退化。

  值得注意的是,随机顺序性能优于逆序。可能是因为逆序与训练时的顺序完全相反,而任何随机顺序都至少保持了一些顺序上的连贯性(即某层i总在另一层j之后,其中i>j)。

  可以并行运行这些层吗?

  如果层的存在,即没有被跳过,比它们执行的顺序更重要,那么是否可以考虑独立地运行这些层,然后将它们的结果合并?如下图所示。

  研究人员进行了一个实验,不是跳过第N到第T-N 层,而是并行运行这些中间层,然后将它们的平均结果传递到最后的N层。

  结果如下图所示,除了 GSM8K 数学题基准外,所有基准测试都表现出缓慢的性能退化。

  有趣的是,并行层的表现优于跳过层,但不如逆序运行层。

  总之,可以并行运行这些层吗?答案是:可以,数学为主的基准测试除外。

  对于某些任务,顺序是否更重要?

  大多数变体(包括逆序、跳过和并行)在抽象推理 ARC 或数学推理 GSM8K 基准中,表现出最快速的性能下降。

  可以解释为逐步推理任务比“语义”任务(如 Winogrande 或 HellaSwag)对层顺序的变化更为敏感。

  这是因为推理任务需要结合结构和语义双重信息,而 HellaSwag 这类任务仅需语义即可完成。

  通过实验,研究人员得出结论:数学和推理任务比“语义”任务更依赖顺序。

  迭代对并行层有帮助吗?

  如果把 Transformer 内部运行机制比作是画一幅画的过程:画布(输入)在一些画家之间传递,一些画家专门画鸟,一些则更擅长画轮子……每个画家都依次从另一位画家手里接过画布,然后决定对这幅画进行补充,还是将其直接传递给下一位画家(使用残差连接)。

  可以想象,某些层在收到适当的输入时才会对画作进行“补充”。例如,如果“画轮子”的画家先看到汽车的车身,才更有可能画上轮子。

  在 Transformer 中,某些层可能只有在接收到适当的输入时才会对前向传递发挥作用,而不是通过残差连接将输入直接传递出去。

  这么来看的话,那么相比于仅执行一次并行层,迭代执行并行层应该会提高性能。

  研究人员通过将并行层的平均输出回馈到同一层并固定迭代次数来进行测试,如下图:

  下图 9 中,研究人员展示了并行层迭代 3 次的结果,这种方法显著优于仅执行一次并行层。

  唯一的例外是在 Llama2-7B 的起始层N为 15 或 BERT 的起始层N为 11 时。在这种情况下,循环并行 3 次的效果相当于仅重复中间层 3 次,此时的并行层等同于完整模型。

  研究人员还用不同的迭代次数重复进行了实验。

  下图展示了 Llama2-7B 的性能随并行层数M和迭代次数的变化而变化。

  每个M的最佳迭代次数用红框表示。除了M=29 和M=31(几乎并行所有层)外,最佳迭代次数大致与并行层数成线性比例。

  因此结论是:迭代对并行层有帮助,最佳迭代次数与并行层数成比例。

  哪些变体对性能损害最小?

  最后,研究人员将实验中的所有不同变体在同一图表上进行了比较。

  结果显示,重复单一层(如上面提到的用同等数量的最中心的层替换中间层)效果最差,性能迅速退化至随机基准线。

  迭代并行和随机层顺序性能退化最小,其中迭代并行在 BERT 和 Llama2-7B 中表现最好。

  论文附录中还补充了更多实验结果,感兴趣的家人们可以查看原论文。

  论文链接:https://arxiv.org/abs/2407.09298v1

  参考链接:https://x.com/A_K_Nain/status/1812684597248831912