Pytorch 注意力分数
0. 环境介绍
环境使用 Kaggle 里免费建立的 Notebook
小技巧:当遇到函数看不懂的时候可以按 Shift+Tab
查看函数详解。
1. 注意力分数
1.1 回顾
注意力汇聚(Pooling)表达式:
f
(
x
)
=
∑
i
α
(
x
,
x
i
)
y
i
=
∑
i
=
1
n
s
o
f
t
m
a
x
(
−
1
2
(
x
−
x
i
)
2
)
y
i
f(x) = \sum_i{\alpha(x, x_i)y_i} = \sum_{i=1}^{n}softmax(-\frac{1}{2}(x-x_i)^2)y_i
f(x)=i∑α(x,xi)yi=i=1∑nsoftmax(−21(x−xi)2)yi
其中高斯核指数部分视为注意力评分函数(attention scoring function),
α
(
x
,
x
i
)
\alpha(x, x_i)
α(x,xi) 视为注意力权重,最后输出
f
(
x
)
f(x)
f(x) 就是基于这些注意力权重的值的加权和。
1.2 扩展到高维
假设 Query q ∈ R q \boldsymbol{q} \in \mathbb{R}^q q∈Rq, m m m 对 Key-Value ( k 1 , v 1 ) , . . . (\boldsymbol{k_1}, \boldsymbol{v_1}), ... (k1,v1),...,这里 k i ∈ R k , v i ∈ R v \boldsymbol{k_i} \in \mathbb{R}^k, \boldsymbol{v_i} \in \mathbb{R}^v ki∈Rk,vi∈Rv,
注意力汇聚层(Pooling):
最关键的就是如何设计
a
(
q
,
k
i
)
a(\boldsymbol{q}, \boldsymbol{k_i})
a(q,ki)。
1.3 加性注意力 Additive Attention
可学习参数:
等价于将 Key 和 Value 合并起来后放入到一个隐藏大小为 h h h,输出大小为 1 1 1 的单隐藏层 MLP。(Query,Key 和 Value 的长度可以不一样)
1.4 缩放点积注意力(Scaling Dot Product Attention)
如果 Query 和 Key 的长度相同,
q
,
k
i
∈
R
d
\boldsymbol{q}, \boldsymbol{k_i} \in \mathbb{R}^d
q,ki∈Rd,那么可以得到:
a
(
q
,
k
)
=
q
⊤
k
/
d
a(\mathbf q, \mathbf k) = \mathbf{q}^\top \mathbf{k} /\sqrt{d}
a(q,k)=q⊤k/d
在实践中,我们通常从小批量的角度来考虑提高效率, 例如基于
n
n
n 个查询和
m
m
m 个键-值对计算注意力, 其中查询和键的长度为
d
d
d,值的长度为
v
v
v。 查询
Q
∈
R
n
×
d
\mathbf Q\in\mathbb R^{n\times d}
Q∈Rn×d、 键
K
∈
R
m
×
d
\mathbf K\in\mathbb R^{m\times d}
K∈Rm×d 和 值
V
∈
R
m
×
v
\mathbf V\in\mathbb R^{m\times v}
V∈Rm×v 的缩放点积注意力是:
f
=
s
o
f
t
m
a
x
(
Q
K
⊤
d
)
V
∈
R
n
×
v
f = \mathrm{softmax}\left(\frac{\mathbf Q \mathbf K^\top }{\sqrt{d}}\right) \mathbf V \in \mathbb{R}^{n\times v}
f=softmax(dQK⊤)V∈Rn×v
1.5 总结
- 注意力分数是 Query 和 Key 的相似度,注意力权重是分数的 s o f t m a x softmax softmax 结果。
- 两种常见的分数计算:
- 将 Query 和 Key 合并起来进入一个单输出单隐藏层的 MLP
- 直接将 Query 和 Key 做内积
2. 代码
2.0 导包
!pip install -U d2l
import math
import torch
from torch import nn
from d2l import torch as d2l
2.1 遮蔽 softmax 操作
s o f t m a x softmax softmax 操作用于输出一个概率分布作为注意力权重,在某些情况下,并非所有的值都应该被纳入到注意力汇聚中。
为了高效处理小批量数据集, 某些文本序列被填充了没有意义的特殊词元。 为了仅将有意义的词元作为值来获取注意力汇聚, 我们可以指定一个有效序列长度(即词元的个数), 以便在计算 s o f t m a x softmax softmax 时过滤掉超出指定范围的位置。
任何超出有效长度的位置都被掩蔽并置为 0 0 0:
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)
该函数的效果:
masked_softmax(torch.rand(2, 2, 4) , torch.tensor([2, 3]))
torch.rand(2, 2, 4)
表示批量大小为
2
2
2 的样本,
torch.tensor([2, 3])
表示第
1
1
1 个样本有效序列为前
2
2
2 个,第
2
2
2 个样本有效序列为前
3
3
3 个。
masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]]))
还可以定义每个样本的每个元素的有效长度。
2.2 加性注意力(Additive Attention)
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)
# 在维度扩展后,
# 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)
用一个小例子来演示上面的 AdditiveAttention
类, 其中查询、键和值的形状为(批量大小,步数或词元序列长度,特征大小), 实际输出为
(
2
,
1
,
20
)
(2,1,20)
(2,1,20)、
(
2
,
10
,
2
)
(2,10,2)
(2,10,2) 和
(
2
,
10
,
4
)
(2,10,4)
(2,10,4)。 注意力汇聚输出的形状为(批量大小,查询的步数,值的维度):
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')
2.3 缩放点积注意力(Scaling Dot Product Attention)
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')
3. Q&A
Q:mask_softmax 是啥意思?
A:有时候一个句子不够长,假如说一个句子 4 个单词,输入格式要求 10 个单词,那么需要填充 6 个无意义的词,然后用 mask_softmax 告诉 Query 计算的时候不需要考虑后 6 个词。