注意力权重，评分函数

注意力评分函数（Attention Scoring Function）

注意力机制的核心在于根据查询（query）和键（key）之间的关系，计算出每个值（value）对应的注意力权重，然后使用这些权重对值进行加权求和，得到最终的输出。

关键概念：

1. 注意力权重：

   • 通过注意力评分函数 $ a(\mathbf{q}, \mathbf{k}_i) $ 计算查询 \(\mathbf{q}\) 和键 \(\mathbf{k}_i\) 的匹配程度。
   • 使用 softmax 函数将评分转化为概率分布（即注意力权重）：
     \[\alpha(\mathbf{q}, \mathbf{k}_i) = \mathrm{softmax}(a(\mathbf{q}, \mathbf{k}_i)) = \frac{\exp(a(\mathbf{q}, \mathbf{k}_i))}{\sum_{j=1}^m \exp(a(\mathbf{q}, \mathbf{k}_j))} \]

2. 注意力汇聚函数：

   • 最终输出是基于注意力权重的值的加权平均：
     \[f(\mathbf{q}, (\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_m, \mathbf{v}_m)) = \sum_{i=1}^m \alpha(\mathbf{q}, \mathbf{k}_i) \mathbf{v}_i \]

3. 注意力评分函数的作用：

   • 不同的评分函数 $ a(\mathbf{q}, \mathbf{k}_i) $ 会生成不同的注意力机制。
   • 常见的评分函数包括：
     • 点积（Dot Product）：适用于向量维度一致的情况。
     • 加性注意力（Additive Attention）：通过一个前馈神经网络计算得分，适合不同维度的向量。

4. 宏观理解：

   • 注意力机制可以看作是从一组“键-值”对中，根据当前查询选择最相关的信息，并进行加权汇总。
   • 这种机制在自然语言处理、图像识别等领域有广泛应用。

---

掩码softmax实现

### X的形状[batch_size, features]
def sequence_mask(X, valid_len, value=0):
    """Mask irrelevant entries in sequences.
"""
    maxlen = X.size(1) ## maxlen = features  
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]

    ### torch.arange((maxlen), dtype=torch.float32, device=X.device)  => 这个range的形状是[features], valid_len形状是[batch_size, 1] , 二者做小于运算会先broadcast到[batchsize, features]维度 , 并得到一个此维度的 bool矩阵
    X[~mask] = value 
    return X



## 这个方法是用在算完x的注意力得分之后的，原理是在计算注意力权重之前，把无效的padding元素的得分清除
def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    ## x.shape = [time_step, batch_size, features]
    ## valid_len.shape = [time_step] 或者  [time_step, batch_size]
    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]) ## 把valid_lens reshape成 [time_steps*batch_size], 
        else:
            valid_lens = valid_lens.reshape(-1)
            ## 把valid_lens reshape成 [time_steps*batch_size], 


        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)## 把X reshape成 [time_step*batch_size, features]把valid_lens reshape成 [time_steps*batch_size], 运算squence_mask，把相当于把三维的concat成二维的做一次运算，而不是用for循环做time_step次运算，更能利用gpu的矩阵乘法的优势
        return nn.functional.softmax(X.reshape(shape), dim=-1)

加性注意力

\[a(\mathbf q, \mathbf k) = \mathbf w_v^\top \text{tanh}(\mathbf W_q\mathbf q + \mathbf W_k \mathbf k) \in \mathbb{R}, \]

class AdditiveAttention(nn.Module):
    """加性注意力"""
    ## q.shape = [batch_size, q.time_step 或者 q的个数, q_featrues]
    ## k.shape = [batch_size, k.time_step 或者 k的个数, k_featrues]
    ## v.shape = [batch_size, k.time_step 或者 k的个数 v_featrues] k和v是一一对应的，k的个数肯定等于v的个数
    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)
        ## 当tensor维度大于2时，比如k.shape = 3, linear层只会作用在后两个维度，即把featrues维map到num_hiddens维
        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)  ## 当tensor维度大于2时，比如k.shape = 3, linear层只会作用在后两个维度，即把featrues维map到num_hiddens维
   
        features = queries.unsqueeze(2) + keys.unsqueeze(1) 
        # 用unsqueeze进行维度扩展后，（q在第3个维度扩展，k在第二个维度扩展）
        # queries的形状：(batch_size，查询的个数，1，num_hidden)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        features = torch.tanh(features)
        ## 现在features的维度为[batch_size, num_q, num_kv, num_hiddens]
        
        scores = self.w_v(features).squeeze(-1)
        # self.w_v仅有一个输出，因此从形状中移除最后那个维度。
        # scores的形状：(batch_size，num_q, num_kv)
        self.attention_weights = masked_softmax(scores, valid_lens)##   # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数), 这里scores中的 num_kv相当于features, mask会遮盖住部分features，即遮盖掉部分不发挥作用的k，即k中的那些padding部分所对应的元素
        ## attention_weights的形状和输入一样，即和scores一样，只是加了掩码后的scores, 即shape为(batch_size，num_q, num_kv)
        ## v.shape = (batch_size, num_kv, v_featrues)
        
        return torch.bmm(self.dropout(self.attention_weights), values)# values的形状：(batch_size，“键－值”对的个数，值的维度)
        ## 则最后bmm后得到的res.shape = (batch_size, num_q, v_featrues)，正是我们想要的对应每个q 的 v的加权和

点积注意力

缩放点积注意力（scaled dot-product attention）评分函数为：

\[a(\mathbf q, \mathbf k) = \mathbf{q}^\top \mathbf{k} /\sqrt{d}. \]

在实践中，我们通常从小批量的角度来考虑提高效率，
例如基于\(n\)个查询和\(m\)个键－值对计算注意力，
其中查询和键的长度为\(d\)，值的长度为\(v\)。
查询\(\mathbf Q\in\mathbb R^{n\times d}\)、
键\(\mathbf K\in\mathbb R^{m\times d}\)和
值\(\mathbf V\in\mathbb R^{m\times v}\)的缩放点积注意力是：

\[\mathrm{softmax}\left(\frac{\mathbf Q \mathbf K^\top }{\sqrt{d}}\right) \mathbf V \in \mathbb{R}^{n\times v}. \]

#@save
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)