注意力权重计算的原理及实现

注意力分数

计算注意力汇聚的输出为值的加权和：把查询query和每一个键key都做一个注意力评分函数a，会得到一个个值，叫注意力分数，把全部这些值经过softmax就可以得到注意力权重，把注意力权重和对应的值value进行相乘再求和就可以得到输出了。

拓展到高维度：

Additive Attention

其等价于将key和query合并起来后放入到一个隐藏大小为h输出大小为1的单隐藏层MLP:

分块矩阵乘法：

意思就是对于每个查询的query，总共有n个，最后都会对应一个一行的长度为v的向量，所以最后是一个n*v的注意力池化矩阵。

向量版本代码实现

增加了批量维度batch_size，每一个批次有batch_size个样本，每个样本有num_queries个查询，每个查询的维度是query_size，并且每个查询都对应了同样的m个key-value pairs ，并且每个k是key_size维的，每个value是value_size维度的。

import math
import torch
from torch import nn
from d2l import torch as d2l
import matplotlib.pyplot as plt

def masked_softmax(X, valid_lens):
    # 在输入X上的最后一个维度/轴上 屏蔽元素 再来进行softmax
    #  X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(input=X, dim=-1)
    else:
        shape = X.shape
        # 如果 valid_lens 是 1D 张量（形状为 (batch_size,)），表示每个样本的有效长度。
        # 如果 valid_lens 是 2D 张量（形状为 (batch_size, num_rows)），表示每个样本的每一行的有效长度。
        if valid_lens.dim() == 1:
            # 这里shape[1]是输入X的行数 就是把每个样本的有效列数重复行数次 使得可以与X的shape相匹配
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            # 降维 变成一行向量
            valid_lens = valid_lens.reshape(-1)
        # 这里sequence_mask 传入的第一个参数是二维的：样本数*需要softmax的特征数，
        # 第二个参数表示有效长度是一维的 有多少个样本就有多少个数据 最后一个参数表示填充值
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)
        # e的-1e6次方才就是很小 接近0
        return nn.functional.softmax(X.reshape(shape), dim=-1)


# print(masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3])))
# print(masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]])))


class AdditiveAttention(nn.Module):
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    # 1. 变量的初始形状
    # 输入的 queries、keys、values 具有以下形状：
    # queries: (batch_size, num_queries, query_size)
    # batch_size：批量大小 样本个数
    # num_queries：每个样本的查询个数
    # query_size：查询向量的维度 q
    # keys: (batch_size, num_kv_pairs, key_size)
    # num_kv_pairs：每个样本的键-值对个数 m
    # key_size：键向量的维度
    # values: (batch_size, num_kv_pairs, value_size)
    # value_size：值的向量维度
    # valid_lens: (batch_size,) 或 (batch_size, num_queries)
    # 表示有效键-值对的个数（用于掩码 softmax）
    def forward(self, queries, keys, values, valid_lens):
        # queries变成(batch_size, num_queries, num_hiddens) keys同理
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在维度扩展后，
        # queries的形状：(batch_size，查询的个数，1，num_hidden)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        features = torch.tanh(features)
        # self.w_v仅有一个输出，因此从形状中移除最后那个维度。
        # scores的形状：(batch_size，查询的个数，“键-值”对的个数)
        scores = self.w_v(features).squeeze(-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # values的形状：(batch_size，“键－值”对的个数，值的维度)
        # 因此最后return出来的结果的shape为(batch_size, 查询的个数, 值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)



queries, keys = torch.normal(mean=0, std=1, size=(2, 1, 20)), torch.ones((2, 10, 2))
values= torch.arange(40, dtype=torch.float32).reshape(1, 10, 4).repeat(2, 1, 1)
valid_lens = torch.tensor([2, 6])

attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8,
                              dropout=0.1)
attention.eval()
print(attention(queries, keys, values, valid_lens))
print(attention.attention_weights)


d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')
plt.show()


class DotProductAttention(nn.Module):
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 交换keys的最后两个维度 变成转置 就是(batch_size，d, “键－值”对的个数)
        scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
        # 注意力权重的shape为 (batch_size, 查询的个数, “键－值”对的个数)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)



queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
print(attention(queries, keys, values, valid_lens))
print(attention.attention_weights)
d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')
plt.show()

参考连接

10.3. 注意力评分函数 — 动手学深度学习 2.0.0 documentation