🎯 课程主题

自注意力机制的矩阵理解方式——用矩阵运算统一表达 QKV 计算与注意力融合。

📝 核心知识点

1. 输入矩阵 X 的构建

  • 概念说明：将序列所有 token 的词向量堆叠成输入矩阵。
  • 关键细节（以 "我喜欢打篮球" 为例）：
    • $A_1$="我", $A_2$="喜欢", $A_3$="打", $A_4$="篮球"
    • 序列长度 $L = 4$
    • 词嵌入维度 $D$（Transformer 中为 512，演示中设为 5）
    • 输入矩阵 $X \in \mathbb{R}^{L \times D} = \mathbb{R}^{4 \times 512}$

2. QKV 矩阵映射

  • 概念说明：输入矩阵 X 与三个权重矩阵相乘，一次性得到所有 Q、K、V。
  • 关键细节：
    • 权重矩阵维度：$W_Q, W_K, W_V \in \mathbb{R}^{D \times d_k}$
      • $D$ 必须与输入维度匹配（用于矩阵相乘）
      • $d_k$ 是人为设置的 Q/K/V 维度
    • $X \times W_Q = Q$，矩阵 $Q \in \mathbb{R}^{L \times d_k}$，包含 $Q_1, Q_2, Q_3, Q_4$
    • $X \times W_K = K$，矩阵 $K \in \mathbb{R}^{L \times d_k}$，包含 $K_1, K_2, K_3, K_4$
    • $X \times W_V = V$，矩阵 $V \in \mathbb{R}^{L \times d_v}$，包含 $V_1, V_2, V_3, V_4$

3. 注意力得分矩阵计算

  • 概念说明：通过 $Q \times K^T$ 一次性计算所有 token 间的注意力得分。
  • 关键细节：
    • $K$ 需转置为 $K^T \in \mathbb{R}^{d_k \times L}$
    • $Q \times K^T = A$，得分矩阵 $A \in \mathbb{R}^{L \times L}$
    • 矩阵 $A$ 中每个元素 $\alpha_{ij}$ 表示第 $i$ 个 token 与第 $j$ 个 token 的关联程度
    • 经 Softmax 归一化后，每行和为 1
    • 需除以 $\sqrt{d_k}$：防止 Softmax 饱和导致梯度消失，提升训练稳定性与收敛速度

4. 信息融合输出

  • 概念说明：注意力得分矩阵与 V 相乘，得到融合后的输出矩阵。
  • 关键细节：
    • $A' \times V = B$，输出矩阵 $B \in \mathbb{R}^{L \times d_v}$
    • $B$ 中每行 $B_i$ 是第 $i$ 个 token 对所有 token 信息的加权融合
    • 输入 $X$ 维度 $L \times D$，输出 $B$ 维度 $L \times d_v$（通常 $d_v = D$）

🧮 核心公式与推导

QKV 矩阵映射：

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

其中 $X \in \mathbb{R}^{L \times D}$，$W_Q, W_K \in \mathbb{R}^{D \times d_k}$，$W_V \in \mathbb{R}^{D \times d_v}$。

缩放点积注意力：

$$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

分步展开：

  1. 注意力得分：$S = \frac{Q K^T}{\sqrt{d_k}} \in \mathbb{R}^{L \times L}$
  2. 归一化：$A = \text{Softmax}(S)$，每行和为 1
  3. 信息融合：$B = A V \in \mathbb{R}^{L \times d_v}$

维度变化：

  • $Q K^T$：$(L \times d_k) \times (d_k \times L) = L \times L$
  • $A V$：$(L \times L) \times (L \times d_v) = L \times d_v$

物理意义：矩阵形式将所有 token 的注意力计算并行化，一次矩阵乘法完成所有关联度计算。

🏗️ 模型架构与数据流向

维度变化全流程： $$X_{L \times D} \xrightarrow{\times W_Q} Q_{L \times d_k} \xrightarrow{\times K^T} S_{L \times L} \xrightarrow{\text{Softmax}} A_{L \times L} \xrightarrow{\times V} B_{L \times d_v}$$

💻 代码实战

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SelfAttention(nn.Module):
    """自注意力机制的矩阵实现"""
    def __init__(self, d_model, d_k):
        """
        d_model: D, 输入词嵌入维度 (如512)
        d_k: Q/K/V的维度
        """
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)  # D×dk
        self.W_K = nn.Linear(d_model, d_k, bias=False)  # D×dk
        self.W_V = nn.Linear(d_model, d_k, bias=False)  # D×dv(这里dv=dk)
        self.d_k = d_k
    
    def forward(self, X):
        """
        X: [batch_size, L, D] 输入矩阵
        return: [batch_size, L, d_k] 输出矩阵
        """
        # 1. QKV映射: [batch, L, D] → [batch, L, dk]
        Q = self.W_Q(X)  # 等价于 X @ W_Q
        K = self.W_K(X)
        V = self.W_V(X)
        
        # 2. 计算注意力得分: Q @ K^T / sqrt(dk)
        #    [batch, L, dk] @ [batch, dk, L] → [batch, L, L]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # 3. Softmax归一化: [batch, L, L]
        attn_weights = F.softmax(scores, dim=-1)
        
        # 4. 信息融合: [batch, L, L] @ [batch, L, dk] → [batch, L, dk]
        output = torch.matmul(attn_weights, V)
        
        return output, attn_weights

# 示例
d_model = 512
d_k = 64
batch_size = 2
L = 4  # "我喜欢打篮球"

attn = SelfAttention(d_model, d_k)
X = torch.randn(batch_size, L, d_model)
output, weights = attn(X)
print(output.shape)   # torch.Size([2, 4, 64])
print(weights.shape) # torch.Size([2, 4, 4])

⚠️ 常见问题与避坑指南

  • 必须除以 $\sqrt{d_k}$：否则点积值过大导致 Softmax 进入饱和区，梯度消失，无法更新 $W_K, W_V$
  • K 矩阵需转置（$K^T$）才能与 Q 做矩阵乘法
  • 矩阵形式与模型结构图形式是同一过程的两种表达，建议结合理解
  • 注意力得分矩阵 $A$ 是 $L \times L$，表示所有 token 两两间的关联度

💡 个人总结与延伸

矩阵理解方式将自注意力机制从逐 token 计算提升为并行矩阵运算，这正是 Transformer 高效并行的数学基础。$\text{Softmax}(\frac{QK^T}{\sqrt{d_k}})V$ 这一公式是整个 Transformer 最核心的表达，后续的多头注意力、交叉注意力均基于此。

最近更新:
Contributors: yangwenguang