注意力机制

注意力机制

在“编码器—解码器（seq2seq）”⼀节⾥，解码器在各个时间步依赖相同的背景变量（context vector）来获取输⼊序列信息。当编码器为循环神经⽹络时，背景变量来⾃它最终时间步的隐藏状态。将源序列输入信息以循环单位状态编码，然后将其传递给解码器以生成目标序列。然而这种结构存在着问题，尤其是RNN机制实际中存在长程梯度消失的问题，对于较长的句子，我们很难寄希望于将输入的序列转化为定长的向量而保存所有的有效信息，所以随着所需翻译句子的长度的增加，这种结构的效果会显著下降。

与此同时，解码的目标词语可能只与原输入的部分词语有关，而并不是与所有的输入有关。例如，当把“Hello world”翻译成“Bonjour le monde”时，“Hello”映射成“Bonjour”，“world”映射成“monde”。在seq2seq模型中，解码器只能隐式地从编码器的最终状态中选择相应的信息。然而，注意力机制可以将这种选择过程显式地建模。

注意力机制框架

Attention 是一种通用的带权池化方法，输入由两部分构成：询问（query）和键值对（key-value pairs）。 Query , attention layer得到输出与value的维度一致 . 对于一个query来说，attention layer 会与每一个key计算注意力分数并进行权重的归一化，输出的向量则是value的加权求和，而每个key计算的权重与value一一对应。

为了计算输出，我们首先假设有一个函数 用于计算query和key的相似性，然后可以计算所有的 attention scores $a_1,\dots ,a_n$by
$$a_i=\alpha (\mathbf{q},{\mathbf{k}}_i)$$
我们使用 softmax函数 获得注意力权重：
$$b_1,\dots ,b_n=\text{softmax}(a_1,\dots ,a_n)$$
最终的输出就是value的加权求和：
$$\mathbf{o}=\sum _{i=1}^n{b}_i{\mathbf{v}}_i$$

不同的attetion layer的区别在于score函数的选择，在本节的其余部分，我们将讨论两个常用的注意层 Dot-product Attention 和 Multilayer Perceptron Attention；随后我们将实现一个引入attention的seq2seq模型并在英法翻译语料上进行训练与测试。

softmax的屏蔽

def SequenceMask(X, X_len,value=-1e6):
    maxlen = X.size(1)
    #print(X.size(),torch.arange((maxlen),dtype=torch.float)[None, :],'\n',X_len[:, None] )
    mask = torch.arange((maxlen),dtype=torch.float)[None, :] >= X_len[:, None]   
    #print(mask)
    X[mask]=value
    return X
def masked_softmax(X, valid_length):
    # X: 3-D tensor, valid_length: 1-D or 2-D tensor
    softmax = nn.Softmax(dim=-1)
    if valid_length is None:
        return softmax(X)
    else:
        shape = X.shape
        if valid_length.dim() == 1:
            try:
                valid_length = torch.FloatTensor(valid_length.numpy().repeat(shape[1], axis=0))#[2,2,3,3]
            except:
                valid_length = torch.FloatTensor(valid_length.cpu().numpy().repeat(shape[1], axis=0))#[2,2,3,3]
        else:
            valid_length = valid_length.reshape((-1,))
        # fill masked elements with a large negative, whose exp is 0
        X = SequenceMask(X.reshape((-1, shape[-1])), valid_length)
 
        return softmax(X).reshape(shape)

masked_softmax(torch.rand((2,2,4),dtype=torch.float), torch.FloatTensor([2,3]))

超出二维矩阵的乘法

X和 Y是维度分别为(b,n,m)和(b, m, k)的张量，进行 b次二维矩阵乘法后得到 , 维度为 (b, n, k)。
$$Z[i,:,:]=dot(X[i,:,:],Y[i,:,:])fori=1,\dots ,n$$

torch.bmm(torch.ones((2,1,3), dtype = torch.float), torch.ones((2,3,2), dtype = torch.float))

点积注意力

The dot product 假设query和keys有相同的维度, 即 . 通过计算query和key转置的乘积来计算attention score,通常还会除去sqrt{d}减少计算出来的score对维度𝑑的依赖性，如下
𝛼(𝐪,𝐤)=⟨𝐪,𝐤⟩/ \sqrt{d}
假设𝐐∈ℝ^{𝑚×𝑑}有m个query， 有n个keys. 我们可以通过矩阵运算的方式计算所有mn个score：
𝛼(𝐐,𝐊)=𝐐𝐊^𝑇/\sqrt{d}
现在让我们实现这个层，它支持一批查询和键值对。此外，它支持作为正则化随机删除一些注意力权重.

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

    # query: (batch_size, #queries, d)
    # key: (batch_size, #kv_pairs, d)
    # value: (batch_size, #kv_pairs, dim_v)
    # valid_length: either (batch_size, ) or (batch_size, xx)
    def forward(self, query, key, value, valid_length=None):
        d = query.shape[-1]
        # set transpose_b=True to swap the last two dimensions of key
        
        scores = torch.bmm(query, key.transpose(1,2)) / math.sqrt(d)
        attention_weights = self.dropout(masked_softmax(scores, valid_length))
        print("attention_weight\n",attention_weights)
        return torch.bmm(attention_weights, value)

多层感知机注意力

将score函数定义:
$$a(\mathbf{s},\mathbf{h})={\mathbf{v}}^{\top }\tanh ({\mathbf{W}}_s\mathbf{s}+{\mathbf{W}}_h\mathbf{h}),$$
. 然后将key 和 value 在特征的维度上合并（concatenate），然后送至 a single hidden layer perceptron 这层中 hidden layer 为 ℎ and 输出的size为 1 .隐层激活函数为tanh，无偏置.

# Save to the d2l package.
class MLPAttention(nn.Module):  
    def __init__(self, units,ipt_dim,dropout, **kwargs):
        super(MLPAttention, self).__init__(**kwargs)
        # Use flatten=True to keep query's and key's 3-D shapes.
        self.W_k = nn.Linear(ipt_dim, units, bias=False)
        self.W_q = nn.Linear(ipt_dim, units, bias=False)
        self.v = nn.Linear(units, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, valid_length):
        query, key = self.W_k(query), self.W_q(key)
        #print("size",query.size(),key.size())
        # expand query to (batch_size, #querys, 1, units), and key to
        # (batch_size, 1, #kv_pairs, units). Then plus them with broadcast.
        features = query.unsqueeze(2) + key.unsqueeze(1)
        #print("features:",features.size())  #--------------开启
        scores = self.v(features).squeeze(-1) 
        attention_weights = self.dropout(masked_softmax(scores, valid_length))
        return torch.bmm(attention_weights, value)

计算背景变量

我们先描述第一个关键点，即计算背景变量。图描绘了注意力机制如何为解码器在时间步2计算背景变量。首先，函数$a$根据解码器在时间步1的隐藏状态和编码器在各个时间步的隐藏状态计算softmax运算的输入。softmax运算输出概率分布并对编码器各个时间步的隐藏状态做加权平均，从而得到背景变量。

具体来说，令编码器在时间步$t$的隐藏状态为${\mathbf{h}}_t$，且总时间步数为$T$。那么解码器在时间步$t^{\prime }$的背景变量为所有编码器隐藏状态的加权平均：

$${\mathbf{c}}_{t^{\prime }}=\sum _{t=1}^T{\alpha }_{t^{\prime }t}{\mathbf{h}}_t,$$

其中给定$t^{\prime }$时，权重${\alpha }_{t^{\prime }t}$在$t=1,\dots ,T$的值是一个概率分布。为了得到概率分布，我们可以使用softmax运算:

$${\alpha }_{t^{\prime }t}=\frac{\exp (e_{t^{\prime }t})}{{\sum }_{k=1}^T\exp (e_{t^{\prime }k})},t=1,\dots ,T.$$

现在，我们需要定义如何计算上式中softmax运算的输入$e_{t^{\prime }t}$。由于$e_{t^{\prime }t}$同时取决于解码器的时间步$t^{\prime }$和编码器的时间步$t$，我们不妨以解码器在时间步$t^{\prime }-1$的隐藏状态${\mathbf{s}}_{t^{\prime }-1}$与编码器在时间步$t$的隐藏状态${\mathbf{h}}_t$为输入，并通过函数$a$计算$e_{t^{\prime }t}$：

$$e_{t^{\prime }t}=a({\mathbf{s}}_{t^{\prime }-1},{\mathbf{h}}_t).$$

这里函数$a$有多种选择，如果两个输入向量长度相同，一个简单的选择是计算它们的内积$a(\mathbf{s},\mathbf{h})={\mathbf{s}}^{\top }\mathbf{h}$。而最早提出注意力机制的论文则将输入连结后通过含单隐藏层的多层感知机变换 [1]：

$$a(\mathbf{s},\mathbf{h})={\mathbf{v}}^{\top }\tanh ({\mathbf{W}}_s\mathbf{s}+{\mathbf{W}}_h\mathbf{h}),$$

其中$\mathbf{v}$、${\mathbf{W}}_s$、${\mathbf{W}}_h$都是可以学习的模型参数。

矢量化计算

我们还可以对注意力机制采用更高效的矢量化计算。广义上，注意力机制的输入包括查询项以及一一对应的键项和值项，其中值项是需要加权平均的一组项。在加权平均中，值项的权重来自查询项以及与该值项对应的键项的计算。

在上面的例子中，查询项为解码器的隐藏状态，键项和值项均为编码器的隐藏状态。
让我们考虑一个常见的简单情形，即编码器和解码器的隐藏单元个数均为$h$，且函数$a(\mathbf{s},\mathbf{h})={\mathbf{s}}^{\top }\mathbf{h}$。假设我们希望根据解码器单个隐藏状态${\mathbf{s}}_{t^{\prime }-1}\in {\mathbb{R}}^h$和编码器所有隐藏状态${\mathbf{h}}_t\in {\mathbb{R}}^h,t=1,\dots ,T$来计算背景向量${\mathbf{c}}_{t^{\prime }}\in {\mathbb{R}}^h$。
我们可以将查询项矩阵$\mathbf{Q}\in {\mathbb{R}}^{1\times h}$设为${\mathbf{s}}_{t^{\prime }-1}^{\top }$，并令键项矩阵$\mathbf{K}\in {\mathbb{R}}^{T\times h}$和值项矩阵$\mathbf{V}\in {\mathbb{R}}^{T\times h}$相同且第$t$行均为${\mathbf{h}}_t^{\top }$。此时，我们只需要通过矢量化计算

$$\text{softmax}(\mathbf{Q}{\mathbf{K}}^{\top })\mathbf{V}$$

即可算出转置后的背景向量${\mathbf{c}}_{t^{\prime }}^{\top }$。当查询项矩阵$\mathbf{Q}$的行数为$n$时，上式将得到$n$行的输出矩阵。输出矩阵与查询项矩阵在相同行上一一对应。

引入注意力机制的S2S

本节中将注意机制添加到sequence to sequence 模型中，以显式地使用权重聚合states。下图展示encoding 和decoding的模型结构，在时间步为t的时候。此刻attention layer保存着encodering看到的所有信息——即encoding的每一步输出。在decoding阶段，解码器的时刻的隐藏状态被当作query，encoder的每个时间步的hidden states作为key和value进行attention聚合. Attetion model的输出当作成上下文信息context vector，并与解码器输入拼接起来一起送到解码器：

下图展示了seq2seq机制的所以层的关系，下面展示了encoder和decoder的layer结构

Decoder

由于带有注意机制的seq2seq的编码器与之前章节中的Seq2SeqEncoder相同，所以在此处我们只关注解码器。我们添加了一个MLP注意层(MLPAttention)，它的隐藏大小与解码器中的LSTM层相同。然后我们通过从编码器传递三个参数来初始化解码器的状态:

• the encoder outputs of all timesteps：encoder输出的各个状态，被用于attetion layer的memory部分，有相同的key和values
• the hidden state of the encoder’s final timestep：编码器最后一个时间步的隐藏状态，被用于初始化decoder 的hidden state
• the encoder valid length: 编码器的有效长度，借此，注意层不会考虑编码器输出中的填充标记（Paddings）

在解码的每个时间步，我们使用解码器的最后一个RNN层的输出作为注意层的query。然后，将注意力模型的输出与输入嵌入向量连接起来，输入到RNN层。虽然RNN层隐藏状态也包含来自解码器的历史信息，但是attention model的输出显式地选择了enc_valid_len以内的编码器输出，这样attention机制就会尽可能排除其他不相关的信息。

class Seq2SeqAttentionDecoder(d2l.Decoder):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        self.attention_cell = MLPAttention(num_hiddens,num_hiddens, dropout)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.LSTM(embed_size+ num_hiddens,num_hiddens, num_layers, dropout=dropout)
        self.dense = nn.Linear(num_hiddens,vocab_size)

    def init_state(self, enc_outputs, enc_valid_len, *args):
        outputs, hidden_state = enc_outputs
#         print("first:",outputs.size(),hidden_state[0].size(),hidden_state[1].size())
        # Transpose outputs to (batch_size, seq_len, hidden_size)
        return (outputs.permute(1,0,-1), hidden_state, enc_valid_len)
        #outputs.swapaxes(0, 1)
        
    def forward(self, X, state):
        enc_outputs, hidden_state, enc_valid_len = state
        #("X.size",X.size())
        X = self.embedding(X).transpose(0,1)
#         print("Xembeding.size2",X.size())
        outputs = []
        for l, x in enumerate(X):
#             print(f"\n{l}-th token")
#             print("x.first.size()",x.size())
            # query shape: (batch_size, 1, hidden_size)
            # select hidden state of the last rnn layer as query
            query = hidden_state[0][-1].unsqueeze(1) # np.expand_dims(hidden_state[0][-1], axis=1)
            # context has same shape as query
#             print("query enc_outputs, enc_outputs:\n",query.size(), enc_outputs.size(), enc_outputs.size())
            context = self.attention_cell(query, enc_outputs, enc_outputs, enc_valid_len)
            # Concatenate on the feature dimension
#             print("context.size:",context.size())
            x = torch.cat((context, x.unsqueeze(1)), dim=-1)
            # Reshape x to (1, batch_size, embed_size+hidden_size)
#             print("rnn",x.size(), len(hidden_state))
            out, hidden_state = self.rnn(x.transpose(0,1), hidden_state)
            outputs.append(out)
        outputs = self.dense(torch.cat(outputs, dim=0))
        return outputs.transpose(0, 1), [enc_outputs, hidden_state,
                                        enc_valid_len]

encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8,
                            num_hiddens=16, num_layers=2)
# encoder.initialize()
decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8,
                                  num_hiddens=16, num_layers=2)
X = torch.zeros((4, 7),dtype=torch.long)
print("batch size=4\nseq_length=7\nhidden dim=16\nnum_layers=2\n")
print('encoder output size:', encoder(X)[0].size())
print('encoder hidden size:', encoder(X)[1][0].size())
print('encoder memory size:', encoder(X)[1][1].size())
state = decoder.init_state(encoder(X), None)
out, state = decoder(X, state)
out.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape