文献阅读：Foundation Transformers

• 文献阅读：Foundation Transformers
  • 1. 文章简介
  • 2. 模型结构
    • 1. Sub-LN
    • 2. Initialization
  • 3. 实验效果
    • 1. NLP任务
      • 1. 语言模型上效果
      • 2. MLM模型上效果
      • 3. 翻译模型上效果
    • 2. Vision任务上效果
    • 3. Speech任务上效果
    • 4. 图文任务上效果
  • 4. 结论 & 思考

• 文献链接：https://arxiv.org/abs/2210.06423
• GitHub: https://github.com/microsoft/unilm

1. 文章简介

这篇工作是DeepNet的一个后续工作，本质上就是deepnet解决了深层transformer的稳定性的情况下，考察其在各类任务下的泛用性。

关于DeepNet的工作，我司已经将其开源了（https://github.com/microsoft/unilm），网上也有不少关于DeepNet的文章解读，我自己也写过一篇拙作（文献阅读：DeepNet: Scaling Transformers to 1,000 Layers）。

本质上来说，DeepNet的工作就是通过优化transformer模型的参数初始化的方式增加模型参数反向传递的稳定性，从而可以达到超多层模型的训练。

模型的参数初始化方式推导其实还是挺繁琐的，不过真到实现上倒还是挺简单的，改改也化不了太多的时间，不过比较苦逼的是我在我们自己的场景下面试了一下，倒是没有看到很好的一个效果，深层的也没train起来，估摸着哪里理解还是有问题吧，不过整体的效果倒是应该没啥怀疑了，而且貌似这玩意也已经被用到了GPT4里面，应该是一个确实有效的策略，后面还是得回去检查检查自己的code……

Anyway，话扯远了，回来说这篇工作吧，除了DeepNet中介绍的参数初始化方式之外，文中还引入了Sub-LN这个简单的策略，从而给出了文中的Foundation Transformer的架构，使得其可以在各类任务的transformer架构当中都获得较好的效果，而在此之前，各类任务当中的主流模型的transformer架构是略有不同的。

下面，我们就来具体看看这篇工作的细节。

2. 模型结构

下面，我们来看一下文中对于模型结构的细节描述，如前所述，主要就是以下两方面的内容介绍：

1. Sub-LN
2. Initialization

下图就是文中的两个技术点的一个整体描述：

下面，我们来分别对两者进行一下考察。

1. Sub-LN

首先，我们来看一下Sub-LN的内容。

如上图所示，所谓的Sub-LN层，本质上就是对self-attention层以及FFN层的Layer Normalization方式进行一定的优化，不再是在最前侧或者最后侧使用一个Layer Normalization，而是在输入与输出的两个线性层之前均加入一个Layer Normalization。

其具体的函数表达如下：

1. Self-Attention

   $$\{\begin{array}{rl}Q& =W^{Q}LN(x)\\ K& =W^{K}LN(x)\\ V& =W^{V}LN(x)\\ y& =x+W^{O}LN(Attn(Q,K,V))\end{array}$$

2. FFN

   $$\{\begin{array}{rl}z& =W^{1}LN(x)\\ y& =x+W^{2}LN(\varphi (z))\end{array}$$

   其中，$\varphi$为激活函数，图示当中使用Relu作为激活函数。

2. Initialization

而关于文中的intialization的部分，其实和前作DeepNet是基本保持一致的，其中的推理部分这里就不过多展开了，事实上也不算看的很明白，但是结果倒是比较简单，如上图所示，只需要满足：

def subln_init(w):
    if w in ["ffn", "v_proj", "out_proj"]:
        nn.init.xavier_normal_(w, gain=γ)
    elif w in ["q_proj", "k_proj"]:
        nn.init.xavier_normal_(w, gain=1.0)

其中，xavier_normal_是pytorch当中的一个初始化函数，其文档链接如下：

• https://pytorch.org/docs/stable/nn.init.html

简言之，它就是一个正态分布$N(\mu ,{\sigma }^2)$，对应参数为：

$$w\sim N(0,\frac{2{\gamma }^2}{f_{in}+f_{out}})$$

其中，$f_{in}$和$f_{out}$为参数矩阵$w$的维度，对于二维矩阵，他就是矩阵的行数和列数。

换用tensorflow的代码实现的话，就是：

import tensorflow as tf

def cal_fan_in_and_fan_out(shape):
    if len(shape) < 2:
        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")
    fan_out, fan_in = shape[:2]
    receptive_field_size = 1
    if len(shape) > 2:
        for d in shape[2:]:
            receptive_field_size *= d
    fan_in, fan_out = fan_in * receptive_field_size, fan_out * receptive_field_size
    return fan_in, fan_out

def xavier_normal_initializer(gain, fan_in, fan_out):
    std = gain * tf.sqrt(2.0 / (fan_in + fan_out))
    return tf.random_normal_initializer(stddev=std)

而具体到模型结构，$\gamma$的取值详见上表，这里转录一下如下：

Architecture	Encoder	Decoder
Encoder-Only (e.g. BERT)	$\sqrt{log2N}$	-
Decoder-Only (e.g. GPT)	-	$\sqrt{log2M}$
Encoder-Decoder (e.g. NMT,T5)	$\sqrt{\frac{1}{3}log3M\cdot log2N}$	$\sqrt{log3M}$

3. 实验效果

下面，我们给出文中的实验结果。

1. NLP任务

1. 语言模型上效果

首先，对于语言模型，也就是文中所说的纯Decoder结构，实验结果如下：

可以看到：

• 对比Pre-LN和Normformer，Foundation Transformer训练得到的LLM模型无论是在zero-shot, one-shot还是few-shot上面都获得了更好的效果。

2. MLM模型上效果

可以看到：

• 在MLM预训练模型当中，Foundation Transformer训练得到的MAGNETO模型同样可以获得更好的模型效果表达。

3. 翻译模型上效果

文中同样在翻译任务上对Encoder-Decoder架构的模型效果进行了考察，得到结果如下：

可以看到：

• MAGNETO模型在翻译任务上同样可以获得更好的模型效果表达。

2. Vision任务上效果

除了NLP任务之外，文中还考察了Foundation Transformer在CV任务当中的效果，即对比了Foundation Transformer以及ViT。

得到结果如下：

可以看到：

• MAGNETO在CV任务当中也同样是有效的。

3. Speech任务上效果

同样的，文中还在语音识别任务当中对Foundation Transformer进行了一下考察，得到结果如下：

可以看到：

• 结论与之前相同，Foundation Transformer在语音识别任务当中同样有效。

4. 图文任务上效果

最后，文中还在多模态任务当中进行了考察，得到结果如下：

结论同样成立，Foundation Transformer依然是work的。

4. 结论 & 思考

整体来说，文中给出的Foundation Transformer框架算是一个更加robust且效果更优的transformer架构，不但可以适用于各类任务，且对应的模型效果也更加优秀。

其对应的关于sub-ln的部分，事实上个人觉得配合苏剑林的下面这篇博客一起看效果更优：

• 为什么Pre Norm的效果不如Post Norm？

不过原理层面坦率地说我倒是一直没有仔细去按照他的思路完整推导过一遍，因此就不发表评论了，只是就结论和代码迁移应用而言，倒是还算简单，后续可以在我们实际的工作当中使用一下看看效果。