0. TL; DR

本文探讨了Transformer模型中前馈层（feed-forward layers）的作用，发现这些层实际上充当着键值记忆（key-value memories）的角色。每个键与输入文本中的特定模式相关联，而每个值则诱导输出词汇表上的分布。研究发现，这些模式是人类可解释的，较低层倾向于捕获浅层模式，而较高层则学习更语义化的模式。此外，输出分布的构建是一个自下而上的过程，通过残差连接逐步细化。这些发现有助于我们更好地理解Transformer模型的内部工作机制，并为未来的研究提供了新的方向。

1. 背景介绍

Transformer模型已成为自然语言处理（NLP）领域的核心，其成功主要归功于自注意力机制。尽管许多研究致力于分析自注意力层的功能，但Transformer模型中大部分参数（约三分之二）实际上是用于前馈层。

然而，这些前馈层在网络中的作用仍然未被充分研究。本文旨在揭示前馈层在Transformer语言模型中的功能，特别是它们如何作为模式检测器和记忆单元工作，以及模型的最终输出分布是如何逐步构建的。

2. 方法介绍

Transformer模型由自注意力层和前馈层交织而成。每个前馈层是一个位置函数，独立处理每个输入向量。前馈层可以表示为：

\[FF(x)=f(x⋅K^\top )⋅V\]

其中$K,V∈R^{d_m×d}$是参数矩阵，$f$ 是非线性函数（如ReLU）。与神经记忆相比，前馈层几乎与键值神经记忆相同，唯一的区别是神经记忆使用softmax作为非线性函数，而标准Transformer的前馈层不使用归一化函数。

假设前馈层中的键向量作为输入序列的模式检测器，每个键向量对应于输入前缀的一个特定模式。使用16层Transformer语言模型进行实验，该模型在WikiText103数据集上进行训练。随机抽取每层的10个键（共160个）进行分析。通过计算每个输入前缀与键的内存系数，并检索触发每个键的最优示例，发现每个键都与人类可识别的模式相关联。

作者进一步分析了与键对应的值向量，将每个值向量转换为输出词汇表上的概率分布，并分析这些分布与相应键的下一个标记的一致性。结果表明，上层的值向量倾向于为下一个标记分配更高的概率，显示出键和值之间的非平凡预测能力。

作者还探讨了多个层中的多个内存单元如何聚合形成模型范围内的预测。每个前馈层通过残差连接组合多个内存单元，形成一个与各个组件内存值分布不同的分布。这些层级分布然后通过残差连接进行细化，最终形成模型的输出。