扒源码系列:GPT / GPT-2 中 proj 的作用

时间:2022-12-04 21:54:37

事情是这样的。前两天翻译了一篇文章图解GPT-2。在翻译的过程中为了防止自己出错,所以参考了一下其他人对于GPT的一些理解,然后就出错了,为了解决这个错误,导致我最后重新扒了一遍GPT-2的源码,在这里跟大家分享一下。


大家先来回顾一下GPT-2的结构,GPT-使用的是类似于transformer的decoder的组件。

就是输入先经过一个masked multi-head attention。在经过一个前馈神经网络。

扒源码系列:GPT / GPT-2 中 proj 的作用

假设我们现在有一个已经训练好的GPT-2模型。我们用它来做预测任务。就是输入一个句子让提炼。它继续往下生成,它是每次都往下生成一个token。

输入序列开始处理之前需要先对其进行表示,转化成模型可以理解的样子。所以要将输入的每个token使用token embedding矩阵和position embedding矩阵进行表示。之后数据会传入masked self-attention子层进行计算。

在masked 自注意力子层中,首先要进行一次一维卷积,其实我们就可以认定是它和一个大的权重矩阵进行计算,将计算结果拆分成三部分得到计算自注意力所需要的query、key和value三个矩阵。之后使用carry key value三个矩阵进行attention计算。将来计算获得的结果拼接起来,就是token的新表示,重新得到的表示中已经包含了token及其上下文的信息。

以GPT-2 small为例, 表示向量长度为768,最后经过Attention计算,拼接起来的向量表示长度也是768。但是看下图淡黄色的那一部分区域,数据出来之后,它和一个768×768的方阵进行了一个projection计算。

扒源码系列:GPT / GPT-2 中 proj 的作用

问题就出在了这里。这个projection计算有什么作用。

先来看一下我要翻译的那篇文章原文是怎么写的:

“But the vector isn't ready to be sent to the next sublayer just yet.  We need to first turn this Frankenstein’s-monster of hidden states into a homogenous representation.”

直译是“但是现在这个向量还不能直接传给下一个子层,我们需要首先把这个隐藏状态的弗兰肯斯坦怪物变成同类表示。”

我在这里产生了疑问。为什么要乘一个768×768的方阵,这样的话对它的张量的大小并没有任何改变,而一般情况下进行矩阵运算都是为了改变它的某一个维度。因为在这里存在疑问,所以我就去查阅了其他人的资料,之后在还算是比较权威的一个公众号的文章中看到了这样一个说法:

扒源码系列:GPT / GPT-2 中 proj 的作用

这篇文章中说因为向量的长度是不对的,所以这里需要进行一个projection投影改变向量长度。但是我翻译的文章中,他画的图是乘以一个方阵,并不会造成任何的向量维度改变。

所以这两个说法是完全矛盾的!!!

因此他们两个人中间必定有一个说法是错误的。所以我回去翻了一下GPT-2的源码。

从第1行代码到第18行代码都是做一个attention的计算。这里我们不用管它是怎么实现的。我们看一下第20行。第20行是把attention最后的计算结果合并起来。第21行代码是把合并起来的结果做了一个projection映射。之后的代码我们在这里就不用管了,就是对输出结果进行不改变维度的处理。。

# attention的计算
    query = self._split_heads(query, self.num_heads, self.head_dim)
    key = self._split_heads(key, self.num_heads, self.head_dim)
    value = self._split_heads(value, self.num_heads, self.head_dim)

    if layer_past is not None:
        past_key, past_value = layer_past
        key = torch.cat((past_key, key), dim=-2)
        value = torch.cat((past_value, value), dim=-2)

    if use_cache is True:
        present = (key, value)
    else:
        present = None

    if self.reorder_and_upcast_attn:
        attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
    else:
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
# 把attention最后的计算结果合并起来
    attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
    attn_output = self.c_proj(attn_output)          # projection映射
    attn_output = self.resid_dropout(attn_output)
# 对输出处理,不改变维度
    outputs = (attn_output, present)
    if output_attentions:
        outputs += (attn_weights,)

    return outputs  # a, present, (attentions)

接下来我们看一下这个 projection都做了什么。

Conv1D先看一下这个类,我们可以看到它注释中写了你就把它认为成立一个线性层就行了,区别就是它的权重和线性层相比进行了一个转置。所以说你就可以把它当成一个进行维度转化的就可以了。可以看到在这里Conv1D(self.embed_dim, self.embed_dim) 第一个参数是输出的维度,第二个参数是输入的维度。我们可以知道它的输入输出维度并没有改变,也就是说我翻译的那篇文章画的图是没有任何错误的。而我在公众号看的那篇文章说维度进行了改变,这个公众号的文章是有错误的。

self.c_proj = Conv1D(self.embed_dim, self.embed_dim)

class Conv1D(nn.Module):
    """
    你把1D卷积当成一个线性层就行了,它和线性层的主要区别在于二者的权重矩阵互为转置。
    GPT和GPT-2中都使用的这样的卷积。
    参数:
        nf 类型为int:输出特征的数量
        nx 类型为int:输入特征的数量
    """

    def __init__(self, nf, nx):
        super().__init__()
        self.nf = nf
        w = torch.empty(nx, nf)
        nn.init.normal_(w, std=0.02)
        self.weight = nn.Parameter(w)
        self.bias = nn.Parameter(torch.zeros(nf))

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf,)
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        x = x.view(size_out)
        return x

那这里为什么要进行一个projection呢?

We’ll let the model learn how to best map concatenated self-attention results into a vector that the feed-forward neural network can deal with. Here comes our second large weight matrix that projects the results of the attention heads into the output vector of the self-attention sublayer:

我们要让模型学习到 如何将自注意力的拼接结果更好地映射成前馈神经网络可以处理的向量 。因此这里要做一步映射。在这就用到了我们的第二大权重矩阵,它将自注意力的拼接结果映射为自注意力子层的输出向量。

再回到这张图。也就是说我们可以理解为在这里的这一部projection是为了做一个平滑作用。让模型学习出一个权重,将self-attention的结果映射到前馈神经网络更好处理的样子。

扒源码系列:GPT / GPT-2 中 proj 的作用