注意力分数
在上一节中,阐述了注意力机制,现在将其拓展到高维度:
假设query q ∈ R q q\in R^q q∈Rq,m对key-value ( k 1 , v 1 ) , ⋯ , (k_1,v_1),\cdots, (k1,v1),⋯,这里 k i ∈ R k , v i ∈ R v k_i \in R^k,v_i\in R^v ki∈Rk,vi∈Rv
注意力池化层:
f
(
q
,
(
k
1
,
v
1
)
,
⋯
,
(
k
m
,
v
m
)
)
=
∑
i
=
1
m
α
(
q
,
k
i
)
v
i
∈
R
v
α
(
q
,
k
i
)
=
s
o
f
t
m
a
x
(
a
(
q
,
k
i
)
)
=
e
x
p
(
a
(
q
,
k
i
)
)
∑
j
=
1
m
e
x
p
(
a
(
q
,
k
j
)
)
∈
R
f(q,(k_1,v_1),\cdots,(k_m,v_m))=\sum ^m_{i=1}\alpha(q,k_i)v_i\in R^v\\ \alpha(q,k_i) = softmax(a(q,k_i))=\frac{exp(a(q,k_i))}{\sum ^m _{j=1} exp(a(q,k_j))}\in R
f(q,(k1,v1),⋯,(km,vm))=i=1∑mα(q,ki)vi∈Rvα(q,ki)=softmax(a(q,ki))=∑j=1mexp(a(q,kj))exp(a(q,ki))∈R
函数a计算出来的就是注意力分数,下面介绍两种注意力评分函数的设计
1.掩蔽softmax操作
有时并非所有的值都被纳入到注意力池化中,例如,某些文本序列被填充了没有意义的特殊词元,所以需要进行掩蔽:
#@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.Additive Attention
加性注意力
一般来说,当查询和键是不同长度的矢量时,可以使用加性注意力作为评分函数。
可学参数:
W
k
∈
R
h
×
k
,
W
q
∈
R
h
×
q
,
v
∈
R
h
W_k\in R^{h\times k},W_q\in R^{h\times q},v\in R^h
Wk∈Rh×k,Wq∈Rh×q,v∈Rh
a
(
k
,
q
)
=
v
T
t
a
n
h
(
W
k
k
+
W
q
q
)
a(k,q) =v^T tanh(W_kk+W_qq)
a(k,q)=vTtanh(Wkk+Wqq)
等价于将key和query合并起来放入到一个隐藏大小为h输出大小为h输出大小为1的单隐藏层MLP,使用tanh作为激活函数,并且不适用偏置项。
#@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) # 最后输出的单隐藏层MLP
self.dropout = nn.Dropout(dropout)
def forward(self, queries, keys, values, valid_lens):
queries, keys = self.W_q(queries), self.W_k(keys)
# 在维度扩展后(unsqueeze是为了能广播)
# 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 torch.bmm(self.dropout(self.attention_weights), values)
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)
'''注意力权重'''
d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
xlabel='Keys', ylabel='Queries')
3. Scaled Dot-Product Attention
缩放点积注意力
如果query和key都是同样的长度
q
,
k
i
∈
R
d
q,k_i\in R ^d
q,ki∈Rd,并假设查询和键的所有元素都是独立的随机变量, 并且都满足零均值和单位方差, 那么两个向量的点积的均值为0,方差为d。 为确保无论向量长度如何, 点积的方差在不考虑向量长度的情况下仍然是1, 我们再将点积除以
d
\sqrt d
d:
a
(
q
,
k
i
)
=
<
q
,
k
i
>
d
a(q,k_i) =\frac{<q,k_i>}{\sqrt{d}}
a(q,ki)=d<q,ki>
在实践中,我们通常从小批量的角度来考虑提高效率, 例如基于n个查询和m个键-值对计算注意力, 其中查询和键的长度为d,值的长度为v。
查询
Q
∈
R
n
×
d
Q\in R^{n\times d}
Q∈Rn×d,键
K
∈
R
m
×
d
K\in R ^{m\times d}
K∈Rm×d和值
V
∈
R
m
×
v
V\in R^{m\times v}
V∈Rm×v的缩放点积注意力是:
s
o
f
t
m
a
x
(
Q
K
T
d
)
V
∈
R
n
×
v
softmax(\frac{QK^T}{\sqrt{d}})V \in R ^{n\times v}
softmax(dQKT)V∈Rn×v
#@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_lens=None):
d = queries.shape[-1]
# 设置transpose_b=True为了交换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)
queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
attention(queries, keys, values, valid_lens)
d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
xlabel='Keys', ylabel='Queries')
注意力分数是query和key的相似度,注意力权重是分数的softmax结果。