python将自注意力得分进行缩放

在自注意力机制中，我们需要对每个query和key计算得分，得分需要进行缩放，这是为了控制得分的大小，避免softmax函数的输入值过大或过小，导致梯度消失或梯度爆炸的问题。

具体来说，我们需要使用一个缩放因子，它是一个常数，通常为$\sqrt{d_k}$，其中$d_k$表示key的维度。对于每个query和key的点积结果，我们将它除以这个缩放因子，即：

$$
\text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
$$

这里的除以$\sqrt{d_k}$就是对得分进行缩放。在实现中，我们可以使用如下代码进行缩放：

import torch.nn.functional as F

class MultiheadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(MultiheadAttention, self).__init__()

        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        bs = q.size(0)

        # 线性变换
        q = self.q_linear(q).view(bs, -1, self.n_heads, self.d_k).transpose(1,2)
        k = self.k_linear(k).view(bs, -1, self.n_heads, self.d_k).transpose(1,2)
        v = self.v_linear(v).view(bs, -1, self.n_heads, self.d_k).transpose(1,2)

        # 计算得分
        scores = F.softmax(q.matmul(k.transpose(-2,-1)) / np.sqrt(self.d_k), dim=-1)

        # 对得分进行缩放
        scores = scores.masked_fill(mask == 0, -1e9)

        # 加权求和
        output = scores.matmul(v).transpose(1,2).contiguous().view(bs, -1, self.n_heads*self.d_k)

        # 线性变换得到最终结果
        output = self.out_linear(output)

        return output