注意力机制（四）：多头注意力

专栏：神经网络复现目录

注意力机制

注意力机制（Attention Mechanism）是一种人工智能技术，它可以让神经网络在处理序列数据时，专注于关键信息的部分，同时忽略不重要的部分。在自然语言处理、计算机视觉、语音识别等领域，注意力机制已经得到了广泛的应用。

注意力机制的主要思想是，在对序列数据进行处理时，通过给不同位置的输入信号分配不同的权重，使得模型更加关注重要的输入。例如，在处理一句话时，注意力机制可以根据每个单词的重要性来调整模型对每个单词的注意力。这种技术可以提高模型的性能，尤其是在处理长序列数据时。

在深度学习模型中，注意力机制通常是通过添加额外的网络层实现的，这些层可以学习到如何计算权重，并将这些权重应用于输入信号。常见的注意力机制包括自注意力机制（self-attention）、多头注意力机制（multi-head attention）等。

总之，注意力机制是一种非常有用的技术，它可以帮助神经网络更好地处理序列数据，提高模型的性能。

文章目录

多头注意力

多头注意力（Multi-Head Attention）是注意力机制的一种扩展形式，可以在处理序列数据时更有效地提取信息。

在标准的注意力机制中，我们计算一个加权的上下文向量来表示输入序列的信息。而在多头注意力中，我们使用多组注意力权重，每组权重可以学习到不同的语义信息，并且每组权重都会产生一个上下文向量。最后，这些上下文向量会被拼接起来，再通过一个线性变换得到最终的输出。

多头注意力是Transformer模型中的一个重要组成部分，被广泛用于各种自然语言处理任务，如机器翻译、文本分类等。

数学逻辑

在实现多头注意力之前，让我们用数学语言将这个模型形式化地描述出来。 给定查询

     q 
    
   
     ∈ 
    
    
    
      R 
     
     
     
       d 
      
     
       q 
      
     
    
   
  
    q\in R^{d_q} 
   
  
q∈Rdq​、 键 
 
  
   
   
     k 
    
   
     ∈ 
    
    
    
      R 
     
     
     
       d 
      
     
       k 
      
     
    
   
  
    k\in R^{d_k} 
   
  
k∈Rdk​和值 
 
  
   
   
     v 
    
   
     ∈ 
    
    
    
      R 
     
     
     
       d 
      
     
       v 
      
     
    
   
  
    v\in R^{d_v} 
   
  
v∈Rdv​， 每个注意力头的计算方法为：
  
   
    
     
     
       h 
      
     
       i 
      
     
    
      = 
     
    
      f 
     
    
      ( 
     
     
     
       W 
      
     
       i 
      
      
      
        ( 
       
      
        q 
       
      
        ) 
       
      
     
    
      q 
     
    
      , 
     
     
     
       W 
      
     
       i 
      
      
      
        ( 
       
      
        k 
       
      
        ) 
       
      
     
    
      k 
     
    
      , 
     
     
     
       W 
      
     
       i 
      
      
      
        ( 
       
      
        v 
       
      
        ) 
       
      
     
    
      v 
     
    
      ) 
     
    
      ∈ 
     
     
     
       R 
      
      
      
        p 
       
      
        v 
       
      
     
    
   
     h_i=f(W_i^{(q)}q,W_i^{(k)}k,W_i^{(v)}v)\in R^{pv} 
    
   
 hi​=f(Wi(q)​q,Wi(k)​k,Wi(v)​v)∈Rpv

其中，可学习的参数包括

      W 
     
    
      i 
     
     
     
       ( 
      
     
       q 
      
     
       ) 
      
     
    
   
  
    W_i^{(q)} 
   
  
Wi(q)​、  
 
  
   
    
    
      W 
     
    
      i 
     
     
     
       ( 
      
     
       k 
      
     
       ) 
      
     
    
   
  
    W_i^{(k)} 
   
  
Wi(k)​和  
 
  
   
    
    
      W 
     
    
      i 
     
     
     
       ( 
      
     
       v 
      
     
       ) 
      
     
    
   
  
    W_i^{(v)} 
   
  
Wi(v)​， 以及代表注意力汇聚的函数 
 
  
   
   
     f 
    
   
  
    f 
   
  
f。  
 
  
   
   
     f 
    
   
  
    f 
   
  
f可以是加性注意力和缩放点积注意力。 多头注意力的输出需要经过另一个线性转换， 它对应着 
 
  
   
   
     h 
    
   
  
    h 
   
  
h个头连结后的结果，因此其可学习参数是  
 
  
   
    
    
      W 
     
    
      o 
     
    
   
  
    W_o 
   
  
Wo​：

实现

在实现过程中通常选择缩放点积注意力作为每一个注意力头。 为了避免计算代价和参数代价的大幅增长， 我们设定

      p 
     
    
      q 
     
    
   
     = 
    
    
    
      p 
     
    
      k 
     
    
   
     = 
    
    
    
      p 
     
    
      v 
     
    
   
     = 
    
    
    
      p 
     
    
      p 
     
    
   
     / 
    
   
     h 
    
   
  
    p_q=p_k=p_v=p_p/h 
   
  
pq​=pk​=pv​=pp​/h。 值得注意的是，如果将查询、键和值的线性变换的输出数量设置为 
 
  
   
    
    
      p 
     
    
      q 
     
    
   
     h 
    
   
     = 
    
    
    
      p 
     
    
      k 
     
    
   
     h 
    
   
     = 
    
    
    
      p 
     
    
      v 
     
    
   
     h 
    
   
     = 
    
    
    
      p 
     
    
      p 
     
    
   
  
    p_qh=p_kh=p_vh=p_p 
   
  
pq​h=pk​h=pv​h=pp​， 则可以并行计算 
 
  
   
   
     h 
    
   
  
    h 
   
  
h个头。 在下面的实现中，是通过参数 
 
  
   
    
    
      p 
     
    
      o 
     
    
   
  
    p_o 
   
  
po​num_hiddens指定的。

#@saveclassMultiHeadAttention(nn.Module):"""多头注意力"""def__init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False,**kwargs):super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)defforward(self, queries, keys, values, valid_lens):# queries，keys，values的形状:# (batch_size，查询或者“键－值”对的个数，num_hiddens)# valid_lens　的形状:# (batch_size，)或(batch_size，查询的个数)# 经过变换后，输出的queries，keys，values　的形状:# (batch_size*num_heads，查询或者“键－值”对的个数，# num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)if valid_lens isnotNone:# 在轴0，将第一项（标量或者矢量）复制num_heads次，# 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)# output的形状:(batch_size*num_heads，查询的个数，# num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)# output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)return self.W_o(output_concat)

为了能够使多个头并行计算， 上面的MultiHeadAttention类将使用下面定义的两个转置函数。 具体来说，transpose_output函数反转了transpose_qkv函数的操作。

#@savedeftranspose_qkv(X, num_heads):"""为了多注意力头的并行计算而变换形状"""# 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)# 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，# num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads,-1)# 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,# num_hiddens/num_heads)
    X = X.permute(0,2,1,3)# 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,# num_hiddens/num_heads)return X.reshape(-1, X.shape[2], X.shape[3])#@savedeftranspose_output(X, num_heads):"""逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0,2,1,3)return X.reshape(X.shape[0], X.shape[1],-1)

代码解释：

这段代码实现了多头注意力机制，其中 MultiHeadAttention 类实现了多头注意力的前向传播， transpose_qkv 函数将输入的 queries, keys, values 通过线性变换并按照 num_heads 进行分组，最终输出变换后的 queries, keys, values，在前向传播中使用这些变换后的 queries, keys, values 来计算注意力权重。在 transpose_qkv 函数的实现中，首先将 queries, keys, values 转换成形状为 (batch_size, queries/keys/values_num, num_hiddens) 的张量，然后根据 num_heads 将最后一维进行分组，变换成形状为 (batch_size, num_heads, queries/keys/values_num, num_hiddens/num_heads) 的张量，最后将第一维和第二维进行交换，输出形状为 (batch_size*num_heads, queries/keys/values_num, num_hiddens/num_heads) 的张量。transpose_output 函数实现了对 MultiHeadAttention 的输出进行逆转换的操作。

这么做的原因是因为多头注意力机制可以将输入张量进行 num_heads 个独立的注意力计算，将计算结果在最后一维拼接起来作为输出，这样可以提高模型的并行性，加快计算速度。同时，通过变换形状将 num_heads 独立处理，也可以增强模型对不同位置和特征的表征能力。

具体来说，这段代码实现的是一个MultiHeadAttention类，其中定义了一个forward方法。这个方法接收一个查询序列queries，一个键序列keys，一个值序列values和一个有效长度序列valid_lens作为输入，然后输出一个加权聚合的结果。

MultiHeadAttention类的初始化方法中，我们定义了几个线性层，以及注意力计算函数，然后用这些组件来定义一个多头注意力层。该层包括将输入queries、keys和values通过三个线性层进行变换，以便将它们的形状变为(batch_size * num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads)，其中num_heads表示注意力头的数量。然后，我们通过调用transpose_qkv函数对这些变换后的输入进行一次变换，以便在注意力计算函数中实现多头并行计算。最后，我们通过调用transpose_output函数将输出重构成(batch_size，查询的个数，num_hiddens)，并通过一个线性层对其进行变换，输出最终结果。

transpose_qkv函数将输入的queries、keys和values通过reshape和permute操作进行变换，以便多头并行计算。具体来说，它将输入变换为(batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)的形状，然后将第2和第3个轴进行交换。最后，它将输出变换为(batch_size * num_heads, 查询或者“键－值”对的个数, num_hiddens/num_heads)的形状。

transpose_output函数将多头并行计算得到的输出通过reshape和permute操作逆转回原来的形状，具体来说，它将输出变换为(batch_size，查询的个数，num_heads, num_hiddens/num_heads)的形状，然后将第2和第3个轴进行交换，最终将输出变换为(batch_size，查询的个数，num_hiddens)的形状。

这里似乎所有的单头都是同一些参数，这样不会导致每个单头的输出都是一样的吗？

这里的确有点难懂, 这里其实是把所有注意力头里面的参数拼起来, 变成了一个大的全连接层