Informer复现--模型之Attention

源代码有一个 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 out

ProbAttention复现

图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?

同学，您有其他方法吗？

参考文献

Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting

AI天天用

关注

40
点赞
踩
32

收藏

觉得还不错? 一键收藏
0
评论
Informer复现--模型之Attention

原文首发在公众号 AI天天用，欢迎关注，共同进步。应朋友所托，复现Informer代码。没想到后来还吃到瓜了。如果您也有需要复现的文章（ns子刊，顶刊顶会为主），麻烦您关注公众号留言。Informer作者提出的ProbSparseAttention被认为是Informer的核心创新点。今天我们一起来揭开其神秘面纱。获取完整代码，请关注公众号。下次内容有瓜吃哇，不要错过！
复制链接

扫一扫