《动手学深度学习 Pytorch版》 10.5 多头注意力

多头注意力(multihead attention):用独立学习得到的 h 组不同的线性投影(linear projections)来变换查询、键和值,然后并行地送到注意力汇聚中。最后,将这 h 个注意力汇聚的输出拼接在一起,并且通过另一个可以学习的线性投影进行变换,以产生最终输出。

对于 h 个注意力汇聚输出,每一个注意力汇聚都被称作一个头(head)。

《动手学深度学习 Pytorch版》 10.5 多头注意力_第1张图片

10.5.1 模型

用数学语言描述多头注意力:

h i = f ( W i ( q ) q , W i ( k ) k , W i ( v ) v ) ∈ R p \boldsymbol{h}_i=f(\boldsymbol{W}_i^{(q)}\boldsymbol{q},\boldsymbol{W}_i^{(k)}\boldsymbol{k},\boldsymbol{W}_i^{(v)}\boldsymbol{v})\in\R^p hi=f(Wi(q)q,Wi(k)k,Wi(v)v)Rp

参数字典:

  • f f f 表示注意力汇聚函数

  • q ∈ R d q \boldsymbol{q}\in\R^{d_q} qRdq k ∈ R d k \boldsymbol{k}\in\R^{d_k} kRdk v ∈ R d v \boldsymbol{v}\in\R^{d_v} vRdv 分别是查询、键和值

  • W i ( q ) ∈ R p d × d q \boldsymbol{W}_i^{(q)}\in\R^{p_d\times d_q} Wi(q)Rpd×dq W i ( k ) ∈ R p k × d k \boldsymbol{W}_i^{(k)}\in\R^{p_k\times d_k} Wi(k)Rpk×dk W i ( v ) ∈ R p v × d v \boldsymbol{W}_i^{(v)}\in\R^{p_v\times d_v} Wi(v)Rpv×dv 均为可学习参数

多头注意力的输出需要经过另一个线性转换:

y = [ h 1 ⋮ h h ] ∈ R p o y= \begin{bmatrix} \boldsymbol{h}_1\\ \vdots\\ \boldsymbol{h}_h \end{bmatrix} \in\R^{p_o} y= h1hh Rpo

import math
import torch
from torch import nn
from d2l import torch as d2l

10.5.2 实现

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

MultiHeadAttention 类将使用下面定义的两个转置函数,transpose_output 函数反转了 transpose_qkv 函数的操作。转来转去是为了避免 for 循环。

#@save
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])


#@save
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)
#@save
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_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)

    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
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.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])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape
torch.Size([2, 4, 100])

练习

(1)分别可视化这个实验中的多个头的注意力权重。

d2l.show_heatmaps(attention.attention.attention_weights.reshape((2, 5, 4, 6)),
                  xlabel='Keys', ylabel='Queries', figsize=(5,5))


《动手学深度学习 Pytorch版》 10.5 多头注意力_第2张图片

(2)假设有一个完成训练的基于多头注意力的模型,现在希望修剪最不重要的注意力头以提高预测速度。如何设计实验来衡量注意力头的重要性呢?

不会,略。

你可能感兴趣的:(《动手学深度学习,Pytorch版》学习笔记,深度学习,pytorch,人工智能)