深入解析 Transformer 内部运作原理：基于 Sakana AI 最新研究
背景介绍
七年前，Google 团队发表的论文《Attention is all you need》提出了 Transformer 架构，彻底颠覆了自然语言处理（NLP）乃至整个深度学习领域。该架构摒弃了传统的循环神经网络（RNN）和卷积神经网络（CNN），完全依赖自注意力机制（Self-Attention）来处理序列数据，实现了并行计算和长距离依赖建模。
如今，从 GPT 系列到 Llama 系列，各家大模型均以 Transformer 架构为基础。然而，Transformer 内部究竟是如何运作的？中间层到底在做什么？这些问题长期以来仍被视为黑盒。去年，Transformer 论文作者之一 Llion Jones 宣布创立人工智能公司 Sakana AI。近期，Sakana AI 发表了一篇题为《Transformer Layers as Painters》的论文，深入探究了预训练 Transformer 中的信息流，并针对仅解码器和仅编码器冻结模型进行了一系列实验。值得注意的是，该研究没有对预训练模型进行任何类型的微调，而是直接观察原始权重下的行为。
论文地址：https://arxiv.org/pdf/2407.09298v1
该研究认为 Transformer 的内部机制（特别是中间层）可以类比画家作画流水线来理解。作画流水线通常是将画布（输入）传递给一系列画家。有些画家擅长画鸟类，而另一些画家则擅长画轮子。每个画家从其下一级画家那里收到画布，然后决定是否给画作添加一些笔画，或者只是将其传递给其上一级画家（使用剩余连接）。这个类比提供了一个直观的思考框架，帮助研究者验证关于 Transformer 层功能的假设。
核心问题与实验设计
受这个类比的启发，该研究测试验证了以下八个关键假设：
- 各层是否都在使用相同的表征空间？
- 所有层都是必要的吗？
- 中间层都执行相同的功能吗？
- 层的顺序重要吗？
- 这些层可以并行运行吗？
- 对于某些任务来说，顺序是否比其他因素更重要？
- 循环有助于层之间并行吗？
- 哪些变体对模型性能影响最小？
该研究对预训练 LLM 进行了一系列实验，包括试验标准 Transformer 执行策略的变化，并在仅解码器 (Llama) 和仅编码器 (BERT) 模型的各种基准上测量这些变化对模型性能的影响。
1. 各层是否都在使用相同的表征空间？
为了回答不同层是否使用相同的表征空间，作者测试了 Transformer 在跳过特定层或切换相邻层的顺序时是否具有稳健性。例如，在 Llama2-7B 中，第 6 层通常期望接收第 5 层的输出。如果给第 6 层以第 4 层的输出，它是否会出现「灾难性」的行为？
在实验中，研究人员发现除了第一层和最后几层之外，Llama2-7B 的各层对跳层或切换层都相当稳健。这表明中间层共享一个表征空间，且与「外围层」（第一层和最后几层）拥有不同的表征空间。
为了进一步验证这一假设，作者测量了基准中模型（Llama2-7B、Llama2-13B 和 BERT-Large）不同层的隐藏状态激活之间的平均余弦相似度。结果显示所有中间层之间的一致性很高。这表明模型可能具有「开始」、「中间」和「结束」层的三个不同的表征空间。
结论： 是的，中间层似乎共享一个共同的表征空间。
2. 所有层都是必要的吗？
为了进一步测试中间层的重定向空间是否真正共享，该研究尝试了「跳过层」，即将第 N 层的输出直接发送到第 N + M 层（其中 M > 1）的输入中，从而「跳过」M − 1 层。该实验是为了看看第 N + M 层是否可以理解第 N 层的激活，尽管它仅根据从第 N + M − 1 层发来的输入进行训练。
实验显示，Llama2-7B 和 BERT-Large 在许多基准测试上性能均出现适度下降，但并未发生灾难性故障。
结论： 不，至少可以删除一些中间层而不会发生灾难性故障。这为模型剪枝提供了理论依据。
3. 中间层都执行相同的功能吗？
如果中间层都共享一个共同的表征空间，这是否意味着除此之外的中间层是多余的呢？为了测试这一点，研究者们重新运行了前一子节中的「跳过」实验，他们将中间层的权重替换为中心层的权重，有效地在被替换的每一层上循环 T - 2N + 1 次，其中 T 是总层数（Llama2-7B 为 32 层，BERT-Large 为 24 层）。
