注意力机制

参考文章:

1. Transformer学习笔记二：Self-Attention（自注意力机制） - 知乎
2. (26 封私信 / 80 条消息) transformer中的attention为什么scaled? - 知乎

1. 什么是注意力机制

注意力机制的核心思想是模仿人类视觉的注意力行为：人类观察复杂场景时，会有选择性地关注其中的关键部分，而忽略其他不重要的区域。同样地，在深度学习中，注意力机制能够让模型根据输入数据动态地分配不同的注意力权重，从而更有效地捕捉关键特征。

2. Scaled Dot-Product Attention

(1) 公式:

缩放点积注意力是 Transformer 的基本单元，用于计算查询 Q、键 K 和值 V 之间的注意力权重并生成输出。

首先给出公式:
$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q\cdot K^T}{\sqrt{d_k}}\right)V$$

• 输入:
  • 向量 Q(query): "我"应该关注什么
  • 向量 K(key): "我"位置的信息特征
  • 向量 V(value): "我"的真实数据
• 输出: 上下文向量, 反映了模型对输入序列的理解，即哪些部分的信息对于回答当前 Query 更为重要。

(2) 举例:

对于图中的 token $a^2$:

1. 首先 $a^2$ 通过运算得到 Q, K, V 向量(后面介绍怎么进行运算)
2. 用 $a^2$ 的Q分别与其他向量的 K 进行点乘: $Q\cdot K^T$, 得到 $a^2$ 与其他 token 之间的相似度(积值越大，表示两者越相似)
   • 两个向量的点乘表示两个向量的相似度
3. 将得到的相似度进行缩放, 并应用 softmax 得到 attention score : $\mathrm{softmax}(\frac{Q\cdot K^T}{\sqrt{d_k}})$
   • $d_k$是 K 的维度, 随着维度的增大，点积的方差会变得越来越大, 会导致 softmax 的梯度值小, 无法进行有效的梯度更新
   • $Q\cdot K^T$ 除以 $\sqrt{d_k}$, 可以将其归一化成均值为 0, 方差为 1 的向量, 防止方差变大
4. 再将得到的 attention score 乘以每个 token 对应的 V, 将得到的结果相加(即加权求和), 便可得到$a^2$通过 attention 模型后所产生的结果

(3) 如何得到 Q, K, V

在 Transformer 模型中，Query (Q)、Key (K) 和 Value (V) 是从输入序列（如嵌入层输出或上一层的输出）通过线性变换得到的

• 图中 $W^Q,W^K,W^V$ 是三个可训练的参数矩阵
• 其中$W^Q,W^K$维度都是 $k_{dim}$, $W^V$维度是$v_{dim}$. 两维度可以不一定相等

(4) 为什么要进行 scaling

上面提到$Q\cdot K^T$ 除以 $\sqrt{d_k}$, 可以将其归一化成均值为 0, 方差为 1 的向量, 防止方差变大, 进而防止 softmax 的梯度值小, 无法进行有效的梯度更新

1. 为什么$d_k$变大, 方差会变大?

假设 Q 和 K 的向量长度为 $d_k$， 服从独立同分布，均值为0， 方差为 1的正态分布。则 Q 和 K 的点积的方差为：
$$\begin{array}{rl}\mathrm{VAR}(Q\cdot K^T)& =\mathrm{VAR}\left[\sum _{i=1}^{d_k}{Q}_i\times K_i\right]\\ & =\sum _{i=1}^{d_k}\mathrm{VAR}\left[Q_i\times K_i\right]\\ & =\sum _{i=1}^{d_k}\mathrm{VAR}\left[Q_i\right]\times \left[K_i\right]\\ & =\sum _{i=1}^{d_k}1\\ & =d_k\end{array}$$
所以, 当 d 变大时，点积的方差变大

2. 为什么方差变大后, softmax 梯度值会变小?

首先有输入向量一个向量$\mathbf{z}=[z_1,z_2,\dots ,z_n]$，softmax 函数的输出是一个向量$y=\mathbf{y}=[y_1,y_2,\dots ,y_n]$

softmax 函数为:
$$y_i=\mathrm{softmax}(z_i)=\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}$$
对其求导, 即有:
$$\frac{\partial y_i}{\partial z_k}=\{\begin{array}{c}{y}_i(1-y_i),i!=k\\ -y_i{y}_k,i=k\end{array}$$
当方差变大, 即 $z_i$ 相对于其他的 $z$ 变得更大或更小

• 当$z_i$相对其它 $z$ 更大, $y_i$ 趋近于1, $y_k$ 趋近于0, 此时以上的两个结果都趋近于0
• 当$z_i$相对其它 $z$ 更小, $y_i$ 趋近于0, $y_k$ 趋近于1, 此时以上的两个结果都趋近于0

即方差变大, 使得 softmax 梯度变小, 无法进行有效的反向传播参数更新

3. Masked Attention

在语言模型的训练过程中，我们希望模型能够预测下一个词，但它只能依赖于当前词和之前的词。即看不到之后的词, 如果没有这种限制，模型在生成每个词时就能够“作弊”，提前看到答案，从而影响训练的正确性

实现原理:

• ‘+’ 号并不是表示相加，只是表示提供了位置覆盖的信息
• MASK 矩阵中标1的地方，就是需要遮蔽的地方，我们把原来的值替换为一个很小的值（比如-1e09），而在标0的地方，我们保留原始的值
• MASK 矩阵的每一行表示对应位置的token。例如在第一行第一个位置是0，其余位置是1，这表示第一个token在attention时，只看到它自己，它右边的tokens是看不到的。以此类推。

4. Multihead Attention

多头注意力就是通过并行地计算多个注意力机制（称为“头”），以获取更丰富的表示。每个头可以关注不同的特征或位置关系。

多头注意力的头可以是上文中的 Scaled Dot-Product Attention, 也可以是 Masked Attention

多头注意力:

假设多头注意力模块有3个头, 每个头的 Q (Query)、K (Key)、V (Value) 矩阵是通过输入 X 乘以全局的 $W^Q,W^K,W^V$，然后在维度上均分，分配给各个头得到:

给定输入 X 的形状为 $(N,T,d_{\text{model}})$，其中：

• N 是批量大小
• T 是序列长度
• $d_{\text{model}}$ 是输入特征维度

为了得到 Q、K、V，使用全局的线性变换矩阵 $W^Q,W^K,W^V$：
$$Q=X{W}^Q,K=X{W}^K,V=X{W}^V$$

• $W^Q,W^K,W^V$ 的形状为 $(d_{\text{model}},d_{\text{model}})$

• Q、K、V 的形状为 $(N,T,d_{\text{model}})$

将 $d_{\text{model}}$ 维度分割为 h 个头，每个头的维度为$d_k=\frac{d_{\text{model}}}{h}$：

• Q 被 reshape 为 $(N,T,h,d_k)$
• K 被 reshape 为 $(N,T,h,d_k)$
• V 被 reshape 为 $(N,T,h,d_k)$

然后将头的维度提前, 得到 Q, K, V 的维度为: $(N,h,T,d_k)$

• $N$：批量大小；
• $h$：头的数量；
• $T$：序列长度；
• $d_k$：每个头的特征维度

然后每个头单独进行运算, 最后将多头输出拼接起来就可以得到最终输出

五. 代码

import math
import torch.nn as nn
import torch

class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0:

            # hidden_size 需要是 num_attention_heads的整数倍
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention heads ({config.num_attention_heads})"
            )
        # 头的数量
        self.num_attention_heads = config.num_attention_heads
        # 头的尺寸
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        # 定义query、key、value的线性变换
        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        """
        将输入张量 `x` 的最后一个维度重新调整大小，然后置换维度，以准备张量用于多头注意力计算。
        参数:
            x (torch.Tensor): 形状为 (batch_size, seq_length, hidden_size) 
        返回:
            torch.Tensor: 形状为 (batch_size, num_attention_heads, seq_length, attention_head_size) 
        """

        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
    ):
        # 生成QKV矩阵
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        # 调整QKV矩阵的形状
        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # 计算注意力分数
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            attention_scores = attention_scores + attention_mask

        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        attention_probs = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs, value_layer)

        # 调整 context_layer 的张量形状，以便将多头注意力的输出合并回原始的隐藏状态维度
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        return outputs