注意力分?jǐn)?shù)

將抽象的注意力機(jī)制進(jìn)行量化

過程描述：已有的數(shù)據(jù)集是若干Key與Value配對，當(dāng)有輸入Query時，將其與每個Key共同輸入函數(shù)a進(jìn)行處理，來比較Query與各個Key的相似度。下圖中的a是softmax。

在實(shí)際問題中，Query經(jīng)常不以一個值的形式而是以向量的形式出現(xiàn)。而函數(shù)a的輸出則往往是一個值，用來評估相似度。

注意力分?jǐn)?shù)的高緯度實(shí)現(xiàn)：

這里Query，即q是一個長度為q的向量，k是一個長度為k的向量。v是一個長度為v的向量。

注意力池化層：將Query和數(shù)據(jù)集的Key-Value對輸入，將q與k輸入α函數(shù)，計(jì)算每個q與key的注意力權(quán)重，再與vi相乘，最后相加，最終得到的數(shù)據(jù)的結(jié)構(gòu)是一個長度為v 的向量。

α函數(shù)的構(gòu)成通常是一個softmax函數(shù)內(nèi)嵌一個注意力分?jǐn)?shù)函數(shù)a。

常用的注意力分?jǐn)?shù)函數(shù)a：

可加性注意力：Additive Attention

原理（下圖中的value應(yīng)為query，tanh為激活函數(shù)）：

當(dāng)query和key長度一致時，可以使用以下方法，其中除以根號d的作用是使數(shù)據(jù)長度d對計(jì)算結(jié)果的影響沒有那么大。

向量化版本：當(dāng)有n個Query和m個Key時，注意力分?jǐn)?shù)函數(shù)為Q與K的轉(zhuǎn)置進(jìn)行相乘。得到一個n*m的矩陣。經(jīng)過注意力池化后，對于每個query都會返回一個長度為v的數(shù)據(jù)（一共n個）

（對m個鍵值對查詢n次，每次查詢得到一個值，最終得到n個值）

總結(jié)：注意力分?jǐn)?shù)是query和key相似度的一個計(jì)算（沒有normalize過）。注意力權(quán)重則是經(jīng)過normalize處理（即softmax）的結(jié)果。

代碼實(shí)現(xiàn)：

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

valid_lens是用來指定數(shù)據(jù)中進(jìn)行softmax運(yùn)算的有效數(shù)位。

#@save
def masked_softmax(X, valid_lens):
    """通過在最后一個軸上掩蔽元素來執(zhí)行softmax操作"""
    # X:3D張量，valid_lens:1D或2D張量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一軸上被掩蔽的元素使用一個非常大的負(fù)值替換，從而其softmax輸出為0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

為了演示此函數(shù)是如何工作的， 考慮由兩個2×4矩陣表示的樣本， 這兩個樣本的有效長度分別為2和3。 經(jīng)過掩蔽softmax操作，超出有效長度的值都被掩蔽為0。

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

tensor([[[0.5980, 0.4020, 0.0000, 0.0000],
         [0.5548, 0.4452, 0.0000, 0.0000]],

        [[0.3716, 0.3926, 0.2358, 0.0000],
         [0.3455, 0.3337, 0.3208, 0.0000]]])

同樣，也可以使用二維張量，為矩陣樣本中的每一行指定有效長度。

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

tensor([[[1.0000, 0.0000, 0.0000, 0.0000],
         [0.4125, 0.3273, 0.2602, 0.0000]],

        [[0.5254, 0.4746, 0.0000, 0.0000],
         [0.3117, 0.2130, 0.1801, 0.2952]]])

下面來實(shí)現(xiàn)加性注意力。

注意：forward函數(shù)中的valid_lens是指在計(jì)算softmax時，每個query計(jì)算多少個數(shù)位的key-value的softmax值。

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

    def forward(self, queries, keys, values, valid_lens):
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在維度擴(kuò)展后，
        # queries的形狀：(batch_size，查詢的個數(shù)，1，num_hidden)
        # key的形狀：(batch_size，1，“鍵－值”對的個數(shù)，num_hiddens)
        # 使用廣播方式進(jìn)行求和
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        features = torch.tanh(features)
        # self.w_v僅有一個輸出，因此從形狀中移除最后那個維度。
        # scores的形狀：(batch_size，查詢的個數(shù)，“鍵-值”對的個數(shù))
        scores = self.w_v(features).squeeze(-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # values的形狀：(batch_size，“鍵－值”對的個數(shù)，值的維度)
        return torch.bmm(self.dropout(self.attention_weights), values)

用一個小例子來演示上面的AdditiveAttention類， 其中查詢、鍵和值的形狀為（批量大小，步數(shù)或詞元序列長度，特征大?。?， 實(shí)際輸出為(2,1,20)、(2,10,2)和(2,10,4)。 注意力匯聚輸出的形狀為（批量大小，查詢的步數(shù)，值的維度）。

queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))
# values的小批量，兩個值矩陣是相同的
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()
attention(queries, keys, values, valid_lens)

tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)

盡管加性注意力包含了可學(xué)習(xí)的參數(shù)，但由于本例子中每個鍵都是相同的， 所以注意力權(quán)重是均勻的，由指定的有效長度決定。

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