超细节的 Self-Attention 知识点 - 附源码解析

系统管理员 2022-12-17 15:00 203阅读 0赞

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

### Q&A ###

**1、 Self-Attention 的核心是什么？**

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

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

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

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

\[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-t9SG6ZU8-1603506048428)(./picture/self\_attention.png)\]

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

**相似度计算** 可以看作大小为 ( n , d ) (n,d) (n,d)和 ( d , n ) (d,n) (d,n)的两个矩阵相乘： ( n , d ) ∗ ( d , n ) = O ( n 2 ⋅ d ) (n,d)\*(d,n)=O(n^2 \\cdot d) (n,d)∗(d,n)=O(n2⋅d)，得到一个 ( n , n ) (n,n) (n,n)的矩阵

**softmax** 就是直接计算了，时间复杂度为 O ( n 2 ) O(n^2) O(n2)

**加权平均** 可以看作大小为 ( n , n ) (n,n) (n,n)和 ( n , d ) (n,d) (n,d)的两个矩阵相乘： ( n , n ) ∗ ( n , d ) = O ( n 2 ⋅ d ) (n,n)\*(n,d)=O(n^2 \\cdot d) (n,n)∗(n,d)=O(n2⋅d)，得到一个 ( n , d ) (n,d) (n,d)的矩阵

因此，Self-Attention的时间复杂度是 O ( n 2 ⋅ d ) O(n^2 \\cdot d) O(n2⋅d)。

这里再提一下 Tansformer 中的 Multi-Head Attention，多头 Attention，简单来说就是多个 Self-Attention 的组合，它的作用类似于CNN中的多核。

多头的实现不是循环的计算每个头，而是通过 transposes and reshapes，用矩阵乘法来完成的。

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

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)的矩阵，其实就相当于(n,a)和(a,n)的两个矩阵相乘，做了m次，时间复杂度是 O ( n 2 ⋅ m ⋅ a ) = O ( n 2 ⋅ d ) O(n^2 \\cdot m \\cdot a)=O(n^2 \\cdot d) O(n2⋅m⋅a)=O(n2⋅d) 。

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

因此Multi-Head Attention时间复杂度也是 O ( n 2 ⋅ d ) O(n^2 \\cdot d ) O(n2⋅d) ，复杂度相较单头并没有变化，主要还是transposes and reshapes 的操作，相当于把一个大矩阵相乘变成了多个小矩阵的相乘。

**3、不考虑多头的原因，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中的多核，去捕捉更丰富的特征/信息成为可能。

**4、在常规 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 这样的参数模式，我认为也是没有问题的，也能减少模型的参数，又不影响多头的实现。

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

### 源码 ###

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

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

**初始化**：包括有几个头，每个头的大小，并初始化QKV三个参数矩阵

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