重新审视MHA与Transformer_手写mha

本文将基于PyTorch源码重新审视MultiheadAttention与Transformer。事实上，早在一年前博主就已经分别介绍了两者：各种注意力机制的PyTorch实现、从零开始手写一个Transformer，但当时的实现大部分是基于d2l教程的，这次将基于PyTorch源码重新实现一遍。

1. MultiheadAttention

1.1 思路

回顾多头注意力，其公式如下：

$$\begin{gathered} \text{MHA}(Q,K,V)=\text{Concat}({\text{head}}_1,\cdots ,{\text{head}}_h)W^O \\ {\text{head}}_i=\text{Attn}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V) \end{gathered}$$

其中 $W_i^Q\in {\mathbb{R}}^{d_{model}\times d_k}$，$W_i^K\in {\mathbb{R}}^{d_{model}\times d_k}$，$W_i^V\in {\mathbb{R}}^{d_{model}\times d_v}$，$W^O\in {\mathbb{R}}^{h{d}_v\times d_{model}}$，且 $d_k=d_v=d_{model}/h$。

如果记 $d_{head}=d_{model}/h$，则 $W_i^Q,W_i^K,W_i^V$ 的形状均为 $(d_{model},d_{head})$，$W^O$ 的形状为 $(d_{model},d_{model})$。

先不考虑batch和mask的情形，在只有一个头的情况下（$h=1$），MHA的计算方式为

class MHA(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.w_q = nn.Parameter(torch.empty(d_model, d_model))
        self.w_k = nn.Parameter(torch.empty(d_model, d_model))
        self.w_v = nn.Parameter(torch.empty(d_model, d_model))
        self.w_o = nn.Parameter(torch.empty(d_model, d_model))

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, query, key, value):
        """
        Args:
            query: (n, d_model)，n是query的个数，m是key-value的个数
            key: (m, d_model)
            value: (m, d_model)
        """
        q = query @ self.w_q
        k = key @ self.w_k
        v = value @ self.w_v

        attn_logits = q @ k.transpose(0, 1) / math.sqrt(q.size(1))  # attn_logits: (n, m)
        attn_probs = F.softmax(attn_logits, dim=-1)
        attn_output = attn_probs @ v  # attn_output: (n, d_model)
        return attn_output, attn_probs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

现在考虑 $h=2$ 的情形，此时一共需要 $3\cdot 2+1=7$ 个参数矩阵

class MHA(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.w_q_1 = nn.Parameter(torch.empty(d_model, d_model // 2))
        self.w_k_1 = nn.Parameter(torch.empty(d_model, d_model // 2))
        self.w_v_1 = nn.Parameter(torch.empty(d_model, d_model // 2))

        self.w_q_2 = nn.Parameter(torch.empty(d_model, d_model // 2))
        self.w_k_2 = nn.Parameter(torch.empty(d_model, d_model // 2))
        self.w_v_2 = nn.Parameter(torch.empty(d_model, d_model // 2))

        self.w_o = nn.Parameter(torch.empty(d_model, d_model))

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, query, key, value):
        """
        Args:
            query: (n, d_model)，n是query的个数，m是key-value的个数
            key: (m, d_model)
            value: (m, d_model)
        """
        q_1 = query @ self.w_q_1
        k_1 = key @ self.w_k_1
        v_1 = value @ self.w_v_1

        q_2 = query @ self.w_q_2
        k_2 = key @ self.w_k_2
        v_2 = value @ self.w_v_2

        attn_logits_1 = q_1 @ k_1.transpose(0, 1) / math.sqrt(q_1.size(1))
        attn_probs_1 = F.softmax(attn_logits_1, dim=-1)
        attn_output_1 = attn_probs_1 @ v_1

        attn_logits_2 = q_2 @ k_2.transpose(0, 1) / math.sqrt(q_2.size(1))
        attn_probs_2 = F.softmax(attn_logits_2, dim=-1)
        attn_output_2 = attn_probs_2 @ v_2

        attn_output = torch.cat([attn_output_1, attn_output_2], dim=-1) @ self.w_o  # attn_output: (n, d_model)
        attn_probs = torch.stack([attn_probs_1, attn_probs_2], dim=0)  # attn_probs: (2, n, m)，其中2是头数

        return attn_output, attn_probs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

可以看到代码量已经增加了不少，如果扩展到 $h$ 个头的情形，则需要 $3h+1$ 个参数矩阵。手动去一个个声明显然不现实，因为 $h$ 是动态变化的，而用for循环创建又略显笨拙，有没有更简便的方法呢？

在上面的代码中，我们用小写 $q$ 来代表查询 $Q$ 经过投影后的结果（$k,v$ 同理），即

$$q_i=Q{W}_i^Q,i=1,2,\cdots ,h$$

其中 $Q$ 的形状为 $(n,d_{model})$，$q_i$ 的形状为 $(n,d_{head})$，且有

$$hea{d}_i=\text{softmax}\left(\frac{q_i{k}_i^T}{\sqrt{d_{head}}}\right)v_i$$

注意到

$$[q_1,q_2,\cdots ,q_h]=Q[W_1^Q,W_2^Q,\cdots ,W_h^Q]\text{\hspace{1em}(1)}$$

如果记 $q\triangleq [q_1,q_2,\cdots ,q_h]$，$W^Q\triangleq [W_1^Q,W_2^Q,\cdots ,W_h^Q]$，则 $W^Q$ 的形状为 $(d_{model},d_{model})$，与 $h$ 无关，$q$ 的形状为 $(n,d_{model})$。这样一来，我们就不需要一个个声明 $W_i^Q$ 了，并且可以一次性存储所有的 $q_i$。

要计算 $hea{d}_1$，我们需要能够从 $q$ 中取出 $q_1$（$k,v$ 同理），所以我们期望 $q$ 的形状是 $(h,n,d_{head})$，从而 $q[1]$ 就是 $q_1$（这里下标从 $1$ 开始）。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/凡人多烦事01/article/detail/351236