MultiHead-Attention和Masked-Attention的机制和原理

一、本文说明

看本文前，需要先彻底搞懂Self-Attention。推荐看我的另一篇博文层层剖析，让你彻底搞懂Self-Attention、MultiHead-Attention和Masked-Attention的机制和原理。本篇文章内容在上面这篇也有，可以一起看。

二. MultiHead Attention

2.1 MultiHead Attention理论讲解

在Transformer中使用的是MultiHead Attention，其实这玩意和Self Attention区别并不是很大。先明确以下几点，然后再开始讲解：

1. MultiHead的head不管有几个，参数量都是一样的。并不是head多，参数就多。
2. 当MultiHead的head为1时，并不等价于Self Attetnion，MultiHead Attention和Self Attention是不一样的东西
3. MultiHead Attention使用的也是Self Attention的公式
4. MultiHead除了 $W^q,W^k,W^v$三个矩阵外，还要多额外定义一个 $W^o$。

好了，知道上面几点，我们就可以开始讲解MultiHeadAttention了。

MultiHead Attention大部分逻辑和Self Attention是一致的，是从求出Q,K,V后开始改变的，所以我们就从这里开始讲解。

现在我们求出了Q, K, V矩阵，对于Self-Attention，我们已经可以带入公式了，用图像表示则为：

为了简单起见，该图忽略了Softmax和 $d_k$ 的计算

而MultiHead Attention在带入公式前做了一件事情，就是拆，它按照“词向量维度”这个方向，将Q,K,V拆成了多个头，如图所示：

这里我的head数为4。既然拆成了多个head，那么之后的计算，也是各自的head进行计算，如图所示：

但这样拆开来计算的Attention使用Concat进行合并效果并不太好，所以最后需要再采用一个额外的$W^o$矩阵，对Attention再进行一次线性变换，如图所示：

到这里也能看出来，head数并不是越多越好。而为什么要用MultiHead Attention，Transformer给出的解释为：Multi-head attention允许模型共同关注来自不同位置的不同表示子空间的信息。反正就是用了比不用好。

2.2. Pytorch实现MultiHead Attention

该代码参考项目annotated-transformer。

首先定义一个通用的Attention函数：

def attention(query, key, value):
    """
    计算Attention的结果。
    这里其实传入的是Q,K,V，而Q,K,V的计算是放在模型中的，请参考后续的MultiHeadedAttention类。

    这里的Q,K,V有两种Shape，如果是Self-Attention，Shape为(batch, 词数, d_model)，
                           例如(1, 7, 128)，即batch_size为1，一句7个单词，每个单词128维

                           但如果是Multi-Head Attention，则Shape为(batch, head数, 词数，d_model/head数)，
                           例如(1, 8, 7, 16)，即Batch_size为1，8个head，一句7个单词，128/8=16。
                           这样其实也能看出来，所谓的MultiHead其实就是将128拆开了。

                           在Transformer中，由于使用的是MultiHead Attention，所以Q,K,V的Shape只会是第二种。

    """

    # 获取d_model的值。之所以这样可以获取，是因为query和输入的shape相同，
    # 若为Self-Attention，则最后一维都是词向量的维度，也就是d_model的值。
    # 若为MultiHead Attention，则最后一维是 d_model / h，h为head数
    d_k = query.size(-1)
    # 执行QK^T / √d_k
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 执行公式中的Softmax
    # 这里的p_attn是一个方阵
    # 若是Self Attention，则shape为(batch, 词数, 次数)，例如(1, 7, 7)
    # 若是MultiHead Attention，则shape为(batch, head数, 词数，词数)
    p_attn = scores.softmax(dim=-1)

    # 最后再乘以 V。
    # 对于Self Attention来说，结果Shape为(batch, 词数, d_model)，这也就是最终的结果了。
    # 但对于MultiHead Attention来说，结果Shape为(batch, head数, 词数，d_model/head数)
    # 而这不是最终结果，后续还要将head合并，变为(batch, 词数, d_model)。不过这是MultiHeadAttention
    # 该做的事情。
    return torch.matmul(p_attn, value)


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model):
        """
        h: head的数量
        """
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        # 定义W^q, W^k, W^v和W^o矩阵。
        # 如果你不知道为什么用nn.Linear定义矩阵，可以参考该文章：
        # https://blog.csdn.net/zhaohongfei_358/article/details/122797190
        self.linears = [
            nn.Linear(d_model, d_model),
            nn.Linear(d_model, d_model),
            nn.Linear(d_model, d_model),
            nn.Linear(d_model, d_model),
        ]

    def forward(self, x):
        # 获取Batch Size
        nbatches = x.size(0)

        """
        1. 求出Q, K, V，这里是求MultiHead的Q,K,V，所以Shape为(batch, head数, 词数，d_model/head数)
            1.1 首先，通过定义的W^q,W^k,W^v求出SelfAttention的Q,K,V，此时Q,K,V的Shape为(batch, 词数, d_model)
                对应代码为 `linear(x)`
            1.2 分成多头，即将Shape由(batch, 词数, d_model)变为(batch, 词数, head数，d_model/head数)。
                对应代码为 `view(nbatches, -1, self.h, self.d_k)`
            1.3 最终交换“词数”和“head数”这两个维度，将head数放在前面，最终shape变为(batch, head数, 词数，d_model/head数)。
                对应代码为 `transpose(1, 2)`
        """
        query, key, value = [
            linear(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for linear, x in zip(self.linears, (x, x, x))
        ]

        """
        2. 求出Q,K,V后，通过attention函数计算出Attention结果，
           这里x的shape为(batch, head数, 词数，d_model/head数)
           self.attn的shape为(batch, head数, 词数，词数)
        """
        x = attention(
            query, key, value
        )

        """
        3. 将多个head再合并起来，即将x的shape由(batch, head数, 词数，d_model/head数)
           再变为 (batch, 词数，d_model)
           3.1 首先，交换“head数”和“词数”，这两个维度，结果为(batch, 词数, head数, d_model/head数)
               对应代码为：`x.transpose(1, 2).contiguous()`
           3.2 然后将“head数”和“d_model/head数”这两个维度合并，结果为(batch, 词数，d_model)
        """
        x = (
            x.transpose(1, 2)
                .contiguous()
                .view(nbatches, -1, self.h * self.d_k)
        )

        # 最终通过W^o矩阵再执行一次线性变换，得到最终结果。
        return self.linears[-1](x)

接下来尝试使用一下：

# 定义8个head，词向量维度为512
model = MultiHeadedAttention(8, 512)
# 传入一个batch_size为2， 7个单词，每个单词为512维度
x = torch.rand(2, 7, 512)
# 输出Attention后的结果
print(model(x).size())

输出为：

torch.Size([2, 7, 512])

三. Masked Attention

3.1 为什么要使用Mask掩码

在Transformer中的Decoder中有一个Masked MultiHead Attention。本节来对其进行一个详细的讲解。

首先我们来复习一下Attention的公式：

$$\begin{array}{rl}{O}_{n\times d_v}=\text{ Attention }(Q_{n\times d_k},K_{n\times d_k},V_{n\times d_v})& =\mathrm{softmax}\left(\frac{Q_{n\times d_k}{K}_{d_k\times n}^T}{\sqrt{d_k}}\right)V_{n\times d_v}\\ & \\ & =A_{n\times n}^{\prime }{V}_{n\times d_v}\end{array}$$

其中：

$$O_{n\times d_v}=\left[\begin{array}{c}{o}_1\\ o_2\\ ⋮\\ o_n\end{array}\right],A_{n\times n}^{\prime }=\left[\begin{array}{cccc}{\alpha }_{1,1}^{\prime }& {\alpha }_{2,1}^{\prime }& \cdots & {\alpha }_{n,1}^{\prime }\\ {\alpha }_{1,2}^{\prime }& {\alpha }_{2,2}^{\prime }& \cdots & {\alpha }_{n,2}^{\prime }\\ ⋮& ⋮& & ⋮\\ {\alpha }_{1,n}^{\prime }& {\alpha }_{2,n}^{\prime }& \cdots & {\alpha }_{n,n}^{\prime }\end{array}\right],V_{n\times d_v}=\left[\begin{array}{c}{v}_1\\ v_2\\ ⋮\\ v_n\end{array}\right]$$

假设 $(v_1,v_2,...v_n)$ 对应着 $(机,器,学,习,真,好,玩)$。那么 $(o_1,o_2,...,o_n)$ 就对应着 $({机}^{\prime },{器}^{\prime },{学}^{\prime },{习}^{\prime },{真}^{\prime },{好}^{\prime },{玩}^{\prime })$。 其中 ${机}^{\prime }$ 包含着 $v_1$ 到 $v_n$ 的所有注意力信息。而计算 ${机}^{\prime }$ 时的 $(机,器,...)$ 这些字的权重就是 $A^{\prime }$ 的第一行的 $({\alpha }_{1,1}^{\prime },{\alpha }_{2,1}^{\prime },...)$。

如果上面的回忆起来了，那么接下来看一下Transformer的用法，假设我们是要用Transformer翻译“Machine learning is fun”这句话。

首先，我们会将“Machine learning is fun” 送给Encoder，输出一个名叫Memory的Tensor，如图所示：

之后我们会将该Memory作为Decoder的一个输入，使用Decoder预测。Decoder并不是一下子就能把“机器学习真好玩”说出来，而是一个词一个词说（或一个字一个字，这取决于你的分词方式），如图所示：

紧接着，我们会再次调用Decoder，这次是传入“<bos> 机”：

依次类推，直到最后输出<eos>结束：

当Transformer输出<eos>时，预测就结束了。

到这里我们就会发现，对于Decoder来说是一个字一个字预测的，所以假设我们Decoder的输入是“机器学习”时，“习”字只能看到前面的“机器学”三个字，所以此时对于“习”字只有“机器学习”四个字的注意力信息。

但是，例如最后一步传的是“<bos>机器学习真好玩”，还是不能让“习”字看到后面“真好玩”三个字，所以要使用mask将其盖住，这又是为什么呢？原因是：如果让“习”看到了后面的字，那么“习”字的编码就会发生变化。

我们不妨来分析一下：

一开始我们只传入了“机”（忽略bos），此时使用attention机制，将“机”字编码为了 $[0.13,0.73,...]$

第二次，我们传入了“机器”，此时使用attention机制，如果我们不将“器”字盖住的话，那“机”字的编码就会发生变化，它就不再是是$[0.13,0.73,...]$了，也许就变成了$[0.95,0.81,...]$。

这就会导致第一次“机”字的编码是$[0.13,0.73,...]$，第二次却变成了$[0.95,0.81,...]$，这样就可能会让网络有问题。所以我们为了不让“机”字的编码产生变化，所以我们要使用mask，掩盖住“机”字后面的字，也就是即使他能attention后面的字，也不让他attention。

许多文章的解释是Mask是为了防止Transformer在训练时泄露后面的它不应该看到的信息，为什么会这么说呢？这个在3.4进行详细的解释。

3.2 如何进行mask掩码

要进行掩码，只需要对scores动手就行了，也就是 $A_{n\times n}^{\prime }$ 。直接上例子：

第一次，我们只有 $v_1$ 变量，所以是：

$$\left[\begin{array}{c}{o}_1\end{array}\right]=\left[\begin{array}{c}{\alpha }_{1,1}^{\prime }\end{array}\right]\cdot \left[\begin{array}{c}{v}_1\end{array}\right]$$

第二次，我们有 $v_1,v_2$ 两个变量：

$$\left[\begin{array}{c}{o}_1\\ o_2\end{array}\right]=\left[\begin{array}{cc}{\alpha }_{1,1}^{\prime }& {\alpha }_{2,1}^{\prime }\\ {\alpha }_{1,2}^{\prime }& {\alpha }_{2,2}^{\prime }\end{array}\right]\left[\begin{array}{c}{v}_1\\ v_2\end{array}\right]$$

此时如果我们不对 $A_{2\times 2}^{\prime }$ 进行掩码的话，$o_1$的值就会发生变化（第一次是 ${\alpha }_{1,1}^{\prime }{v}_1$，第二次却变成了${\alpha }_{1,1}^{\prime }{v}_1+{\alpha }_{2,1}^{\prime }{v}_2$）。那这样看，我们只需要将 ${\alpha }_{2,1}^{\prime }$ 盖住即可，这样就能保证两次的 $o_1$ 一致了。

所以第二次实际就为：

$$\left[\begin{array}{c}{o}_1\\ o_2\end{array}\right]=\left[\begin{array}{cc}{\alpha }_{1,1}^{\prime }& 0\\ {\alpha }_{1,2}^{\prime }& {\alpha }_{2,2}^{\prime }\end{array}\right]\left[\begin{array}{c}{v}_1\\ v_2\end{array}\right]$$

依次类推，如果我们执行到第$n$次时，就应该变成：

$$\left[\begin{array}{c}{o}_1\\ o_2\\ ⋮\\ o_n\end{array}\right]=\left[\begin{array}{cccc}{\alpha }_{1,1}^{\prime }& 0& \cdots & 0\\ {\alpha }_{1,2}^{\prime }& {\alpha }_{2,2}^{\prime }& \cdots & 0\\ ⋮& ⋮& & ⋮\\ {\alpha }_{1,n}^{\prime }& {\alpha }_{2,n}^{\prime }& \cdots & {\alpha }_{n,n}^{\prime }\end{array}\right]\left[\begin{array}{c}{v}_1\\ v_2\\ ⋮\\ v_n\end{array}\right]$$

3.3 为什么是负无穷而不是0

按照上面的说法，mask掩码是0，但为什么源码中的掩码是 $-1e9$ （负无穷）。Attention部分源码如下：

if mask is not None:
    scores = scores.masked_fill(mask == 0, -1e9)

p_attn = scores.softmax(dim=-1)

你仔细看，我们上面说的$A_{n\times n}^{\prime }$ 是什么，是softmax之后的。而源码中呢， 源码是在softmax之前进行掩码，所以才是负无穷，因为将负无穷softmax后就会变成0了。

3.4. 训练时的掩码

通常我们在网上看Masked Attention相关的文章时，会说mask的目的是为了防止网络看到不该看到的内容。本节主要来解释一下这句话。

首先，我们需要了解一下Transformer的训练过程。

在Transformer推理时，我们是一个词一个词的输出，但在训练时这样做效率太低了，所以我们会将target一次性给到Transformer（当然，你也可以按照推理过程做），如图所示：

从图上可以看出，Transformer的训练过程和推理过程主要有以下几点异同：

1. 源输入src相同：对于Transformer的inputs部分(src参数)一样，都是要被翻译的句子。
2. 目标输入tgt不同：在Transformer推理时，tgt是从<bos>开始，然后每次加入上一次的输出（第二次输入为<bos> 我）。但在训练时是一次将“完整”的结果给到Transformer，这样其实和一个一个给结果上一致。这里还有一个细节，就是tgt比src少了一位，src是7个token，而tgt是6个token。这是因为我们在最后一次推理时，只会传入前n-1个token。举个例子：假设我们要预测<bos> 我 爱 你 <eos>（这里忽略pad），我们最后一次的输入tgt是<bos> 我 爱 你（没有<eos>），因此我们的输入tgt一定不会出现目标的最后一个token，所以一般tgt处理时会将目标句子删掉最后一个token。
3. 输出数量变多：在训练时，transformer会一次输出多个概率分布。例如上图，我就的等价于是tgt为<bos>时的输出，爱就等价于tgt为<bos> 我时的输出，依次类推。当然在训练时，得到输出概率分布后就可以计算loss了，并不需要将概率分布再转成对应的文字。注意这里也有个细节，我们的输出数量是6，对应到token就是我 爱 你 <eos> <pad> <pad>，这里少的是<bos>，因为<bos>不需要预测。计算loss时，我们也是要和的这几个token进行计算，所以我们的label不包含<bos>。代码中通常命名为tgt_y。

其实总结一下就一句话：Transformer推理时是一个一个词预测，而训练时会把所有的结果一次性给到Transformer，但效果等同于一个一个词给，而之所以可以达到该效果，就是因为对tgt进行了掩码，防止其看到后面的信息，也就是不要让前面的字具备后面字的上下文信息。

可能看了这句总结还是很难理解，所以我们接下来来做个实验，我们的实验内容为：首先模拟Transformer的推理过程，然后再模拟Transformer的训练过程，看看训练时一次性给到所有的tgt和推理时一个一个给的结果是否一致。

这里我们要用到Pytorch中的nn.Transformer，用法可参考这篇文章。

首先我们来定义模型：

# 词典数为10， 词向量维度为8
embedding = nn.Embedding(10, 8)
# 定义Transformer，注意一定要改成eval模型，否则每次输出结果不一样
transformer = nn.Transformer(d_model=8, batch_first=True).eval()

接下来定义我们的src和tgt：

# Encoder的输入
src = torch.LongTensor([[0, 1, 2, 3, 4]])
# Decoder的输入
tgt = torch.LongTensor([[4, 3, 2, 1, 0]])

然后我们将[4]送给Transformer进行预测，模拟推理时的第一步：

transformer(embedding(src), embedding(tgt[:, :1]),
            # 这个就是用来生成阶梯式的mask的
            tgt_mask=nn.Transformer.generate_square_subsequent_mask(1))

tensor([[[ 1.4053, -0.4680,  0.8110,  0.1218,  0.9668, -1.4539, -1.4427,
           0.0598]]], grad_fn=<NativeLayerNormBackward0>)

然后我们将[4, 3]送给Transformer，模拟推理时的第二步：

transformer(embedding(src), embedding(tgt[:, :2]), tgt_mask=nn.Transformer.generate_square_subsequent_mask(2))

tensor([[[ 1.4053, -0.4680,  0.8110,  0.1218,  0.9668, -1.4539, -1.4427,
           0.0598],
         [ 1.2726, -0.3516,  0.6584,  0.3297,  1.1161, -1.4204, -1.5652,
          -0.0396]]], grad_fn=<NativeLayerNormBackward0>)

这个时候你有没有发现，输出的第一个向量和上面那个一模一样。

最后我们再将tgt一次性送给transformer，模拟训练过程：

transformer(embedding(src), embedding(tgt), tgt_mask=nn.Transformer.generate_square_subsequent_mask(5))

tensor([[[ 1.4053, -0.4680,  0.8110,  0.1218,  0.9668, -1.4539, -1.4427,
           0.0598],
         [ 1.2726, -0.3516,  0.6584,  0.3297,  1.1161, -1.4204, -1.5652,
          -0.0396],
         [ 1.4799, -0.3575,  0.8310,  0.1642,  0.8811, -1.3140, -1.5643,
          -0.1204],
         [ 1.4359, -0.6524,  0.8377,  0.1742,  1.0521, -1.3222, -1.3799,
          -0.1454],
         [ 1.3465, -0.3771,  0.9107,  0.1636,  0.8627, -1.5061, -1.4732,
           0.0729]]], grad_fn=<NativeLayerNormBackward0>)

看到没，前两个tensor和模拟推理时的输出结果一模一样。所以使用mask时，我们可以保证前面的词不会具备后面词的信息，这样就可以保证Transformer的输出不会因为传入词的多少而改变，从而我们就可以做到在训练时一次将tgt全部给到Transformer，却不会出现问题。这也就是人们常说的，防止网络训练时看到不该看到的内容。

可以尝试思考下为什么输出不会变，原因其实就是因为神经网络的本质就是不断的进行矩阵相乘，例如：$X{W}_1{W}_2{W}_3\cdots W_n\to O$，$X$ 为输入， $O$ 为输出。在这之中，$X$ 的第二个行向量本身就不会让你的第一个行向量的结果改变。在Transformer中多个行向量会互相影响是因为Attention机制，因为里面存在有$X$自身的运算，类似于 $X\cdot X$，但我们通过mask可以保证 $X\cdot Mask\_X$ 的第二个行向量不要影响到第一个行向量。这里就不展开讲解了，可以尝试用纸笔算一下。

参考资料

Pytorch中 nn.Transformer的使用详解与Transformer的黑盒讲解： https://blog.csdn.net/zhaohongfei_358/article/details/126019181文章来源地址https://www.toymoban.com/news/detail-817470.html

到了这里，关于MultiHead-Attention和Masked-Attention的机制和原理的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！