注意力机制

1. 词向量

  词向量是自然语言处理中的一种核心技术，用于将词语或文本表示为向量。通过词向量，计算机可以以数值形式理解和处理语言。它们通过将词语映射到向量空间，将相似语义的词表示为距离接近的向量，从而实现更丰富的语言语义分析。
  词向量表示词语的好处：1. 语义相似的词在向量空间中靠近，比如“猫”和“狗”的向量接近，而与“车”较远。2. 词向量可以通过线性关系表达某些语义规则，如King - Man + Woman ≈ Queen。
  一段文本首先要经过分词（文本编码）得到词索引，再经过词嵌入得到词向量。这里暂时不介绍分词和词嵌入的方法。

2. 掩码softmax

2.1 填充掩码softmax

  在词向量技术还未兴起之前，填充掩码Softmax用于处理输入序列中由于长度不一致而造成的填充部分。在自然语言处理任务中，句子通常具有不同的长度，因此在将它们批量处理时，较短的句子会被填充为与最长句子相同的长度。这种填充通常使用0表示。在计算Softmax时，我们希望忽略这些填充的部分，以避免影响模型的学习和输出。

2.2 因果掩码softmax

  因果掩码主要用于自回归任务，如序列生成和机器翻译中的解码器部分。它的作用是确保模型在当前时间步只能关注到先前的时间步，避免看到未来的信息，从而符合序列生成的因果关系。

3. 注意力

  与RNN相比，注意力的计算复杂度更低，并且可以通过并行计算缩短计算时间。此外，注意力还可以动态地决定哪些输入部分对当前输出最重要，从而对重要信息给予更多的关注，提高了信息的利用效率。
  本节将主要介绍Transformer模型中使用的注意力，为理解Transformer模型打好基础。

3.1 注意力机制的$Q$，$K$，$V$

3.1.1 非自主性提示和自主性提示

  注意力过程中有两个重要的组件：自主性提示和非自主性提示。
  非自主性提示是基于环境的易见性和突出性。假如我们面前有五个物品： 一份报纸、一篇研究论文、一杯咖啡、一本笔记本和一本书。 所有纸制品都是黑白印刷的，但咖啡杯是红色的。 这个咖啡杯在这种视觉环境中是突出和显眼的， 不由自主地引起人们的注意。 所以我们会把视力最敏锐的地方放到咖啡上， 如下图所示。

  喝咖啡后，我们会变得兴奋并想读书， 所以转过头，重新聚焦眼睛，然后看看书， 就像下图中描述那样。 与上图中由于突出性导致的选择不同， 此时选择书是受到了认知和意识的控制， 因此注意力在基于自主性提示去辅助选择时将更为谨慎。人的主观意愿推动，选择的力量也就更强大。

3.1.2 如何用自主性提示和非自主性提示理解$Q$，$K$，$V$

  上面提到的非自主性提示就像注意力机制中的$K$，而自主性提示就像注意力机制中的$Q$，人接受到的感官输入（这里是视觉）就像注意力机制中的$V$。
  不妨想象一下这样的一个场景：你和朋友约定在一个地方见面，然后一起去进行今日的计划。朋友可能告诉了你今日他穿的衣服和等待的位置，这些是自主性提示($Q$)，你会集中注意力搜索符合这些特征的区域和对象上。此外，朋友看到了你，可能进行挥手或者大声喊你，这些动作和场景内其他引人注目的事物是非自主性提示($K$)，你的注意力被吸引，尽管可能不是你的朋友，但是这种自然而然的反应可以帮你缩小寻找的范围。自主性提示($Q$)和非自主性提示($K$)结合，并对感官输入($V$)作出反应，直至找到你的朋友。

3.2 注意力机制的计算过程

  下图中给出了一个查询在注意力机制中根据键值对计算输出的过程。查询$q_i$是一个形状为$query\_size\times 1$的向量，键$k_j$是一个形状为$key\_size\times 1$的向量，值$v_j$是一个形状为$value\_size\times 1$的向量。在注意力机制中键、值的数量必须一样。
  $\alpha$是注意力评分函数，查询$q_i$和键$k_j$的注意力评分函数值为$a(q_i,k_j)$，查询$q_i$和键$k_j$的注意力权重为$w_a(q_i,k_j)=\frac{\alpha (q_i,k_j)}{\sum _j\alpha (q_i,k_j)}$，最终的输出为$o=\sum _j{w}_a(q_i,k_j)v_j$。

3.3 加性注意力

  加性注意力的评分函数为$e_{ij}=w^{T}tanh(W_q{q}_i+W_k{k}_j)$，其中查询$q_i$为$query\_size\times 1$的向量，$W_q$的形状为$h\times query\_size$，键$k_j$为$key\_size\times 1$的向量，$W_k$的形状为$h\times key\_size$，$w$为$h\times 1$的向量。
  加性注意力层的PyTorch实现如下。

import torch
import torch.nn as nn

class AdditiveAttention(nn.Module):
    def __init__(self, query_size, key_size, num_hidden, dropout):
        super().__init__()
        self.wq = nn.Linear(query_size, num_hidden, bias=False)
        self.wk = nn.Linear(key_size, num_hidden, bias=False)
        self.wv = nn.Linear(num_hidden, 1, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, queries, keys, values, mask=None):
    	assert keys.shape[1] == values.shape[1], "The sequence length of keys and values must be equal"
        queries = self.wq(queries)
        keys = self.wk(keys)
        scores = self.wv(nn.Tanh()(queries.unsqueeze(2) + keys.unsqueeze(1))).squeeze(-1)
        if mask is not None:
            scores.masked_fill_(mask, -float("inf"))
        scores = nn.Softmax(-1)(scores)
        return torch.bmm(self.dropout(scores), values)

3.4 缩放点积注意力

  缩放点积注意力$Q$的维度query_size必须和$K$的维度key_size相等，令$d_k=query\_size=key\_size$。缩放点积注意力的评分函数为$e_{ij}=\frac{q_i^T{k}_j}{\sqrt{d_k}}$，其中查询$q_i$为$d_k\times 1$的向量，键$k_j$为$d_k\times 1$的向量。v
  缩放点积注意力的矩阵表示为$A=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$，其中查询$Q$的形状为$num\_query\times d_k$，键$K$的形状为$seq\_len\times d_k$，值$V$的形状为$seq\_len\times value\_size$。

  缩放点积注意力的PyTorch实现如下。

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, queries, keys, values, mask=None):
    	assert keys.shape[1] == values.shape[1], "The sequence length of keys and values must be equal"
    	assert queries.shape[2] == keys.shape[2], "query_size must equal key_size"
        scores = torch.bmm(queries, keys.transpose(1, 2)) / queries.shape[-1]
        if mask is not None:
        	# mask matrix's shape: (1, num_query, seq_len) or (1, queries.shape[1], keys.shape[1])
            scores.masked_fill_(mask, -float("inf"))
        scores = nn.Softmax(-1)(scores)
        return torch.bmm(self.dropout(scores), values)

3.5 多头注意力

  多头注意力有$h$个注意力头，这里的注意力头使用缩放点积注意力。查询$Q$的形状为$num\_query\times query\_size$，键$K$的形状为$seq\_len\times key\_size$，值$V$的形状为$seq\_len\times value\_size$。$W_q$的形状为$query\_size\times d$，$W_k$的形状为$key\_size\times d$，$W_v$的形状为$value\_size\times d$($d$是$h$的倍数)，$W_o$的形状为$d\times d$。
  多头注意力的计算过程：
  1.$Q=Q{W}_q$，$K=K{W}_k$，$V=V{W}_v$。
  2.经过上一步$Q$、$K$、$V$形状的第二维都为$d$，$Q$的形状重塑成$num\_query\times h\times \frac{d}{h}$
，$K$的形状重塑成$seq\_len\times h\times \frac{d}{h}$，$V$的形状重塑成$seq\_len\times h\times \frac{d}{h}$，重塑后的$Q$、$K$、$V$第二维和第一维互换。
  3.经历重塑和互换后的$Q$、$K$、$V$进行缩放点积注意力计算，得到输出O，形状为$h\times num\_query\times \frac{d}{h}$。
  4.$O$的第一维和第二维互换，然后形状重塑成$num\_query\times d$。
  5.最终输出为$O{W}_o$。

  多头注意力的PyTorch实现如下。

class MultiHeadAttention(nn.Module):
    def __init__(self, query_size, key_size, value_size, num_hidden, num_head, dropout):
        super().__init__()
        assert num_hidden % num_head == 0, "num_hidden must be divisible by num_head"
        self.num_head = num_head
        self.wq = nn.Linear(query_size, num_hidden, bias=False)
        self.wk = nn.Linear(key_size, num_hidden, bias=False)
        self.wv = nn.Linear(value_size, num_hidden, bias=False)
        self.attention = ScaledDotProductAttention(dropout)
        self.wo = nn.Linear(num_hidden, num_hidden)
        
    def transpose_qkv(self, x, num_head):
        x = x.reshape(x.shape[0], x.shape[1], num_head, -1)
        x = x.permute(0, 2, 1, 3)
        return x.reshape(-1, x.shape[2], x.shape[3])
    
    def transpose_output(self, x, num_head):
        x = x.reshape(-1, num_head, x.shape[1], x.shape[2])
        x = x.permute(0, 2, 1, 3)
        return  x.reshape(x.shape[0], x.shape[1], -1)
    
    def forward(self, queries, keys, values, mask=None):
        queries = self.transpose_qkv(self.wq(queries), self.num_head)
        keys = self.transpose_qkv(self.wk(keys), self.num_head)
        values = self.transpose_qkv(self.wv(values), self.num_head)
        # mask matrix's shape: (1, num_query, seq_len) or (1, queries.shape[1], keys.shape[1])
        scores = self.attention(queries, keys, values, mask)
        scores_concat = self.transpose_output(scores, self.num_head)
        return self.wo(scores_concat)

参考

Ashish Vaswani, Noam Shazeer, Niki Parmar and et al. Attention Is All You Need.
动手学深度学习 10.注意力机制