在当前的 NLP 领域，Transformer / BERT 已然成为基础应用，而 Self-Attention  则是两者的核心部分，下面尝试用 Q&A 和源码的形式深入 Self-Attention 的细节。

一 Q&A

1. Self-Attention 的核心是什么？

Self-Attention 的核心是用文本中的其它词来增强目标词的语义表示，从而更好的利用上下文的信息。

2. Self-Attention 的时间复杂度是怎么计算的？

Self-Attention 时间复杂度：O(*d)，这里n 是序列的长度，d 是 embedding 的维度，不考虑 batch 维。

Self-Attention 包括三个步骤：相似度计算，softmax 和加权平均。

它们分别的时间复杂度是：

• 相似度计算 可以看作大小为(n, d) 和 （d, n） 的两个矩阵相乘, (n, d) * (d, n) = O(*d), 得到一个(n, n)的矩阵
• softmax 就是直接计算了，时间复杂度为 O()
• 加权平均 可以看作大小为(n, n)和(n, d)的两个矩阵相乘 (n, n) * (n, d) = O(*d), 得到一个(n, d) 的矩阵。

因此，Self-Attention 的时间复杂度是 O(*d)

2.1  为什么在进行softmax之前需要对attention进行scaled（为什么除以 sdk的平方根）?

      假设 Q 和 K 的均值为0，方差为1。它们的矩阵乘积将有均值为0，方差为dk，因此使用dk的平方根被用于缩放，因为，Q 和 K 的矩阵乘积的均值本应该为 0，方差本应该为1，这样可以获得更平缓的softmax。当维度很大时，点积结果会很大，会导致softmax的梯度很小。为了减轻这个影响，对点积进行缩放。

2.3  计算attention的时候为何选择点乘而不是加法？两者计算复杂度和效果上有什么区别？

K和Q的点乘是为了得到一个attention score 矩阵，用来对V进行提纯。K和Q使用了不同的W_k, W_Q来计算，可以理解为是在不同空间上的投影。正因为 有了这种不同空间的投影，增加了表达能力，这样计算得到的attention score矩阵的泛化能力更高。

3. Transformer为何使用多头注意力机制?

In practice, the multi-headed attention are done with transposes and reshapes rather than actual separate tensors. —— 来自 google BERT 源代码注释

Tansformer 中的 Multi-Head Attention，简单来说就是多个 Self-Attention 的组合，它的作用类似于 CNN 中的多核。多头的实现不是循环的计算每个头，而是通过 transposes and reshapes，用矩阵乘法来完成的。多头可以使参数矩阵形成多个子空间，矩阵整体的size不变，只是改变了每个head对应的维度大小，这样做使矩阵对多方面信息进行学习，但是计算量和单个head差不多。

将原有的高维空间转化为多个低维空间并再最后进行拼接，形成同样维度的输出，借此丰富特性信息，降低了计算量

Transformer/BERT 中把 d ，也就是 hidden_size/embedding_size 这个维度做了 reshape 拆分，可以去看 Google 的 TF 源码或者上面的 pytorch 源码

 hidden_size (d) = num_attention_heads (m) * attention_head_size (a)，也即 d=m*a。

并将 num_attention_heads 维度 transpose 到前面，使得 Q 和 K 的维度都是 (m,n,a)，这里不考虑 batch 维度。

这样点积可以看作大小为 (m,n,a) 和 (m,a,n) 的两个张量相乘，得到一个 (m,n,n) 的矩阵，其实就相当于 m 个头，时间复杂度是:

=

张量乘法时间复杂度分析参见：矩阵、张量乘法的时间复杂度分析 [1]。

因此 Multi-Head Attention 时间复杂度就是 ，而实际上，张量乘法可以加速，因此实际复杂度会更低一些。

4. 不考虑多头的原因，self-attention中词向量不乘QKV参数矩阵，会怎么样？

对于 Attention 机制，都可以用统一的 query/key/value 模式去解释，而对于  self-attention，一般会说它的 q=k=v，这里的相等实际上是指它们来自同一个基础向量，而在实际计算时，它们是不一样的，因为这三者都是乘了 QKV 参数矩阵的。那如果不乘，每个词对应的 q,k,v 就是完全一样的。

在 self-attention 中，sequence 中的每个词都会和 sequence 中的每个词做点积去计算相似度，也包括这个词本身。

在相同量级的情况下，qi 与 ki 点积的值会是最大的（可以从“两数和相同的情况下，两数相等对应的积最大”类比过来）。

那在 softmax 后的加权平均中，该词本身所占的比重将会是最大的，使得其他词的比重很少，无法有效利用上下文信息来增强当前词的语义表示。

而乘以 QKV 参数矩阵，会使得每个词的 q,k,v 都不一样，能很大程度上减轻上述的影响。

当然，QKV 参数矩阵也使得多头，类似于 CNN 中的多核，去捕捉更丰富的特征/信息成为可能。

5. 在常规 attention 中，一般有 K=V，那 self-attention 可以嘛？

Self-Attention 实际只是 attention 中的一种特殊情况，因此 k=v 是没有问题的，也即 K，V 参数矩阵相同。

扩展到 Multi-Head Attention 中，乘以 Q、K 参数矩阵之后，其实就已经保证了多头之间的差异性了，在 q 和 k 点积 +softmax 得到相似度之后，从常规 attention 的角度，觉得再去乘以和 k 相等的 v 会更合理一些。

在 Transformer / BERT 中，完全独立的 QKV 参数矩阵，可以扩大模型的容量和表达能力。

但采用 Q，K=V 这样的参数模式，我认为也是没有问题的，也能减少模型的参数，又不影响多头的实现。

当然，上述想法并没有做过实验，为个人观点，仅供参考。

6. Q和K使用不同的权重矩阵生成，为何不能使用同一个值进行自身的点乘？

答：请求和键值初始为不同的权重是为了解决可能输入句长与输出句长不一致的问题。并且假如QK维度一致，如果不用Q，直接拿K和K点乘的话，你会发现attention score 矩阵是一个对称矩阵。因为是同样一个矩阵，都投影到了同样一个空间，所以泛化能力很差。

二 源码

在整个 Transformer / BERT 的代码中，(Multi-Head Scaled Dot-Product) Self-Attention 的部分是相对最复杂的，也是 Transformer / BERT 的精髓所在，这里给出 Pytorch 版本的实现 [2]，并对重要的代码加上了注释和维度说明。

话不多说，都在代码里，它主要有三个部分：

初始化：包括有几个头，每个头的大小，并初始化 Q K V 三个参数矩阵。

class SelfAttention(nn.Module):
    def __init__(self, config):
        super(SelfAttention, self).__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
        # 在Transformer/BERT中，这里的 all_head_size 就等于 config.hidden_size
        # 应该是一种简化，为了从embedding到最后输出维度都保持一致
        # 这样使得多个attention头合起来维度还是config.hidden_size
        # 而 attention_head_size 就是每个attention头的维度，要保证可以整除
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
 
        # 三个参数矩阵
        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)
 
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

transposes and reshapes：这个函数主要是把维度大小为 [batch_size * seq_length * hidden_size] 的 q,k,v 向量变换成 [batch_size * num_attention_heads * seq_length * attention_head_size]，便于后面做 Multi-Head Attention。

    def transpose_for_scores(self, x):
        """
        shape of x: batch_size * seq_length * hidden_size
        这个操作是把hidden_size分解为 self.num_attention_heads * self.attention_head_size
        然后再交换 seq_length 维度 和 num_attention_heads 维度
        为什么要做这一步：因为attention是要对query中的每个字和key中的每个字做点积，即是在 seq_length 维度上
        query和key的点积是 [seq_length * attention_head_size] * [attention_head_size * seq_length]=[seq_length * seq_length]
        """
        # 这里是一个维度拼接：(1,2)+(3,4) -> (1, 2, 3, 4)
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

前向计算： 乘以 QKV 参数矩阵 —> transposes and reshapes —> 做 scaled —> 加 attention mask —> Softmax —> 加权平均 —> 维度恢复。

 def forward(self, hidden_states, attention_mask):
        # shape of hidden_states and mixed_*_layer: batch_size * seq_length * hidden_size
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)
 
        # shape of *_layer: batch_size * num_attention_heads * seq_length * attention_head_size
        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)
 
        # Take the dot product between "query" and "key" to get the raw attention scores.
        # shape of attention_scores: batch_size * num_attention_heads * seq_length * seq_length
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
 
        # 这里就是做 Scaled，将方差统一到1，避免维度的影响
        attention_scores /= math.sqrt(self.attention_head_size)
 
        # shape of attention_mask: batch_size * 1 * 1 * seq_length. 它可以自动广播到和attention_scores一样的维度
        # 我们初始输入的attention_mask是：batch_size * seq_length，做了两次unsqueeze之后得到当前的attention_mask
        attention_scores = attention_scores + attention_mask
 
        # Normalize the attention scores to probabilities. Softmax 不改变维度
        # shape of attention_scores: batch_size * num_attention_heads * seq_length * seq_length
        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        attention_probs = self.dropout(attention_probs)
 
        # shape of value_layer: batch_size * num_attention_heads * seq_length * attention_head_size
        # shape of first context_layer: batch_size * num_attention_heads * seq_length * attention_head_size
        # shape of second context_layer: batch_size * seq_length * num_attention_heads * attention_head_size
        # context_layer 维度恢复到：batch_size * seq_length * hidden_size
        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
        return context_layer

Attention is all you need ! 希望这篇文章能让你对 Self-Attention 有更深的理解。

参考文献

[1] https://liwt31.github.io/2018/10/12/mul-complexity/

[2] https://github.com/hichenway/CodeShare/tree/master/bert_pytorch_source_code

超细节！从源代码剖析Self-Attention知识点-CSDN博客

21个Transformer面试题的简单回答