动手学深度学习——注意力机制读书笔记 注意力机制简单来说就是面对一个查询query从一堆键值对key-value中根据查询与各个键的匹配程度对值做加权平均。俩个很容易混淆的函数一.注意力汇聚Nadaraya-Watson 核回归假如我们有50个点横坐标记为X,纵坐标记为Y。现在我们有50个只知道横坐标的点横坐标记为x,我们需要预测他们的纵坐标y。这里我们就可以采用注意力汇聚x可以看成查询XY看出key-value,我们根据x和X的关系对Y做加权平均来预测y的值。创作一些点和一个函数。n_train 50 # 训练样本数 x_train, _ torch.sort(torch.rand(n_train) * 5) # 排序后的训练样本 def f(x): return 2 * torch.sin(x) x**0.8 y_train f(x_train) torch.normal(0.0, 0.5, (n_train,)) # 训练样本的输出 x_test torch.arange(0, 5, 0.1) # 测试样本 y_truth f(x_test) # 测试样本的真实输出 n_test len(x_test) # 测试样本数 n_test1.1非参数注意力汇聚x是一个查询将查询和键之间的关系建模为阿尔法称为注意力权重attention weight# X_repeat的形状:(n_test,n_train), # 每一行都包含着相同的测试输入例如同样的查询 X_repeat x_test.repeat_interleave(n_train).reshape((-1, n_train)) # x_train包含着键。attention_weights的形状(n_test,n_train), # 每一行都包含着要在给定的每个查询的值y_train之间分配的注意力权重 attention_weights nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim1) # y_hat的每个元素都是值的加权平均值其中的权重是注意力权重 y_hat torch.matmul(attention_weights, y_train) plot_kernel_reg(y_hat)运行结果1.2批量矩阵乘法在训练时假设所有查询都是针对一个key-value对我们使用的是一个批量每个序列都有自己的注意力权重和key-value对。注意力权重假设为批量大小q的数量key-value对的数量序列值假设为批量大小key-value对的数量1。这里为了更快速的实注意力权重矩阵和序列值矩阵的相乘我们可以使用批量矩阵乘法。X torch.ones((2, 1, 4)) Y torch.ones((2, 4, 6)) torch.bmm(X, Y).shapetensor([[[ 4.5000]], [[14.5000]]])也就是看成2个1*4和2个4*6的矩阵互相相乘再结合起来1.3带参数注意力汇聚模型class NWKernelRegression(nn.Module): def __init__(self, **kwargs): super().__init__(**kwargs) self.w nn.Parameter(torch.rand((1,), requires_gradTrue)) def forward(self, queries, keys, values): # queries和attention_weights的形状为(查询个数“键值”对个数) queries queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1])) self.attention_weights nn.functional.softmax( -((queries - keys) * self.w)**2 / 2, dim1) # values的形状为(查询个数“键值”对个数) return torch.bmm(self.attention_weights.unsqueeze(1), values.unsqueeze(-1)).reshape(-1)训练这里我们再之前数据的基础上在带参数的注意力汇聚模型中 任何一个训练样本的输入都会和除自己以外的所有训练样本的“键值”对进行计算 从而得到其对应的预测输出。也就是说每个查询的key-value都不同。return torch.bmm(self.attention_weights.unsqueeze(1), values.unsqueeze(-1)).reshape(-1)queriesvalue和attention_weights的形状为(查询个数“键值”对个数)value是每个查询*每个查询的value行权重是每个查询*每个查询和key的分数。不能直接运算我们要对维度进行变化才能运算。如权重54变成514value(5,4)变成541就可以通过批量矩阵乘法运算。# X_tile的形状:(n_trainn_train)每一行都包含着相同的训练输入 X_tile x_train.repeat((n_train, 1)) # Y_tile的形状:(n_trainn_train)每一行都包含着相同的训练输出 Y_tile y_train.repeat((n_train, 1)) # keys的形状:(n_trainn_train-1) keys X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1)) # values的形状:(n_trainn_train-1) values Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))net NWKernelRegression() loss nn.MSELoss(reductionnone) trainer torch.optim.SGD(net.parameters(), lr0.5) animator d2l.Animator(xlabelepoch, ylabelloss, xlim[1, 5]) for epoch in range(5): trainer.zero_grad() l loss(net(x_train, keys, values), y_train) l.sum().backward() trainer.step() print(fepoch {epoch 1}, loss {float(l.sum()):.6f}) animator.add(epoch 1, float(l.sum()))预测# keys的形状:(n_testn_train)每一行包含着相同的训练输入例如相同的键 keys x_train.repeat((n_test, 1)) # value的形状:(n_testn_train) values y_train.repeat((n_test, 1)) y_hat net(x_test, keys, values).unsqueeze(1).detach() plot_kernel_reg(y_hat)二.注意力评分函数就是衡量查询和key之间关系强度然后分配注意力权重。2.1掩蔽softmax操作正如上面提到的softmax操作用于输出一个概率分布作为注意力权重。 在某些情况下并非所有的值都应该被纳入到注意力汇聚中。 例如为了在高效处理小批量数据集 某些文本序列被填充了没有意义的特殊词元。 为了仅将有意义的词元作为值来获取注意力汇聚 可以指定一个有效序列长度即词元的个数 以便在计算softmax时过滤掉超出指定范围的位置。长度有一维B和二维BQ一维每个批量的查询共享一个长度需要先扩展到每个查询。然后将所有长度展成一维将输入x变化为B*Q,k进行遮蔽用很大的负值对无效位置进行替换。#save def masked_softmax(X, valid_lens): 通过在最后一个轴上掩蔽元素来执行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) # 最后一轴上被掩蔽的元素使用一个非常大的负值替换从而其softmax输出为0 X d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value-1e6) return nn.functional.softmax(X.reshape(shape), dim-1)2.2加性注意力加性注意力Bahdanau attention先用两个线性层分别把 query 和 key 映射到同一隐空间相加后过tanh再用一个线性层打出标量分数softmax 得到权重最后对values加权求和#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, biasFalse) self.W_q nn.Linear(query_size, num_hiddens, biasFalse) self.w_v nn.Linear(num_hiddens, 1, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, queries, keys, values, valid_lens): queries, keys self.W_q(queries), self.W_k(keys) # 在维度扩展后 # queries的形状(batch_size查询的个数1num_hidden) # key的形状(batch_size1“键值”对的个数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 torch.bmm(self.dropout(self.attention_weights), values)代码非常难懂我们先看传入的参数形状。举个例子B 2 # 2 个句子Q 1 # decoder 当前 1 步1 个 queryK 3 # encoder 输出 3 个词padding 到 3d_q d_k d_v 4 # 每个向量 4 维queries.shape (2, 1, 4) # 2 个句子每句 1 个 query每个 4 维keys.shape (2, 3, 4) # 2 个句子每句 3 个 key每个 4 维values.shape (2, 3, 4) # 同上通常 keys 和 values 是同一个东西valid_lens tensor([3, 2]) # 第 0 句 3 个词都有效第 1 句只有前 2 个有效首先介绍线性层形状变化queries: (2, 1, 4) ──W_q──► (2, 1, h)keys: (2, 3, 4) ──W_k──► (2, 3, h)假设h8广播相加就是让每一个q和所有key相加要借助广播机制之后经过上面的步骤我们就得到每一个q和所有key的注意力分数。之后经过softmax得到权重用权重和值做矩阵乘法得到2.3缩放点积注意力使用点积可以得到计算效率更高的评分函数但是点积操作要求查询和键具有相同的长度d代码如下#save 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_lensNone): d queries.shape[-1] # 设置transpose_bTrue为了交换keys的最后两个维度 scores torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d) self.attention_weights masked_softmax(scores, valid_lens) return torch.bmm(self.dropout(self.attention_weights), values)三.Bahdanau注意力2014 年 Bahdanau 等人提出用来解决老式 Seq2Seq 的毛病。解码器每说一个词之前先回头看一眼原文决定该重点关注原文哪个位置。3.1例子用I love cats→我爱猫举例从头到尾走一遍。outputs, hidden_state enc_outputsh1,h2,h3outputs最后一层的所有隐状态,初始状态hidden_state[-1]取最后一层。3.2代码class Seq2SeqAttentionDecoder(AttentionDecoder): def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout0, **kwargs): super(Seq2SeqAttentionDecoder, self).__init__(**kwargs) self.attention d2l.AdditiveAttention( num_hiddens, num_hiddens, num_hiddens, dropout) self.embedding nn.Embedding(vocab_size, embed_size) self.rnn nn.GRU( embed_size num_hiddens, num_hiddens, num_layers, dropoutdropout) self.dense nn.Linear(num_hiddens, vocab_size) def init_state(self, enc_outputs, enc_valid_lens, *args): # outputs的形状为(batch_sizenum_stepsnum_hiddens). # hidden_state的形状为(num_layersbatch_sizenum_hiddens) outputs, hidden_state enc_outputs return (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens) def forward(self, X, state): # enc_outputs的形状为(batch_size,num_steps,num_hiddens). # hidden_state的形状为(num_layers,batch_size, # num_hiddens) enc_outputs, hidden_state, enc_valid_lens state # 输出X的形状为(num_steps,batch_size,embed_size) X self.embedding(X).permute(1, 0, 2) outputs, self._attention_weights [], [] for x in X: # query的形状为(batch_size,1,num_hiddens) query torch.unsqueeze(hidden_state[-1], dim1) # context的形状为(batch_size,1,num_hiddens) context self.attention( query, enc_outputs, enc_outputs, enc_valid_lens) # 在特征维度上连结 x torch.cat((context, torch.unsqueeze(x, dim1)), dim-1) # 将x变形为(1,batch_size,embed_sizenum_hiddens) out, hidden_state self.rnn(x.permute(1, 0, 2), hidden_state) outputs.append(out) self._attention_weights.append(self.attention.attention_weights) # 全连接层变换后outputs的形状为 # (num_steps,batch_size,vocab_size) outputs self.dense(torch.cat(outputs, dim0)) return outputs.permute(1, 0, 2), [enc_outputs, hidden_state, enc_valid_lens] property def attention_weights(self): return self._attention_weights这里context作用所以context的唯一用途就是给 GRU 提供“原文的相关信息”。GRU 同时看到“当前词”和“原文相关部分”才能决定下一个词是什么。测试encoder d2l.Seq2SeqEncoder(vocab_size10, embed_size8, num_hiddens16, num_layers2) encoder.eval() decoder Seq2SeqAttentionDecoder(vocab_size10, embed_size8, num_hiddens16, num_layers2) decoder.eval() X torch.zeros((4, 7), dtypetorch.long) # (batch_size,num_steps) state decoder.init_state(encoder(X), None) output, state decoder(X, state) output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape运行结果四.多头注意力在实践中当给定相同的查询、键和值的集合时 我们希望模型可以基于相同的注意力机制学习到不同的行为 然后将不同的行为作为知识组合起来 捕获序列内各种范围的依赖关系 例如短距离依赖和长距离依赖关系。 因此允许注意力机制组合使用查询、键和值的不同子空间表示representation subspaces可能是有益的。也即是输入 q/k/v: (B, N, H)拆头后: (B*num_heads, N, H/num_heads) ← 每个头当独立 batch注意力后: (B*num_heads, N, H/num_heads)拼回: (B, N, H)线性变换后: (B, N, H)可以把每个头看出独立的批次进行计算计算完后再拼回去。切分工具函数#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_sizenum_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, biasFalse, **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, biasbias) self.W_k nn.Linear(key_size, num_hiddens, biasbias) self.W_v nn.Linear(value_size, num_hiddens, biasbias) self.W_o nn.Linear(num_hiddens, num_hiddens, biasbias) def forward(self, queries, keys, values, valid_lens): # querieskeysvalues的形状: # (batch_size查询或者“键值”对的个数num_hiddens) # valid_lens 的形状: # (batch_size)或(batch_size查询的个数) # 经过变换后输出的querieskeysvalues 的形状: # (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, repeatsself.num_heads, dim0) # 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)最后我们讲不同头的输入拼起来后形状已经batch_size,查询格式num_hiddens了为什么我们还需要经过一个num_hiddens,num_hiddens的线性层。形状测试代码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