Transformer 中的信息流动机制:最新研究揭开其内部运作
Transformer 架构中的信息流动机制,近期被一项名为'Transformer Layers as Painters'的研究深入揭示。该研究由 AI 初创公司 Sakana AI 和 Emergence AI 的团队完成,通过一系列实验解答了关于 Transformer 层结构的关键问题。
核心问题与研究动机
研究团队从 Transformer 内部工作原理出发,重点探讨了以下问题:
- 所有层都是必要的吗?
- 中间层在做同样的事吗?
- 层的顺序重要吗?
- 如果跳过一些层(如第 4 层输出接到第 6 层)会怎样?
- 随机打乱层的顺序(如 4-6-5-7)又会怎样?
实验动机源于一个假设:由于训练过程中使用了残差连接,神经网络的中间层可能使用了一个共同的表征空间。如果层能共享一个表征空间,将对后续条件计算或向预训练 Transformer 模型动态添加新知识及下游应用产生重要影响。
实验模型选择与基准
实验在 decoder-only 和 encoder-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 用于测困惑度。
- BERT: 采用 GLUE 基准并遵循其评估指标。
性能评估方面,Llama2 提供了基准测试的标准化中位数,将性能从 0 到 1(模型最优性能)进行量化。
关于 Transformer 的 8 大问题解析
1. 层是否使用相同的表征空间?
为确定不同层是否共享相同的表征空间,研究人员检验了 Transformer 对于跳过特定层或更改相邻层顺序的鲁棒性。
实验发现,除了最前面的和最末尾的几层,Llama2-7B 跳过或改变层序表现出很好的鲁棒性。这表明中间层共享一个表征空间,而中间层与'外层'(最前面的和最末尾的几层)具有独立的表征空间。
为了进一步证实这一假设,研究人员测量了不同模型中不同层的隐藏状态激活之间的平均余弦相似度。结果显示所有中间层之间的一致性较高。例如,底部第四层的激活与顶部第四层的激活高度相似。对于 40 层的 Llama2-13B,可以看到这些层按相似性可划分成 4-5 个组:第 0 层,1-3 层,中间层,然后是最后一两层。
这表明模型可能对'开始'、'中间'和'结束'层具有三个不同的表征空间。高余弦相似度证明有一个共享的表征空间,低相似度表明这些空间不是共享的。这进一步证明了至少中间层的表征空间是共享的。
2. 所有层都是必要的吗?
为了进一步验证中间层的表征空间真正共享,研究人员进行了层跳过实验(未进行任何微调)。
具体来说,将第 N 层的输出直接传递为第 N+M 层的输入(M>1),从而'跳过'了 M-1 层。此类实验中,研究人员正常执行第一层和最后 N-1 层,而跳过或修改第 N+1 到第 T-N 层(T 是模型总层数)。
结果显示,在多个基准测试中,Llama2-7B 和 BERT-Large 的性能均逐渐下降(随着跳过层数递增)。这一结果揭示了:不是所有层都是必要的,至少省略部分中间层不会对整体性能造成严重影响。
3. 中间层是否都执行相同的功能?
如果中间层共享一个共同的表征空间,这些层是否多余?
研究人员重新进行了'跳过'实验,但这次不是跳过中间层,而是用最中心层的权重替换了所有这些中间层的权重。实际上就是在最中心层上循环执行了 T-2N+1 次。
结果基准测试中,随着被替换的层数增加,模型性能迅速下降。而且性能下降速度比仅仅跳过某些层要严重得多,这种权重替换极具破坏性。
因此,中间层各执行不同的功能并非多余,中间层之间共享权重会产生灾难性后果。


