Transformer的知识在哪里：从FFN出发

前情提要：

前阵子，或许更早之前，笔者关于 MoE 的一些研究——如果称得上是研究，而非民科的话——让人不禁思考 Transformer 的知识到底存在哪里？这个问题的价值不言而喻，用来打发时间也是个不错的选择。

可惜的是，目前我关于它的思考仍旧十分细碎，或许也称不上有什么新意，这里只是做下暂且的梳理，一方面以作个人备忘，一方面也会考虑可读性，希望能帮助看到此篇的人度过愉快的厕所时光。

矩阵乘法

最简单的神经网络，或许就是矩阵 -> 激活函数 -> 矩阵 -> 激活函数的堆叠了吧。

假如这个网络的输入是个向量，那么在每一个矩阵层，进行的计算无非是一个向量和一个矩阵的乘法。

即便以笔者匮乏的数学知识，也可以知道，向量和矩阵的乘法，可以看成是以向量为权重，对矩阵的所有行（或者列）进行加权求和。

从这个角度去理解的结果是什么呢？我们可以轻松的把一个神经网络分为头重脚轻的两部分：权重生成器，和知识存储器（其实就是个矩阵）。输入经过好多好多层复杂的运算，为的就是把最后的一层混合起来。

（你或许想到了残差连接这个反例，别急别急。）

那么然后呢？这就有了一丝丝可解释性了。假如我们可以评估知识存储器里的每个知识对于最终预测的影响，那么我们就可以通过查看权重，看看模型是否按照我们的理解在工作。

相关工作：Backpack Language Models 与 残差连接

好吧，口说无凭，还会被当成民科，这里就摆出一篇相关论文：Backpack Language Models -- 23 ACL Outstanding Papers。

这篇文章的出发点是可解释性。用我们刚才提出的视角，这篇文章的方法，完全可以看成是：权重生成器——一整个Transformer，和知识存储器——一个矩阵。

提前做个剧透，它的方法，一言以蔽之：给定一个输入的向量，经过 Transformer 的计算，得到这个输入对于某些向量（知识）的喜爱程度，然后根据这个喜爱程度，把知识混合起来。

那么，具体是怎么做的呢？下面就是繁琐的细节了。

背景知识：Transformer，或者说 decoder 的结构，我们可以看成是:

• embedding layer
• Transformer Blocks
• embedding layer (decoding layer)

这篇的文章呢，在这里面插了两层，变成了：

• embedding layer
• Transformer Blocks
• Attention Layer
• Backpack Layer
• embedding layer (decoding layer)

这插入的两层被加粗了，名字是我随便取的。

大家都知道， Transformer Blocks 的输出是一个 token embedding 的序列。这里假设有 N 个 token，每个 embedding 的维度是 d，那么输出就是个矩阵，维度是 N * d。

Attention Layer 做的工作，就是把这个矩阵拿过来，计算两两 token 之间的 k 个权重，最终得到一个 N * N * k 的矩阵。（为什么是 k 个？怎么得到 k 个权重？无视也行，完全不重要。做法很简单，参考多头注意力。）

通俗点理解，每个 token 都需要给出它对其它所有 token 的喜欢程度，而且还要用 k 个标准给出 k 个不同的分数。所以每个 token 的权重矩阵是 N * k 。

Backpack Layer 里面，有一个预先准备好的矩阵。在这个矩阵里，每个词有 k 个向量。假设词表大小是 V，那么这个矩阵的大小就是 V * k * d——我们只需要用到一起出现的那些词的向量，矩阵大小是 N * k * d。

对于一个 token，它根据权重，把 Backpack Layer 里的那些向量混合，得到 token 最终的输出向量。

繁琐的细节终于解释了，总结下！

Backpack Language model 里，embedding layer + Tranformer Blocks + Attention Layer可以看成是一个活生生的权重生成器——非常重量级的权重生成器。然后，Backpack Layer 里存储的向量，就是所谓的知识啦。

这么训练的模型呢，也就具备了一点点可解释性。具体还是看论文吧。

残差连接

笔者呢，前些时日，发了个想法。根据我们提出的权重生成器和知识存储器这个角度，来理解残差连接的作用。

残差连接的直觉
如果有个神经网络，输入是向量，输出也是向量。

通常来说（或者假定），没有残差的情况下，最后一次的计算是向量 x 与一个矩阵相乘 A。比如：

y = xA。

那么直观点来看，这个操作就是将 A 的各行进行混合。

也就是说，我们可以把神经网络拆成两部分。第一部分根据输入得到一个权重，第二部分根据这个权重把 A 的各行进行混合。

如果我们假定模型里有知识，那么我们只能认为知识存储在了A矩阵里。多么悲观。

可是有了残差就不一样了。在模型里进行的每一次矩阵乘法（每一次权重矩阵的混合）的结果都会被传入到下一层。也就意味着每一层的权重或多或少都贡献了知识。

如何验证，残差是否是按照上面的假设发挥作用的呢？

一个很简单的做法，就是移除掉模型，比如Transformer，最后一层的残差。仅仅把向量与矩阵相乘的结果作为输出。

大意就是，模型的能力应该会受到最后面的知识存储器的限制。而残差连接，巧妙的把每一层都当做了知识存储器。怎么验证这一点呢？

就是移除掉模型，比如Transformer，最后一层的残差。仅仅把向量与矩阵相乘的结果作为输出。

没错，Backpack Language Models 就是个验证。当它把 k 设置为 1 时，模型的性能特别差。不无是个例证。

未完待续

什么，你问我你看完了这个收获了什么？嘛，总该不是完全没有收获吧。

写累了，祝大家吃好喝好，下次再见。