目录
第七步：加mask并得到背景|context|输出的底板信息
原文首发在公众号 AI天天用，欢迎关注，共同进步。
应朋友所托，复现Informer代码。没想到后来还吃到瓜了。如果您也有需要复现的文章（ns子刊，顶刊顶会为主），麻烦您关注公众号留言。
Informer作者提出的ProbSparseAttention被认为是Informer的核心创新点。今天我们一起来揭开其神秘面纱。
获取完整代码，请关注公众号。下次内容有瓜吃哇，不要错过！
原始Attention：卷王
原始Attention在2017年由Vaswani等人提出 （Attention is all you need）：
Attention(Q,K,V)=softmax(QKTdk)V
Informer的作者们发现， QKT 形成的自注意力呈现长尾特性（图1）。
图1：自注意力的长尾特性[1]
卷王卷了半天，原来有效的内容很少啊！
Q和K可以分解为
Q=[q1,q2,⋯,qm],K=[k1,k2,⋯,kn]
换言之，少部分Q和K的内积贡献了大部分attention的值，其他的都可以被忽略。
因此，作者希望用少量的Q来得到注意力attention。
Informer: 拒绝躺平
然而，手心手背都是肉，选哪些 qi ，舍弃哪些 qi 呢？真让人头大。
Infomer的作者灵机一动，来了information。
如果把 qi 与所有 k∈K 的内积定义为一个分布 p(kj|qi) 。
如果这个分布躺平了，变成了均匀分布， 也就是说不管 qi 是什么， p(kj|qi)=1LK 。 导致的结果是，自注意力变成的V的简单求和（某种程度上的均值）。
Informer的作者们认为一定不能躺平啊，分布 p(kj|qi) 要离均匀分布越远越好。
怎么计算两个分布之间的距离呢？这个话题就太大了，但一般而言，用KL散度来度量就行。 经过简单，对于 qi ， 可以度量其与均与分布间的距离（sparsity)
M(qi,K)=ln∑j=1LKeqikjTd−1LK∑j=1LKqikjTd
这样做，就简单，只需要把每个 qi 和 K 计算其sparsity，然后选择进行排序，选择 M(qi,K) 大的几个就行。
选择出来 qi 以后，再用选出来的少量 qi 与K和V做运算，得到ProbSparseAttention。简直Amazing！
同学，读到这里，是不是有点问号？
每个 qi 都和 K 运算（这不就是原始的attention吗？），然后再选择，这比原始attention这个卷王还卷啊！ 上边的计算方式，相当于计算了一遍原始attention，然后选出来少量 qi ，再计算一遍probsparse attention。
换言之，拒绝躺平走向了另一个极端，比卷王还要卷啊。
Informer龙场悟道： K 也要选一选
如何避免比卷王还卷呢？ 既然Q能选，为什么我老K不能选呢？选起来。 Informer通过随机选取少量 kj ，然后与 qi 进行运算，选出来 qi ，再与 K 算自注意力。这样，即避免了躺平，也避免了比卷王还卷。
到底选多少个：少量；些许；一些
要选多少 qi ，多少 kj 呢？ Informer的原文指出，选 qi ，可以通过$u = c \ln m$来选。 选$k_j$可以通过$U=m \ln n$来选。$c$是一个作者指定的值（默认为5），$u$和$U$是要选的$q_i$和$k_j$。$m$和$n$分别是Q和K的长度。
细心的同学可能发现了，Q和K的长度一般比3长，也就是$\ln n \geq 1$， $m\ln n \geq m$。一般$m=n$，也就是说，选择的$k_j$要比K的长度还要长。Amazing！
说好的随机选取少量$k_j$呢？U，你变了。
其实，在具体实现过程中，作者们是根据$U = c \ln n$来选择$k_j$的，避免了尬尴。
OK，以上便是Informer的核心创新工作了。接下来复现。
AI天天用（不负任何责任）乱评：聪明的同学，你有没有更好的方法来选取少量$q_i$呢？
代码复现：talk is cheap
Attention: 原代码中是FullAttention
几点说明：
- 源代码有一个
output_attention的参数，这里的复现采用self.attn=None，attn随着模型走，不需要返回或者用参数决定是否返回。 - use_mask和factor看似多余，实际上是作者为了和提出的ProbAttention模块保持一致。
- 输入尺寸(N, L, D)，输出尺寸(N, L, D)
class VanillaAttention(nn.Module):
"""
vanilla attention
"""
def __init__(self, use_mask=False, dropout=0.1, factor=5):
super().__init__()
self.use_mask = use_mask
self.dropout = nn.Dropout(dropout)
self.factor = factor
self.attn = None
# self.proj_q = nn.Linear(d_model, d_model)
# self.proj_k = nn.Linear(d_model, d_model)
# self.proj_v = nn.Linear(d_model, d_model)
# self.num_heads = num_heads
def forward(self, q, k, v, mask=None):
"""
q, k, v -- (N, H, L, D)
"""
# q, k, v = self.proj_q(q), self.proj_k(k), self.proj_v(v)
# q, k, v = map(lambda item: rearrange(item, "N L (H d) -> N H L d", H=self.num_heads), (q, k, v))
attn = torch.einsum("nhid, nhjd -> nhij", q, k) * (k.shape[-1] ** -0.5)
if self.use_mask:
if mask is None:
shape = (attn.shape[0], 1, attn.shape[2], attn.shape[3])
mask = torch.triu(torch.ones(shape, dtype=torch.bool), diagonal=1).to(attn.device)
attn.masked_fill_(mask, -np.inf)
attn = torch.softmax(attn, dim=-1)
self.attn = self.dropout(attn)
out = torch.einsum("nhij, nhjd -> nhid", attn, v)
# out = rearrange(out, "n h i d -> n i (h d)")
return outProbAttention复现
图2： ProbAttention实现步骤[1]
根据图2复现ProbAttention。
第一步： 计算$u$和$U$
num_q, num_k = [int(self.factor * np.ceil(np.log(length))) for length in [q_len, k_len]]
num_q, num_k = np.minimum(num_q, q_len), np.minimum(num_k, k_len)num_q表示$u$， num_k表示$U$。
第二步：随机选取少量K
k_expanded = k.unsqueeze(-3).expand(-1, -1, q_len, -1, -1)
random_index = torch.randint(k_len, size=(q_len, num_k))
k_sampled = k_expanded[:,:,torch.arange(q_len).unsqueeze(1), random_index, :]有同学可能会问，为什么不直接随机算K呢，还要expand？
图3：随机选择
如图3所示，对于每个 qi ，需要选取不同的 kj 来验证其sparsity。 q1 可能随机选 k1 和 k3 ， q2 可能随机选 k2 和 k4 。 因此需要expand操作，确保每个 qi 随机算的 kj 不一样。 这里用到了python的advance indexing操作。
第三步：计算pre-attention
pre_attn = torch.einsum("bhid,bhijd->bhij", q, k_sampled)torch.einsum的出现，极大地减少了脑细胞的死亡数量。
第四步：得到用来选少量 qi 的M
measure = pre_attn.max(-1)[0] - pre_attn.sum(-1) / k_len其实，这里不能除以 K 的长度 k_len ，而是num_k。Anyway，就是这么随机。
第五步：选少量 qi
q_selected_index = measure.topk(num_q, sorted=False)[1]
q_selected = q.gather(-2, q_selected_index.unsqueeze(-1).expand(-1, -1, -1, q.shape[-1]))终于选到你！同学，你放弃了吗？
第六步：计算attention
attn = torch.einsum("bhid,bhjd->bhij", q_selected, k) * (k.size(-1) ** -0.5)第七步：加mask并得到背景|context|输出的底板信息
if self.use_mask:
assert q_len == v_len
mask = ProbMask(v.shape[0], v.shape[1], q_len, q_selected_index, attn)
attn.masked_fill_(mask.mask, -np.inf)
# set the uniform information as the background (context)
background = v.cumsum(dim=-2) # Step 7
else:
v_mean = v.mean(dim=-2) # Step 7
background = v_mean.unsqueeze(-2).expand(-1, -1, q_len, v_mean.shape[-1]).clone()
attn = torch.softmax(attn, dim=-1)
self.attn = self.dropout(attn)有同学问，在计算background的过程中，为什么一个用了cumsum，令一个用了mean？在github issue中，作者回复了，后边跟的有normalization layer，所以用mean或者cumsum结果应该差不多。
AI天天用（不负任何责任）乱评：cumsum可以保持维度啊！mean那个后边维度进行了扩充和expand，用起来很合理。
第八步：得到自注意力结果
out = torch.einsum("bhij,bhjd->bhid", self.attn, v)
q_selected_index = q_selected_index.unsqueeze(-1)
# recovery the shape of out
out_scatter_index = q_selected_index.expand(-1, -1, -1, out.shape[-1])
out = background.scatter(2, out_scatter_index, out)
attn_init = torch.ones(v.shape[0], v.shape[1], v_len, v_len, dtype=attn.dtype, device=attn.device) / v_len
attn_scatter_index = q_selected_index.expand(-1, -1, -1, attn_init.shape[-1])
self.attn = attn_init.scatter(2, attn_scatter_index, self.attn)通过scatter函数，将得到的out恢复为卷王attention的尺寸。类似地，将attn也恢复为正常attn的尺寸。
Attention 复现
有了VanillaAttention和ProbAttention，就可以得到Multi-head (Prob) Self-Attenetion啦！
class AttentionBlock(nn.Module):
def __init__(self, attn_type, d_model, num_heads, use_mask=False, dropout=0.1, factor=5):
super().__init__()
self.proj_q = nn.Linear(d_model, d_model)
self.proj_k = nn.Linear(d_model, d_model)
self.proj_v = nn.Linear(d_model, d_model)
self.num_heads = num_heads
assert attn_type in ["vanilla", "prob"]
if attn_type == "vanilla":
self.attention = VanillaAttention(use_mask=use_mask, dropout=dropout, factor=factor)
else:
self.attention = ProbAttention(use_mask=use_mask, dropout=dropout, factor=factor)
self.norm = nn.LayerNorm(d_model)
def forward(self, q, k, v, mask=None):
q_, k_, v_ = self.proj_q(q), self.proj_k(k), self.proj_v(v)
q_, k_, v_ = map(lambda item: rearrange(item, "N L (H d) -> N H L d", H=self.num_heads), (q_, k_, v_))
out = self.attention(q_, k_, v_, mask)
out = rearrange(out, "N H L D -> N L (H D)")
# for self-attention, q = k = v = x
# for cross-attention, q = x
# In "encoder-decoder attention" layers, the queries come from the previous decoder layer,
# and the memory keys and values come from the output of the encoder.
# Attention is all you need, Vaswani et al., 2017
return self.norm(out + q)通过一个attn_type的参数，就能控制要使用的attention的类型。 同时，保留了attention中的残差连接（Informer的attention AttentionLayer中无残差连接）。
Any other solutions?
同学，您有其他方法吗？