李沐动手学深度学习V2-多头注意力机制和代码实现

一. 多头自注意力

1. 介绍

在实践中，当给定相同的查询、键和值的集合时， 希望模型可以基于相同的注意力机制学习到不同的行为， 然后将不同的行为作为知识组合起来， 捕获序列内各种范围的依赖关系 （例如短距离依赖和长距离依赖关系）。 因此允许注意力机制组合使用查询、键和值的不同子空间表示（representation subspaces）可能是有益的。
为此与其只使用单独一个注意力汇聚， 我们可以用独立学习得到的 ℎ 组不同的线性投影（linear projections）来变换查询、键和值。 然后，这 ℎ 组变换后的查询、键和值将并行地送到注意力汇聚中。 最后，将这 ℎ 个注意力汇聚的输出拼接在一起， 并且通过另一个可以学习的线性投影进行变换， 以产生最终输出，这种设计被称为多头注意力（multihead attention）。 对于 ℎ 个注意力汇聚输出，每一个注意力汇聚都被称作一个头（head）。 下图展示了使用全连接层来实现可学习的线性变换的多头注意力。

2. 模型

在实现多头注意力之前，用数学语言将这个模型形式化地描述出来。给定查询$\mathbf{q}\in {\mathbb{R}}^{d_q}$、键$\mathbf{k}\in {\mathbb{R}}^{d_k}$和值$\mathbf{v}\in {\mathbb{R}}^{d_v}$，每个注意力头${\mathbf{h}}_i$（$i=1,\dots ,h$）的计算方法为：
$${\mathbf{h}}_i=f({\mathbf{W}}_i^{(q)}\mathbf{q},{\mathbf{W}}_i^{(k)}\mathbf{k},{\mathbf{W}}_i^{(v)}\mathbf{v})\in {\mathbb{R}}^{p_v},$$
其中，可学习的参数包括${\mathbf{W}}_i^{(q)}\in {\mathbb{R}}^{p_q\times d_q}$、${\mathbf{W}}_i^{(k)}\in {\mathbb{R}}^{p_k\times d_k}$和${\mathbf{W}}_i^{(v)}\in {\mathbb{R}}^{p_v\times d_v}$，以及代表注意力汇聚的函数$f$。$f$可以是加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换，它对应着$h$个头连结后的结果，因此其可学习参数是${\mathbf{W}}_o\in {\mathbb{R}}^{p_o\times h{p}_v}$：
$${\mathbf{W}}_o\left[\begin{array}{c}{\mathbf{h}}_1\\ ⋮\\ {\mathbf{h}}_h\end{array}\right]\in {\mathbb{R}}^{p_o}.$$

基于这种设计，每个头都可能会关注输入的不同部分，可以表示比简单加权平均值更复杂的函数。

3. 实现

在实现过程中，选择缩放点积注意力作为每一个注意力头。为了避免计算代价和参数代价的大幅增长，设定$p_q=p_k=p_v=p_o/h$。值得注意的是，如果将查询、键和值的线性变换的输出数量设置为$p_{q}h=p_{k}h=p_{v}h=p_o$，则可以并行计算$h$个头。在下面的实现中，$p_o$是通过参数num_hiddens指定的。

class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self,query_size,key_size,value_size,num_hiddens,num_heads,dropout,bias=False):
        super(MultiHeadAttention,self).__init__()
        self.num_heads = num_heads
        self.attention = d2l.torch.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)
    def forward(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 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)

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

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])

def transpose_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类。 多头注意力输出的形状是（batch_size，num_queries，num_hiddens）。

num_hiddens,num_heads = 100,5
multiHeadAttention = MultiHeadAttention(num_hiddens,num_hiddens,num_hiddens,num_hiddens,5,0.5)
multiHeadAttention.eval()

输出结果如下：
MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)

batch_size,num_queries = 2,4
num_kvpairs,valid_lens = 6,torch.tensor([3,2])
Y = torch.ones(size=(batch_size,num_kvpairs,num_hiddens))
X = torch.ones(size=(batch_size,num_queries,num_hiddens))
multiHeadAttention(X,Y,Y,valid_lens).shape

输出结果如下：
torch.Size([2, 4, 100])

4. 小结

• 多头注意力融合了来自于多个注意力汇聚的不同知识，这些知识的不同来源于相同的查询、键和值的不同的子空间（经过不同的全连接线性层的变换）表示。
• 基于适当的张量操作，可以实现多头注意力的并行计算。

5. 全部代码

import torch
import d2l.torch
from torch import nn


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])


def transpose_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)


class MultiHeadAttention(nn.Module):
    """多头注意力"""

    def __init__(self, query_size, key_size, value_size, num_hiddens, num_heads, dropout, bias=False):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.attention = d2l.torch.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)

    def forward(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 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)


num_hiddens, num_heads = 100, 5
multiHeadAttention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, num_hiddens, 5, 0.5)
multiHeadAttention.eval()
batch_size, num_queries = 2, 4
num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
Y = torch.ones(size=(batch_size, num_kvpairs, num_hiddens))
X = torch.ones(size=(batch_size, num_queries, num_hiddens))
multiHeadAttention(X, Y, Y, valid_lens).shape

6. 为什么需要多头注意力

• 所谓自注意力机制就是通过某种运算来直接计算得到句子在编码过程中每个位置上的注意力权重；然后再以权重和的形式来计算得到整个句子的隐含向量表示。
• 自注意力机制的缺陷就是：模型在对当前位置的信息进行编码时，会过度的将注意力集中于自身的位置， 因此作者提出了通过多头注意力机制来解决这一问题。

7. 相关链接

注意力机制第一篇：李沐动手学深度学习V2-注意力机制
注意力机制第二篇：李沐动手学深度学习V2-注意力评分函数
注意力机制第三篇：李沐动手学深度学习V2-基于注意力机制的seq2seq
注意力机制第四篇：李沐动手学深度学习V2-自注意力机制之位置编码
注意力机制第五篇：李沐动手学深度学习V2-自注意力机制
注意力机制第六篇：李沐动手学深度学习V2-多头注意力机制和代码实现