Self-Attention 机制深度解析：Q、K、V 的数学直觉与 TypeScript 实现

Michael Meng· 2026年7月23日· ◷ 9 分钟阅读

一、问题引入：为什么需要 Self-Attention？

你在读一篇 2000 字的技术文章，读到第三段时突然要理解一个「指代关系」——「如上所述，该方案……」中的「该」指的是第二段提到的某个概念。人的大脑天然能「回顾」之前看过的所有内容，并找到最有关系的部分。

但传统的 RNN/LSTM 模型做不到这一点：信息必须按顺序一步步传递，距离越远的信息越容易被「遗忘」。

场景	RNN/LSTM 处理方式	Self-Attention 处理方式
理解长句中代词指代	逐 token 传递隐状态，100 步前的信息可能丢失	每个 token 直接查看所有 token，路径长度 O(1)
翻译 "The animal didn't cross...because it was too tired"	难以捕捉 "it" 与 "animal" 的跨句依赖	注意力权重直接关联 "it" 和 "animal"
训练速度	必须串行计算，无法并行	所有位置的注意力可并行计算

Transformer 论文（Attention Is All You Need, 2017）用 Self-Attention 彻底解决了这个问题：让序列中每个位置都能直接「关注」到所有其他位置，计算路径长度恒为 O(1)。

二、核心概念速览

Self-Attention 的核心模型只有三个矩阵——Q（Query）、K（Key）、V（Value），这组命名来源于信息检索隐喻：

矩阵	全称	含义	类比
Q (Query)	查询向量	「我在找什么？」	搜索引擎的搜索词
K (Key)	键向量	「我有什么可以被查？」	文档的关键词索引
V (Value)	值向量	「找到后返回什么内容？」	文档的正文内容

计算流程：

输入序列 X (n × d_model)
    ↓ × WQ / WK / WV
Q = X·WQ, K = X·WK, V = X·WV
    ↓
Score = Q·Kᵀ / √dk      ← 缩放点积注意力
    ↓
Weights = softmax(Score)
    ↓
Output = Weights · V

关键点：√dk 的缩放因子专门解决大维度下 Softmax 梯度消失问题（详见第四章）。

三、完整可运行代码

以下是 Self-Attention 的 TypeScript 完整实现，可在浏览器或 Node.js 中直接运行：

/**
 * Self-Attention 完整实现
 * 参考: "Attention Is All You Need" (Vaswani et al., 2017)
 * 
 * 数学公式:
 *   Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V
 */

// ---------- 工具函数 ----------

/** 创建指定形状的随机矩阵 (rows × cols) */
function randMatrix(rows: number, cols: number): number[][] {
  return Array.from({ length: rows }, () =>
    Array.from({ length: cols }, () => Math.random() * 0.02 - 0.01)
  );
}

/** 矩阵乘法: A (m×n) × B (n×p) → (m×p) */
function matMul(A: number[][], B: number[][]): number[][] {
  const m = A.length;
  const n = A[0].length;
  const p = B[0].length;
  if (n !== B.length) {
    throw new Error(`维度不匹配: [${m}×${n}] × [${B.length}×${p}]`);
  }
  const result: number[][] = Array.from({ length: m }, () => new Array(p).fill(0));
  for (let i = 0; i < m; i++) {
    for (let k = 0; k < n; k++) {
      const aik = A[i][k];
      if (aik === 0) continue; // 稀疏优化：跳过零值
      const rowB = B[k];
      for (let j = 0; j < p; j++) {
        result[i][j] += aik * rowB[j];
      }
    }
  }
  return result;
}

/** 矩阵转置: m×n → n×m */
function transpose(M: number[][]): number[][] {
  const rows = M.length;
  const cols = M[0].length;
  const result: number[][] = Array.from({ length: cols }, () => new Array(rows).fill(0));
  for (let i = 0; i < rows; i++) {
    for (let j = 0; j < cols; j++) {
      result[j][i] = M[i][j];
    }
  }
  return result;
}

/** Softmax: 对矩阵的每行做 Softmax，返回概率分布 */
function softmaxRows(matrix: number[][]): number[][] {
  return matrix.map(row => {
    // 1. 减去最大值 → 数值稳定性（防止 exp 溢出）
    const maxVal = Math.max(...row);
    const exps = row.map(v => Math.exp(v - maxVal));
    const sumExp = exps.reduce((a, b) => a + b, 0);
    return exps.map(v => v / sumExp);
  });
}

// ---------- 核心: Self-Attention ----------

interface AttentionConfig {
  d_model: number;      // 模型隐藏维度（论文: 512）
  d_k: number;          // Q/K 维度（论文: 64，8 头下 d_model/8）
  d_v: number;          // V 维度
}

class SelfAttention {
  private WQ: number[][];
  private WK: number[][];
  private WV: number[][];
  private config: AttentionConfig;

  constructor(config: AttentionConfig) {
    this.config = config;
    // 权重矩阵初始化为小的随机值
    this.WQ = randMatrix(config.d_model, config.d_k);
    this.WK = randMatrix(config.d_model, config.d_k);
    this.WV = randMatrix(config.d_model, config.d_v);
  }

  /**
   * 前向传播: X[n, d_model] → Output[n, d_v]
   * @param X - 输入序列，shape: [seq_len, d_model]
   * @param mask - 可选的注意力掩码（Decoder 用），1=可见 0=遮蔽
   * @returns {{ output, attentionWeights }}
   */
  forward(
    X: number[][],
    mask?: number[][]
  ): { output: number[][]; attentionWeights: number[][] } {
    const seqLen = X.length;
    const { d_k } = this.config;

    // Step 1: 线性投影 X → Q, K, V
    const Q = matMul(X, this.WQ);  // [seqLen, d_k]
    const K = matMul(X, this.WK);  // [seqLen, d_k]
    const V = matMul(X, this.WV);  // [seqLen, d_v]

    // Step 2: Score = Q · Kᵀ → [seqLen, seqLen]
    const KT = transpose(K);
    const scores = matMul(Q, KT);

    // Step 3: 缩放 —— 除以 √d_k（防止 Softmax 进入饱和区）
    const scale = Math.sqrt(d_k);
    const scaledScores = scores.map(row =>
      row.map(v => v / scale)
    );

    // Step 4: 应用 Mask（Decoder 中防止泄露未来信息）
    if (mask) {
      for (let i = 0; i < seqLen; i++) {
        for (let j = 0; j < seqLen; j++) {
          if (!mask[i] || mask[i][j] === 0) {
            scaledScores[i][j] = -1e9; // 相当于 softmax 后概率趋近 0
          }
        }
      }
    }

    // Step 5: Softmax → 注意力权重
    const attentionWeights = softmaxRows(scaledScores);

    // Step 6: Output = AttentionWeights · V → [seqLen, d_v]
    const output = matMul(attentionWeights, V);

    return { output, attentionWeights };
  }
}

// ---------- 测试: 简单序列的自注意力 ----------

function test() {
  const config: AttentionConfig = { d_model: 512, d_k: 64, d_v: 64 };
  const attn = new SelfAttention(config);

  // 模拟 4 个 token 的输入序列，每个 512 维（实际场景会用 Embedding）
  const seqLen = 4;
  const X: number[][] = randMatrix(seqLen, config.d_model);

  const { output, attentionWeights } = attn.forward(X);

  console.log(`输入: [${seqLen}, ${config.d_model}]`);
  console.log(`注意力权重: [${seqLen}, ${seqLen}]`);
  console.log(`输出: [${output.length}, ${output[0].length}]`);

  // 打印 4×4 注意力矩阵（每行和为 1）
  console.log('\n注意力矩阵（行=Query, 列=Key）:');
  attentionWeights.forEach((row, i) => {
    const sum = row.reduce((a, b) => a + b, 0);
    console.log(
      ` Token${i}: [${row.map(v => v.toFixed(3)).join(', ')}]  sum=${sum.toFixed(4)}`
    );
  });
}

test();

运行结果示例：

输入: [4, 512]
注意力权重: [4, 4]
输出: [4, 64]

注意力矩阵（行=Query, 列=Key）:
 Token0: [0.248, 0.251, 0.249, 0.252]  sum=1.0000
 Token1: [0.250, 0.254, 0.248, 0.248]  sum=1.0000
 Token2: [0.252, 0.247, 0.249, 0.251]  sum=1.0000
 Token3: [0.248, 0.252, 0.249, 0.251]  sum=1.0000

四、逐行精讲（面试能讲出来就是高手）

4.1 为什么 Q、K、V 需要三个不同的投影矩阵？

如果 WQ = WK = WV，Self-Attention 会退化为对称注意力：token i 对 token j 的关注度等于 token j 对 token i 的关注度。但真实语言中这不对称——「猫追老鼠」中「追」对「猫」的关注应该不同于「猫」对「追」的关注。

// ❌ 错误: 只用 WQ 投影
const V = matMul(X, WQ); // V 和 Q 使用相同投影，丢失不对称性

// ✅ 正确: Q/K/V 独立权重
const Q = matMul(X, this.WQ);
const K = matMul(X, this.WK);
const V = matMul(X, this.WV);

4.2 缩放因子 √d_k 为什么不可或缺？

直观理解：假设 Q 和 K 的各分量独立同分布，均值为 0、方差为 1，则点积 Q·Kᵀ 的方差 = d_k。当 d_k=64 时，点积值可能达到 ±20，此时 Softmax 输出接近 one-hot（梯度 → 0）：

// 演示：大维度下不加缩放的后果
const bigDot = 20;  // 64 维点积的典型值
const smallDot = 0.5;
console.log(Math.exp(bigDot) / (Math.exp(bigDot) + Math.exp(smallDot) + Math.exp(0.3) + Math.exp(0.1)));
// → 0.999... 梯度极低，模型学不动

// 缩放后
const scaled = bigDot / Math.sqrt(64); // 20 / 8 = 2.5
console.log(Math.exp(scaled) / (Math.exp(scaled) + Math.exp(smallDot) + Math.exp(0.3) + Math.exp(0.1)));
// → 0.768 梯度健康

4.3 Softmax 数值稳定性：减去最大值

直接计算 exp(x) 在 x 较大时会溢出：Math.exp(800) = Infinity。标准做法是减去每行最大值：

function stableSoftmax(row: number[]): number[] {
  const maxVal = Math.max(...row);
  const exps = row.map(v => Math.exp(v - maxVal)); // exp(x-max) ≤ exp(0) = 1
  const sumExp = exps.reduce((a, b) => a + b, 0);
  return exps.map(v => v / sumExp);
}

这等价于原始 Softmax（分子分母约掉 exp(maxVal)），但避免了溢出。

五、常见问题与踩坑记录

Q1：Q 和 K 的维度 d_k 必须等于 d_v 吗？
A：不必须。论文中 d_k = d_v = d_model / h（h 为头数），但理论上 d_k 和 d_v 可以不同。只要确保 Q 和 K 的最后一个维度一致即可（否则无法做点积）。V 的维度决定了输出维度。

Q2：序列长度 n=1 时 Self-Attention 还有意义吗？
A：有的。即使只有 1 个 token，Self-Attention 仍然在做 WQ/WK/WV 投影和 Softmax（退化为恒等映射），相当于对输入做了线性变换。这在 Decoder 的生成阶段（逐个 token 输出）时持续有用。

Q3：实际训练中 Self-Attention 的显存瓶颈在哪里？
A：瓶颈在 Q·Kᵀ 矩阵，shape 为 [n, n]。当 n=4096 时，仅这个矩阵就占 4096² × 4 bytes ≈ 64MB（单头），多头时更严重。这就是 Flash Attention 等优化算法的切入点。

Q4：为什么 Mask 要用 -1e9 而不是 -Infinity？
A：Number.NEGATIVE_INFINITY 在 JavaScript 中合法，但 Math.exp(-Infinity) = 0，0/0 = NaN 在某些实现中会导致问题。用 -1e9 足够小（exp(-1e9) ≈ 0）且安全。

六、决策框架（一张图搞定）

你需要理解/实现 Attention 吗？
├── 只是面试 → 记住 Q·Kᵀ / √dk → Softmax → ·V
│   └── 追问：为什么缩放？→ 防止 Softmax 饱和
│
├── 要写代码 → 选择维度实现
│   ├── 单头 Self-Attention → 本文代码直接用
│   │   └── 是否需要 Mask？
│   │       ├── 是（Decoder）→ 传入 mask 参数
│   │       └── 否（Encoder）→ mask = undefined
│   │
│   └── 多头 Attention → 并行 h 个 Self-Attention + Concat + W_O
│       └── 典型配置: h=8, d_k=d_v=64, d_model=512
│
└── 要做性能优化 → Flash Attention / PagedAttention
    ├── n < 1024 → 标准实现就够了
    ├── n < 4096 → 考虑低精度（FP16）
    └── n > 4096 → Flash Attention 或 Sliding Window

七、面试速记

Q：Self-Attention 的时间复杂度是多少？
A：O(n²·d)。其中 n 是序列长度、d 是维度。Q·Kᵀ 矩阵乘法是 n×d × d×n = n²·d，这是瓶颈。对比 RNN 的 O(n·d²)，当 n < d 时 Self-Attention 更快（实际中通常成立：n≈512, d≈512-2048）。

Q：Self-Attention 和 Cross-Attention 有什么区别？
A：Self-Attention 中 Q、K、V 都来自同一个序列（同源），用于捕捉序列内部依赖。Cross-Attention 中 Q 来自 Decoder、K/V 来自 Encoder，用于 Decoder 从 Encoder 输出中提取相关信息。这是机器翻译等 seq2seq 任务的核心机制。

Q：位置编码为什么要加到 Embedding 上而不是拼接到 Embedding 后面？
A：加法等价于在 d_model 维空间中做位置偏移，保留了 Embedding 的语义空间结构。拼接会增加维度，引入额外参数且使模型更难泛化到未见过的位置。

八、总结

Self-Attention 的核心公式只有一行：softmax(Q·Kᵀ / √d_k) · V，但蕴含了信息检索、线性投影、数值稳定三个层次的工程智慧。
Q/K/V 分离不是冗余设计，而是必要的不对称性——语言中「A 关注 B」≠「B 关注 A」。
√d_k 的缩放是保证 Softmax 梯度健康的「隐形英雄」，这个 1 行代码能决定模型能否收敛。
Mask 机制使注意力既能用于双向理解（BERT Encoder），也能用于因果生成（GPT Decoder）。
以上 TypeScript 实现可直接在浏览器中运行并可视化注意力矩阵，适合教学和原型验证。

标签：AI 面试 TypeScript 深度学习算法

Comments 留言讨论

还没有评论，来抢个沙发，聊聊你的看法～