LLM系列：Transfomer

2017年《Attention Is All You Need》发表，Transfomer模型已成为NLP领域的首选模型。 Transfomer抛弃RNN的顺序结构，采用self-attention机制，模型可以并行化训练，而且能充分利用训练资料的全局信息。加入Transformer的Seq2Seq模型在NLP各个任务上表现显著提升。

Transformer由self-attention和FFN（Feed Forward Neural Network）组成，包括encoder、decoder各6层。Transformer核心机制是Self-Attention，本质是人类视觉注意力机制。

Attention注意力

输入序列$x^1,x^2...x^n$，经过embedding后变成向量$a^1,a^2...a^n$，对于输入向量$a^i$通过矩阵$W^q,W^k,W^v$线性变换分别得到Query向量$q^i$，Key向量$k^i$，Value向量$v^i$，简单矩阵表示即Q、K、V。

Attention最核心的公式： $Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$

Attention结构如下：

以下分步骤解析公示：

1 $XX^T$是什么

self-attention最开始形态： $softmax({XX^T})X$

假设$X=[x_1^T;x_2^T;x_3^T;]$，$x_i^T$是一个行向量，X是一个二维矩阵，如下图模拟$XX^T$

矩阵$XX^T$是一个方阵，里面保存每个向量与自己和其他向量的內积结果，即与其他向量相似程度。

两个向量的內积表征两个向量的夹角，表征一个向量在另一个向量的投影。投影的值越大，说明两个向量相关度越高。

2 $\sqrt{d_k}$与Softmax

假设Q、K的元素均值为0，方差为1，那么$A^T=Q^TK$中元素的均值为0，方差为d（具体推导省略）。当d变得很大时，A中元素的方差也会变得很大，Softmax(A)的分布会趋于陡峭（分布方差大，分布集中在绝对值大的区域）。A中每个元素除以$\sqrt{d_k}$后，方差变为1，这使得Softmax(A)分布陡峭程度与d解耦，使得训练过程中梯度保持稳定。

Softmax操作的作用是：归一化。Softmax后数字和为1。 $Softmax(z)=\frac{e^z}{\sum_{c = 1}^{C} e^{z_c}}$

3 加权求和

$softmax({XX^T})X$ 可以用如下举例表示，取一个行向量，与X的一个列向量相乘，得到一个新的行向量（维度与X相同），表示“早”词向量经过注意力机制加权求和。

4 多头模型

什么叫多头模型？

由于不同的 Attention 的权重侧重点不一样，所以将这个任务交给不同的 Attention 一起做，最后取综合结果会更好，有点像 CNN 中的 Kernel。在Transformer的论文中指出，将 Q、K、V 通过一个线性映射后，分成 h 份，对每份进行 Scaled Dot-Product Attention 效果更好， 再把这几个部分 Concat 起来，过一个线性层的效果更好，可以综合不同位置的不同表征子空间的信息。

多头attention模型的实现：

class MultiHead(nn.Module):
    def __init__(self, n_head, model_dim, drop_rate):
        super().__init__()
        self.head_dim = model_dim // n_head
        self.n_head = n_head
        self.model_dim = model_dim
        self.wq = nn.Linear(model_dim, n_head * self.head_dim)
        self.wk = nn.Linear(model_dim, n_head * self.head_dim)
        self.wv = nn.Linear(model_dim, n_head * self.head_dim)

        self.o_dense = nn.Linear(model_dim, model_dim)
        self.o_drop = nn.Dropout(drop_rate)
        self.layer_norm = nn.LayerNorm(model_dim)
        self.attention = None

    def forward(self, q, k, v, mask, training):
        # residual connect
        residual = q

        # linear projection
        key = self.wk(k)  # [n, step, num_heads * head_dim]
        value = self.wv(v)  # [n, step, num_heads * head_dim]
        query = self.wq(q)  # [n, step, num_heads * head_dim]

        # split by head
        query = self.split_heads(query)  # [n, n_head, q_step, h_dim]
        key = self.split_heads(key)
        value = self.split_heads(value)  # [n, h, step, h_dim]
        context = self.scaled_dot_product_attention(query, key, value, mask)  # [n, q_step, h*dv]
        o = self.o_dense(context)  # [n, step, dim]
        o = self.o_drop(o)

        o = self.layer_norm(residual + o)
        return o

    def split_heads(self, x):
        x = torch.reshape(x, (x.shape[0], x.shape[1], self.n_head, self.head_dim))
        return x.permute(0, 2, 1, 3)

    def scaled_dot_product_attention(self, q, k, v, mask=None):
        dk = torch.tensor(k.shape[-1]).type(torch.float)
        score = torch.matmul(q, k.permute(0, 1, 3, 2)) / (torch.sqrt(dk) + 1e-8)  # [n, n_head, step, step]
        if mask is not None:
            # change the value at masked position to negative infinity,
            # so the attention score at these positions after softmax will close to 0.
            score = score.masked_fill_(mask, -np.inf)
        self.attention = softmax(score, dim=-1)
        context = torch.matmul(self.attention, v)  # [n, num_head, step, head_dim]
        context = context.permute(0, 2, 1, 3)  # [n, step, num_head, head_dim]
        context = context.reshape((context.shape[0], context.shape[1], -1))
        return context  # [n, step, model_dim]

5 残差连接 (redidual connection)

残差网络通过加入 shortcut connections，变得更加容易被优化。包含一个 shortcut connection 的几层网络被称为一个残差块（residual block）。残差块分成两部分：直接映射部分和残差部分。 残差网络由残差块组成，一个残差块可以表示为：

残差网络有什么好处？
因为增加了 x 项，那么该网络求 x 的偏导的时候，多了一项常数 1，所以反向传播过程，梯度连乘，也不会造成梯度消失。

残差网络实现：

 def residual(sublayer_fn,x): 
      return sublayer_fn(x)+x 

5 Layer normalization

Normalization 有很多种，但是它们都有一个共同的目的，那就是把输入转化成均值为 0 方差为 1 的数据。

在把数据送入激活函数之前进行 Normalization（归一化），因为我们不希望输入数据落在激活函数的饱和区(平缓区梯度几乎为0, 造成梯度消失问题)。 如果没有BN，使用Sigmod激活函数会有严重的梯度消失问题。

BN的作用

Transform模型

1 编码器(encoder)

Encoder的输入：词编码矩阵$I\in{R^{n\times{l}\times{d}}}$，位置编码矩阵$P\in{R^{n\times{l}\times{d}}}$。n表示句子数，l表示一句话seq length，d表示词向量维度。

为什么需要位置编码矩阵？ 因为 Self-Attention 计算注意力分布的时候只能给出输出向量和输入向量之间的权重关系，但是不能给出词在一句话里的位置信息，所以需要在输入里引入位置编码矩阵。 位置编码向量生成方法有很多，常见是利用三角函数对位置进行编码。

输入矩阵I+P通过线性变换生成矩阵Q、K、V，实际编程中是将输入I+P直接赋值给Q、K、V。如果输入单词长度小于最大长度需要填充0，相应引入mask矩阵。

class EncoderLayer(nn.Module):

    def __init__(self, n_head, emb_dim, drop_rate):
        super().__init__()
        self.mh = MultiHead(n_head, emb_dim, drop_rate)
        self.ffn = PositionWiseFFN(emb_dim, drop_rate)

    def forward(self, xz, training, mask):
        # xz: [n, step, emb_dim]
        context = self.mh(xz, xz, xz, mask, training)  # [n, step, emb_dim]
        o = self.ffn(context)
        return o

2 解码器(decoder)

Decoder的输入：词编码矩阵$O\in{R^{n_1\times{l_1}\times{d}}}$，位置编码矩阵$P^O\in{R^{n_1\times{l_1}\times{d}}}$。因为 Decoder 的输入是具有时顺序关系的（即上一步的输出为当前步输入）所以还需要输入 Mask 矩阵以便计算注意力分布。

class DecoderLayer(nn.Module):
    def __init__(self, n_head, model_dim, drop_rate):
        super().__init__()
        self.mh = nn.ModuleList([MultiHead(n_head, model_dim, drop_rate) for _ in range(2)])
        self.ffn = PositionWiseFFN(model_dim, drop_rate)

    def forward(self, yz, xz, training, yz_look_ahead_mask, xz_pad_mask):
        dec_output = self.mh[0](yz, yz, yz, yz_look_ahead_mask, training)  # [n, step, model_dim]
        dec_output = self.mh[1](dec_output, xz, xz, xz_pad_mask, training)  # [n, step, model_dim]
        dec_output = self.ffn(dec_output)  # [n, step, model_dim]

        return dec_output

总结Decoder和Encoder中的self-attention层区别：

1 在Decoder中，Self-Attention层只允许关注到输出序列中早于当前位置之前的单词。具体做法是：在 Self-Attention 分数经过 Softmax 层之前，屏蔽当前位置之后的那些位置（将attention score设置成-inf）。
2 Decoder Attention层是使用前一层的输出来构造Query矩阵，而Key矩阵和 Value矩阵来自于Encoder最终的输出。

3 线性层和softmax

Decoder 最终的输出是一个向量，其中每个元素是浮点数。我们怎么把这个向量转换为单词呢？ 线性层和softmax完成词向量到单词的转换。

线性层就是一个普通的全连接神经网络，可以把解码器输出的向量，映射到一个更大的向量，这个向量称为 logits 向量：假设我们的模型有 10000 个英语单词（模型的输出词汇表），此 logits 向量便会有 10000 个数字，每个数表示一个单词的分数。然后，Softmax 层会把这些分数转换为概率（把所有的分数转换为正数，并且加起来等于 1）。然后选择最高概率的那个数字对应的词，就是这个时间步的输出单词。

4 损失函数

Transformer训练的时候，需要将解码器的输出和label一同送入损失函数，以获得loss，最终模型根据loss进行方向传播。

只要Transformer解码器预测了一组概率，我们就可以把这组概率和正确的输出概率做对比，然后使用反向传播来调整模型的权重，使得输出的概率分布更加接近整数输出。

以简单的用两组概率向量的的空间距离作为loss（向量相减，然后求平方和，再开方），当然也可以使用交叉熵(cross-entropy)]和KL 散度(Kullback–Leibler divergence)。

相关概念

greedy decoding和beam search

• greedy decoding：模型每个时间步只产生一个输出。模型从概率分布选择概率最大的次，并且丢弃其他词。
• beam search：每个时间步保留k个最高概率的输出词。举例k=1，第一个位置概率最高的两个词保留；第二个位置根据上个位置计算词的概率分布，再从中取出k=2个概率最高的词，如此重复。

如何通俗的理解beam search

参考：
zhuanlan.zhihu.com/p/137615798
zhuanlan.zhihu.com/p/476585349
mp.weixin.qq.com/s/sNeBiQKpH… mp.weixin.qq.com/s/ZllvtpGfk…