NLP学习记录十：多头注意力

一、单头注意力

单头注意力的大致流程如下：

① 查询编码向量、键编码向量和值编码向量分别经过自己的全连接层（Wq、Wk、Wv）后得到查询Q、键K和值V；

② 查询Q和键K经过注意力评分函数（如：缩放点积运算）得到值权重矩阵；

③ 权重矩阵与值向量相乘，得到输出结果。

图1 单头注意力模型

二、多头注意力

2.1 使用多头注意力的意义

看了一些对多头注意力机制解释的视频，我自己的浅显理解是：在实践中，我们会希望查询Q能够从给定内容中尽可能多地匹配到与自己相关的语义信息，从而得到更准确的预测输出。而多头注意力将查询、键和值分成不同的子空间表示（representation subspaces）（有点类似于子特征？），使得匹配过程更加细化。

2.2 代码实现

也许直接看代码能更快地理解这个过程：

import torch
from torch import nn
from attentionScore import DotProductAttention

# 多头注意力模型
class MultiHeadAttention(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_headsself.attention = 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)# queries：(batch_size，查询的个数，query_size)# keys：(batch_size，“键－值”对的个数，key_size)# values：(batch_size，“键－值”对的个数，value_size)def forward(self, queries, keys, values, valid_lens):# queries，keys，values的形状：(batch_size，查询或者“键－值”对的个数，num_hiddens)queries = self.W_q(queries)keys = self.W_k(keys)values = self.W_v(values)# 经过变换后，输出的queries，keys，values的形状：(batch_size*num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads)queries = transpose_qkv(queries, self.num_heads)keys = transpose_qkv(keys, self.num_heads)values = transpose_qkv(values, self.num_heads)# valid_lens的形状：(batch_size，)或(batch_size，查询的个数)if valid_lens is not None:# 在轴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)

# 为了多注意力头的并行计算而变换形状
def transpose_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])

# 逆转transpose_qkv函数的操作
def transpose_output(X, num_heads):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)

可以发现，前面的处理流程和单头注意力的第①步是一样的，都是使用全连接层计算查询Q、键K、值V。但在进行点积运算之前，模型使用transpose_qkv函数对QKV进行了切割变换，下图可以帮助理解这个过程：