文献阅读：Foundation Transformers

最新推荐文章于 2024-09-15 13:26:49 发布

Espresso Macchiato

最新推荐文章于 2024-09-15 13:26:49 发布

阅读量1.2k

点赞数

分类专栏：文献阅读文章标签： Transformer Magneto SubLN NLP LLM

本文链接：https://blog.csdn.net/codename_cys/article/details/131274469

版权

文献阅读专栏收录该内容

55 篇文章 2 订阅

订阅专栏

文献阅读：Foundation Transformers

文献链接：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. 模型结构

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

Sub-LN
Initialization

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

在这里插入图片描述

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

1. Sub-LN

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

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

其具体的函数表达如下：

Self-Attention

$\left\{ \begin{aligned} 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{aligned} \right.$
FFN

$\left\{ \begin{aligned} z &= W^{1} LN(x) \\ y &= x + W^{2} LN( \phi(z)) \end{aligned} \right.$

其中， $\phi$ 为激活函数，图示当中使用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)$ ，对应参数为：

$\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}$