$$ f(x) = sum_(i = 1)^n “softmax”(-1 / 2 ((x - x_i)w)^2)y_i $$
该公式括号内部可称为”评分函数“
masked_softmax
1def masked_softmax(X, valid_lens):2 """在最后一个轴上掩蔽元素后,执行softmax操作"""3 # X:3D张量,valid_lens:1D或2D张量4 # 使用时,X.shape = (batch_size, num_queries, num_keys),表示注意力权重。此函数掩盖部分 keys 的注意力权重5 if valid_lens is None:6 return nn.functional.softmax(X, dim=-1)7 else:8 shape = X.shape9 if valid_lens.dim() == 1:10 valid_lens = torch.repeat_interleave(valid_lens, shape[1])11 else:12 valid_lens = valid_lens.reshape(-1)13 # 最后一轴上被掩蔽的元素使用一个非常大的负值替换,从而其softmax输出为014 X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,15 value=-1e6)1 collapsed line
16 return nn.functional.softmax(X.reshape(shape), dim=-1)
1# 对于第一个 batch,每个 num_keys 保留 22# 对于第二个 batch,每个 num_keys 保留 33masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))4>>>5tensor([[[0.5980, 0.4020, 0.0000, 0.0000],6 [0.5548, 0.4452, 0.0000, 0.0000]],7
8 [[0.3716, 0.3926, 0.2358, 0.0000],9 [0.3455, 0.3337, 0.3208, 0.0000]]])
1masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]]))2>>>3tensor([[[1.0000, 0.0000, 0.0000, 0.0000],4 [0.4125, 0.3273, 0.2602, 0.0000]],5
6 [[0.5254, 0.4746, 0.0000, 0.0000],7 [0.3117, 0.2130, 0.1801, 0.2952]]])
加性注意力
-
学习不同维度的查询与键的注意力。下面输出为实数. q 的大小为 query_size,k 大小为 key_size
-
-
1class AdditiveAttention(nn.Module):2 """加性注意力"""3 def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):4 super(AdditiveAttention, self).__init__(**kwargs)5 self.W_k = nn.Linear(key_size, num_hiddens, bias=False)6 self.W_q = nn.Linear(query_size, num_hiddens, bias=False)7 self.w_v = nn.Linear(num_hiddens, 1, bias=False)8 self.dropout = nn.Dropout(dropout)9
10 def forward(self, queries, keys, values, valid_lens):11 # 输入的 queries 的形状: (batch_size, num_queries, query_size)12 # 输入的 queries 的形状: (batch_size, num_keys, key_size)13
14 queries, keys = self.W_q(queries), self.W_k(keys)15
17 collapsed lines
16 # - 在维度扩展 unsqueeze 后,17 # queries的形状:(batch_size,查询的个数,1,num_hidden)18 # key的形状:(batch_size,1,“键-值”对的个数,num_hiddens)19 # 使用广播方式进行求和20 # - unsqueeze(n) 在第 n 维添加一维,长度为 121 features = queries.unsqueeze(2) + keys.unsqueeze(1)22 features = torch.tanh(features)23
24 # self.w_v仅有一个输出,因此从形状中移除最后那个维度。25 # scores的形状:(batch_size,查询的个数,“键-值”对的个数)26 scores = self.w_v(features).squeeze(-1)27
28 self.attention_weights = masked_softmax(scores, valid_lens)29 # attension_weights.shape = (batch_size, num_queries, num_keys)30 # values的形状:(batch_size,“键-值”对的个数,值的维度)31 return torch.bmm(self.dropout(self.attention_weights), values)32 # return 的形状: (batch_size, num_queries, value_dim)
1# - queries: 2 个批量,分别 1 个 query,分别长度为 202# - 这个 query 查询对所有 key 的权重,都是一样的(虽然参数随机生成,但过程中用到的参数一样)3# - mask 仅保留前两个 key,所以其权重都是 1 / 2,结果 value 为前两个 key 所存 value 的均值4queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))5# values的小批量,两个值矩阵是相同的6values = torch.arange(40, dtype=torch.float32).reshape(1, 10, 4).repeat(7 2, 1, 1)8valid_lens = torch.tensor([2, 6])9attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8,10 dropout=0.3)11attention.eval()12attention(queries, keys, values, valid_lens), values13
14>>>15(tensor([[[ 2.0000, 3.0000, 4.0000, 5.0000]],22 collapsed lines
16 [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>),17 tensor([[[ 0., 1., 2., 3.],18 [ 4., 5., 6., 7.],19 [ 8., 9., 10., 11.],20 [12., 13., 14., 15.],21 [16., 17., 18., 19.],22 [20., 21., 22., 23.],23 [24., 25., 26., 27.],24 [28., 29., 30., 31.],25 [32., 33., 34., 35.],26 [36., 37., 38., 39.]],27
28 [[ 0., 1., 2., 3.],29 [ 4., 5., 6., 7.],30 [ 8., 9., 10., 11.],31 [12., 13., 14., 15.],32 [16., 17., 18., 19.],33 [20., 21., 22., 23.],34 [24., 25., 26., 27.],35 [28., 29., 30., 31.],36 [32., 33., 34., 35.],37 [36., 37., 38., 39.]]]))
- 这两个 batch 对 10 个 key 的注意力分配:
缩放点积注意力
- 原文很简洁:
1class DotProductAttention(nn.Module):2 """缩放点积注意力"""3 def __init__(self, dropout, **kwargs):4 super(DotProductAttention, self).__init__(**kwargs)5 self.dropout = nn.Dropout(dropout)6
7 # queries的形状:(batch_size,查询的个数,d: 即 query_size)8 # keys的形状:(batch_size,“键-值”对的个数,d)9 # values的形状:(batch_size,“键-值”对的个数,值的维度)10 # valid_lens的形状:(batch_size,)或者(batch_size,查询的个数)11 def forward(self, queries, keys, values, valid_lens=None):12 d = queries.shape[-1]13 # 设置transpose_b=True为了交换keys的最后两个维度14 scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)15 self.attention_weights = masked_softmax(scores, valid_lens)1 collapsed line
16 return torch.bmm(self.dropout(self.attention_weights), values)
- [done]