一、Self-Attention原理

Self-Attention（自注意力）机制是Transformer架构中的核心组件，主要用于捕捉序列中不同位置元素之间的依赖关系。其核心思想是通过计算序列中每个元素与其他元素之间的相关性（注意力权重），然后根据这些权重对序列信息进行加权聚合。计算过程主要分为以下几步：

1. 输入表示：
   对于输入序列$X$中的每个元素，通过三个不同的权重矩阵$W^Q$、$W^K$、$W^V$，分别生成：

   查询向量（Query）：$Q=X\cdot W^Q$
   键向量（Key）：$K=X\cdot W^K$
   值向量（Value）：$V=X\cdot W^V$
   其中，Q、K、V 承担不同的角色。
   Q (Query): 查询向量，表示当前需要关注的内容。例如在机器翻译中，解码器当前位置的隐状态作为 Query，用于查询源语言的相关信息。
   K (Key): 键向量，表示待匹配的索引。Key 与 Value 关联，用于计算与 Query 的相似度。
   V (Value): 值向量，存储实际需要提取的信息。相似度计算后，Value 会按权重聚合生成输出。

2. 计算注意力分数：
   通过计算$Q$与$K$的点积，得到元素间的相似度分数： $$\text{Scores}=Q\cdot K^{\top }$$ 为了稳定梯度，分数会除以缩放因子$\sqrt{d_k}$（$d_k$是$K$的维度）。

3. 生成注意力权重：
   对分数应用Softmax函数，得到归一化的注意力权重： $$\text{AttentionWeight}=\text{Softmax}\left(\frac{Q\cdot K^{\top }}{\sqrt{d_k}}\right)$$

4. 加权聚合输出：
   用注意力权重对$V$加权求和，得到最终输出： $$\text{Output}=\text{AttentionWeight}\cdot V$$

综上所述，Self-Attention的完整公式如下：
$$\text{Attention(Q,K,V)}=\text{Softmax}\left(\frac{Q\cdot K^{\top }}{\sqrt{d_k}}\right)\cdot V$$
核心作用：
Self-Attention 让模型能够动态关注输入序列的不同部分（例如句子中与当前词相关的其他词），从而更好地理解上下文依赖关系。

二、Self-Attention代码实现

1.单头注意力

单头注意力的代码实现

代码如下（示例）：

class Attention(nn.Module):
    def __init__(self, d_model, head_size, context_length, dropout=0.1):
    '''
    d_model为输入序列的语义维度
    head_size为一个注意力头在语义维度占据的大小，head_size = d_model/num_heads
    context_len为输入序列的长度
    '''
        super().__init__()
        self.head_size = head_size
        self.Wq = nn.Linear(d_model, head_size, bias=False)
        self.Wk = nn.Linear(d_model, head_size, bias=False)
        self.Wv = nn.Linear(d_model, head_size, bias=False)
        self.register_buffer('mask', torch.tril(torch.ones(context_length, context_length)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        q = self.Wq(x)
        k = self.Wk(x)
        v = self.Wv(x)
        weights = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_size)
        weights = weights.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
        weights = F.softmax(weights, dim=-1)

        return weights @ v

2.多头注意力

多头注意力的核心思想：使用多组（$h$ 个头）不同的查询 $Q$、键 $K$、值 $V$ 投影（线性变换），让模型能够并行地从不同的表示子空间（子空间维度 $d_k$, $d_v$, $d_{model}/h$）学习信息。每个头学习不同的关注模式。

多头注意力的计算流程

假设：
输入维度：$d_{model}$（例如 512）
头数：$h$（例如 8）
每个头的维度：$d_k=d_v=d_{model}/h$（例如 512 / 8 = 64）
步骤：

1. 线性投影（生成 h 组 Q, K, V）：
   对原始的查询 $Q$、键 $K$、值 $V$（维度均为 $d_{model}$）分别应用 $h$ 组不同的线性变换（权重矩阵）。
   得到 $h$ 组投影后的查询 $Q_i$、键 $K_i$、值 $V_i$，每组维度为 $d_k$, $d_k$, $d_v$（通常 $d_k=d_v=d_{model}/h$）。
   $$Q_i=Q{W}_i^Q,K_i=K{W}_i^K,V_i=V{W}_i^V\text{for }i=1,...,h$$ 其中 $W_i^Q,W_i^K,W_i^V$ 是可学习的投影矩阵。

2. 并行计算缩放点积注意力：
   对每一组投影后的 $Q_i,K_i,V_i$，独立计算缩放点积注意力： $${\text{head}}_i=\text{Attention}(Q_i,K_i,V_i)=\text{softmax}(\frac{Q_i{K}_i^T}{\sqrt{d_k}})V_i$$

3. 拼接多头输出：
   将 $h$ 个注意力头计算的结果 ${\text{head}}_1,{\text{head}}_2,...,\text{head}h$（每个维度为 $d_v$）拼接起来，得到一个维度为 $h\times d_v=dmodel$ 的向量。 $$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,{\text{head}}_2,...,{\text{head}}_h)$$

4. 最终线性投影（可选，但常用）：
   将拼接后的结果通过另一个线性变换 $W^O$ 投影到最终的输出维度（通常是 $d_{model}$）。 $$\text{Output}=\text{MultiHead}(Q,K,V)W^O$$ 其中 $W^O$ 是维度为 $d_{model}\times d_{model}$ 的可学习权重矩阵。

多头注意力的代码实现

多头注意力的代码实现方法十分简单，只需使用到上面单头注意力的代码即可。

代码如下：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, head_size, context_length, dropout=0.1):
        super().__init__()
        self.heads = nn.ModuleList([
            Attention(d_model, head_size, context_length, dropout)
            for _ in range(num_heads)
        ])
        self.projection_layer = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        head_outputs = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.dropout(self.projection_layer(head_outputs))

总结

以上就是今天要讲的内容，本文仅仅简单介绍了Self-Attention的原理和简单实现，Self-Attention是大模型的核心机制需要认真学习。