**注意力机制（Attention Mechanism）**会根据输入动态计算一组权重，再用这些权重组合 Value 向量。Query 负责与允许访问的 Key 匹配，匹配分数归一化后决定每个 Value 对输出的贡献。它更接近一种可微的软检索，不是对人类注意力的机械复刻。

核心公式很短：

text 
Attention(Q, K, V, M) = softmax((QK^T / sqrt(d_k)) + M) V

真正容易出错的是公式周围的工程契约：张量形状、Mask 语义、Softmax 数值稳定性、位置信息、多头布局、显存占用、推理状态，以及如何证明实现没有悄悄关注错误位置。

核心结论

  • 每个 Query 行都会在允许访问的 Key 上形成一组分布，对应输出是 Value 行的加权和。
  • 自注意力、因果自注意力与交叉注意力使用相同核心计算，区别在 Q/K/V 的来源和可见范围。
  • Q、K、V 都是训练得到的投影。「我想找什么」等比喻适合建立直觉，但不是固定语义标签。
  • 稠密注意力需要计算目标长度 L 与来源长度 S 的成对关系，分数计算随 L * S 增长。
  • FlashAttention 优化数据搬运和中间结果存储，在数值误差允许范围内保持精确注意力语义；稀疏或线性注意力则改变了计算本身。
  • 注意力权重可以辅助诊断，但不能单独证明模型答案的因果来源。

注意力机制解决了什么问题

注意力允许模型直接从一组表示中读取信息，而不必先把整段输入压缩成一个固定向量。早期神经机器翻译工作中，Bahdanau、Cho 与 Bengio提出可微的软对齐，让 Decoder 在生成不同目标词时读取不同的 Encoder 状态。2017 年的 Transformer 论文进一步把缩放点积注意力作为序列内部与序列之间传递信息的核心机制。

理解这段历史很重要，因为 Attention 不等于 Transformer Self-Attention：

类型 Query 来源 Key 与 Value 来源 常见可见范围
Encoder Self-Attention Encoder 表示流 同一表示流 所有非 Padding 位置
Causal Self-Attention Decoder 表示流 同一表示流 当前与更早位置
Cross-Attention Decoder 或目标表示流 Encoder、图像、音频或其他来源 策略允许的来源位置

Attention 让远距离位置在一层内就能直接交互，也支持训练阶段并行计算多个 Token。但它不保证模型一定学会所需的远距离关系；稠密 Attention 还会为全部允许位置对支付计算成本，即使其中大部分没有价值。

QKV 与张量形状

Q、K、V 分别承担匹配与内容传递接口。给定目标表示 X_q 和来源表示 X_kv，模型通过训练得到的矩阵计算：

text 
Q = X_q  W_Q
K = X_kv W_K
V = X_kv W_V

批量多头注意力可以采用以下形状契约：

张量 形状 含义
Q [B, H_q, L, E] L 个发起信息查询的目标位置
K [B, H_kv, S, E] S 个用于匹配的来源位置
V [B, H_kv, S, E_v] 来源位置携带的内容
Score [B, H_q, L, S] Query-Key 成对 Logit
Output [B, H_q, L, E_v] 每个 Query 位置、每个 Head 的上下文向量

其中 B 是 Batch Size，H 是 Head 数，L 是目标长度，S 是来源长度，E 是 Query/Key 的 Head Dimension。普通 Multi-Head Attention 的 Query Head 与 KV Head 数相同；Multi-Query Attention 和 Grouped-Query Attention 则让多个 Query Head 共享更少的 KV Head。

对单个 Query 行 q_i，计算可拆成三步：

  1. 为每个允许的 Key j 计算 q_i dot k_j / sqrt(d_k)；
  2. 沿 Key 轴归一化这些 Logit；
  3. 返回 sum_j(weight_ij * v_j)。

输出混合的是 Value，而不是原始 Token。之后还会经过 Output Projection、Residual Path、Normalization、Feed-Forward Block 和更多层，因此不能把某个 Attention Weight 直接等同于最终答案的重要性。

为什么除以 sqrt(d_k)

如果 Query 和 Key 的各分量近似零均值、单位方差，点积方差会随 d_k 增长。过大的 Logit 会让 Softmax 变得极尖，进而产生很弱的梯度。原始 Transformer 推导使用 1 / sqrt(d_k) 缩放，使初始化阶段的 Logit 范围更稳定。

这是建立在分布假设上的解释，不是对所有训练后表示的普遍定律。生产 Kernel 可能使用更高精度累加、算子融合或模型专属缩放，最终必须以具体模型和 Runtime 为准。

Mask 是注意力计算契约的一部分

Mask 决定每个 Query 允许读取哪些 Key。它必须在 Softmax 之前作用到 Logit；如果先做 Softmax 再把部分权重置零，剩余位置的概率和就不再正确。

Mask 类型 目的 常见形状
Causal Mask 阻止生成位置读取未来 Token [L, S] 或可广播等价形式
Padding Mask 排除变长 Batch 中的 Padding [B, 1, 1, S] 或 API 专属形式
Local/Window Mask 限制每个 Query 的可见邻域 [L, S] 或 Kernel 参数
Application Mask 执行架构定义的可见策略 取决于 Runtime

最常见的三个实现错误是：

  1. **极性相反：**一个 API 使用 true = 可见，另一个 API 使用 true = 屏蔽。
  2. **Softmax 轴错误：**沿 Query 或 Head 维度归一化，而不是沿 Key 维度。
  3. **整行被屏蔽：**所有 Logit 都变成负无穷，最终产生 NaN，却没有显式报错。

当前 PyTorch scaled_dot_product_attention 的布尔 attn_mask 使用 true 表示允许参与，而 MultiheadAttention.key_padding_mask 使用 true 表示忽略。函数式 SDPA 还会在 dropout_p > 0 时始终执行 Dropout，评估阶段需要显式传入零。这些是 API 契约，不是数学公式自身的属性。

flowchart LR Q["Query B,H,L,E"] --> S["Q 乘 K 转置"] K["Key B,H,S,E"] --> S S --> C["缩放并应用 Mask"] C --> W["沿 S 稳定 Softmax"] W --> O["Weight 乘 V"] V["Value B,H,S,Ev"] --> O O --> R["输出 B,H,L,Ev"]

可运行的 Go 实现

下面的无依赖程序实现一个 Attention Head。它会校验形状，约定 true 表示允许访问，拒绝整行 Mask，先减去行最大值再计算指数，并返回 Weight 供测试验证。示例从 Q/K/V 开始；完整可训练层还需要实现投影参数与梯度。

go 
package main

import (
	"errors"
	"fmt"
	"log"
	"math"
)

func matrixWidth(name string, matrix [][]float64) (int, error) {
	if len(matrix) == 0 || len(matrix[0]) == 0 {
		return 0, fmt.Errorf("%s must not be empty", name)
	}
	width := len(matrix[0])
	for i, row := range matrix {
		if len(row) != width {
			return 0, fmt.Errorf("%s row %d has width %d, want %d", name, i, len(row), width)
		}
	}
	return width, nil
}

func scaledDotProductAttention(
	q, k, v [][]float64,
	allowed [][]bool,
) ([][]float64, [][]float64, error) {
	qWidth, err := matrixWidth("q", q)
	if err != nil {
		return nil, nil, err
	}
	kWidth, err := matrixWidth("k", k)
	if err != nil {
		return nil, nil, err
	}
	vWidth, err := matrixWidth("v", v)
	if err != nil {
		return nil, nil, err
	}
	if qWidth != kWidth {
		return nil, nil, errors.New("q and k head dimensions must match")
	}
	if len(k) != len(v) {
		return nil, nil, errors.New("k and v must have the same source length")
	}
	if len(allowed) != len(q) {
		return nil, nil, errors.New("mask must have one row per query")
	}

	scale := math.Sqrt(float64(qWidth))
	output := make([][]float64, len(q))
	weights := make([][]float64, len(q))

	for i := range q {
		if len(allowed[i]) != len(k) {
			return nil, nil, fmt.Errorf("mask row %d has the wrong source length", i)
		}

		logits := make([]float64, len(k))
		maxLogit := math.Inf(-1)
		for j := range k {
			if !allowed[i][j] {
				logits[j] = math.Inf(-1)
				continue
			}
			var dot float64
			for d := range q[i] {
				dot += q[i][d] * k[j][d]
			}
			logits[j] = dot / scale
			if logits[j] > maxLogit {
				maxLogit = logits[j]
			}
		}
		if math.IsInf(maxLogit, -1) {
			return nil, nil, fmt.Errorf("query row %d is fully masked", i)
		}

		weights[i] = make([]float64, len(k))
		var denominator float64
		for j := range logits {
			if allowed[i][j] {
				weights[i][j] = math.Exp(logits[j] - maxLogit)
				denominator += weights[i][j]
			}
		}
		output[i] = make([]float64, vWidth)
		for j := range weights[i] {
			weights[i][j] /= denominator
			for d := range v[j] {
				output[i][d] += weights[i][j] * v[j][d]
			}
		}
	}
	return output, weights, nil
}

func rowSum(row []float64) float64 {
	var total float64
	for _, value := range row {
		total += value
	}
	return total
}

func main() {
	q := [][]float64{{1, 0}, {0, 1}}
	k := [][]float64{{1, 0}, {0, 1}}
	v := [][]float64{{2, 1}, {0, 3}}
	causal := [][]bool{{true, false}, {true, true}}

	output, weights, err := scaledDotProductAttention(q, k, v, causal)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("row sums: %.1f %.1f\n", rowSum(weights[0]), rowSum(weights[1]))
	fmt.Printf("masked future weight: %.1f\n", weights[0][1])
	fmt.Printf("output shape: %d x %d\n", len(output), len(output[0]))
}

// row sums: 1.0 1.0
// masked future weight: 0.0
// output shape: 2 x 2

这段代码是教学参考，不能替代框架中的融合 Kernel。它展示了生产测试必须保留的几个不变量：

  • 每个未被完全屏蔽的 Weight 行之和约等于 1；
  • 被屏蔽位置的权重为 0；
  • 输出形状是 [Query Length, Value Dimension]；
  • 维度不匹配或整行 Mask 时显式失败；
  • 大 Logit 不会让指数计算溢出。

自注意力、交叉注意力与因果注意力

这些名称描述的是数据来源与可见性，而不是不同的打分公式。

Self-Attention 使用同一表示流生成 Q、K、V。Encoder 风格的自注意力通常允许每个真实 Token 读取所有其他真实 Token。若模型不引入位置 Embedding、Rotary Transform、Relative Bias 或其他位置信号，该操作本身无法区分排列顺序。

Causal Self-Attention 在自注意力上增加位置 t 不能读取未来位置的规则。Teacher Forcing 训练时，三角 Mask 让所有目标位置可以并行计算，又不会泄露未来答案；自回归推理时，模型逐 Token 生成，并通过 KV Cache复用历史 Key 和 Value。

Cross-Attention 的 Query 来自一个表示流，Key/Value 来自另一个表示流。Encoder-Decoder 翻译中，Decoder 查询 Encoder 输出；多模态模型中，文本 Token 可以查询图像或音频表示。因此 Cross-Attention 的目标长度 L 与来源长度 S 不必相等。

Transformer 架构指南进一步解释 Attention 如何与残差连接、归一化、前馈网络和位置信息协作。

MHA、MQA 与 GQA

Multi-Head Attention 为每个 Head 使用独立学习投影，并行计算 Attention，再拼接各 Head 输出并通过 Output Projection 混合。多个 Head 提供不同表示子空间，但架构不会预先规定某个 Head 必须负责语法、另一个负责指代；专业化、冗余与 Head 间交互都要在训练后测量。

text 
head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O

Decoder 推理还需要回答另一个问题：需要存储和读取多少组不同的 Key/Value Head？

变体 Query Head KV Head 主要权衡
Multi-Head Attention（MHA） 多个 与 Query 相同 每个 Head 的 KV 灵活性最大，KV 状态也最大
Multi-Query Attention（MQA） 多个 1 个 降低 Decode 带宽与 Cache 容量，质量取舍取决于训练
Grouped-Query Attention（GQA） 多个 介于 1 与 Query Head 数之间 在质量和容量之间取中间方案

MQA 论文针对增量 Decode 反复读取 K/V 的显存带宽成本；GQA 论文定义了中间形态，并在论文给定模型与算力条件下评估 Uptraining。二者都不是可以对任意 MHA Checkpoint 无损开启的 Serving Flag，模型架构和权重必须支持。

复杂度与高效 Attention

不写清维度时，只说「Attention 是 O(n²)」不够准确。设目标长度为 L、来源长度为 S、模型宽度为 d_model、单 Head 宽度为 d_k：

  • Q/K/V Projection 含有与序列长度乘 d_model² 成正比的项；
  • Score 与 Value Aggregation 每个 Head 含有与 L * S * d_k 成正比的项；
  • 显式保存所有 Score 或 Probability 需要 O(B * H * L * S) 内存；
  • 因果自回归 Decode 可以复用历史 K/V Projection，但仍要读取不断增长的 Cache。

这也能区分经常被混在一起的四类优化：

技术 是否改变数学 Attention 主要目标
Fused SDPA / FlashAttention 不改变，受数值行为约束 显存访问、Kernel Launch、中间结果存储
MQA / GQA 改变训练后的 Head 共享 Decode 带宽与 KV Cache 容量
Sparse / Local Attention 改变可见位置对 Pair 数量与长序列扩展
KV Cache / Prefix Reuse 对兼容状态复用不改变 重复 Projection 或 Prefill 工作

FlashAttention是 IO-aware 的精确 Attention，通过 Tiling 与 Recomputation 减少 GPU 不同内存层级之间的读写。FlashAttention-2进一步优化 GPU 工作划分。论文性能数字只属于给定硬件、模型、Shape 和基线，不能直接当作任意系统的固定加速倍数。当 Dtype、Shape、Mask 或硬件不满足条件时，框架也可能回退到其他 Backend。

设计长上下文系统时，要明确区分「保持公式不变的 Kernel 优化」和「改变可见关系的算法」。长上下文与 RAG 决策指南讨论了什么时候扩大模型上下文才是正确选择。

注意力权重不是因果解释

Attention Heatmap 展示的是某次 Forward Pass 中某层、某个 Head 的归一化路由系数。它没有包含：

  • Value 向量的大小与方向；
  • 当前 Head 的 Output Projection；
  • 绕过当前 Head 的 Residual Stream 信息；
  • 其他 Head 与后续 Layer 的交互；
  • 改变被关注 Token 后，最终预测是否真的改变。

论文 Attention is not Explanation在其 NLP 模型实验中发现，Attention Weight 与部分 Feature Importance 指标相关性较弱，而且显著不同的 Attention Distribution 可能得到相似预测。应把结论限定在证据范围内：原始 Weight 本身不能建立因果重要性。

可以用 Heatmap 发现 Mask 错误、位置模式、异常集中或提出机制假设，但要用 Token/Head Ablation、Activation Patching、受控反事实实验和任务输出变化继续验证。

如何测试 Attention 实现

Attention 实现应先通过语义测试，再谈性能优化：

  1. **Shape Contract：**覆盖 L != S、多 Batch、多 Head，确认输出和 Score 形状。
  2. **Mask Contract：**断言屏蔽位置概率为 0，允许行概率和为 1。
  3. **Causality：**修改未来 Token，证明更早位置的因果输出不变。
  4. **Cross-Attention Provenance：**修改 Source Stream，确认只有设计允许的路径受影响。
  5. **Numerical Stability：**测试大 Logit、低精度输入和整行 Mask。
  6. **Reference Parity：**用小型确定性输入与可信框架实现对比。
  7. **Gradient Check：**训练实现需要在小张量上对比解析梯度与有限差分。
  8. **Backend Audit：**记录实际运行的 Fused Kernel，以及发生 Fallback 的原因。
  9. **Workload Benchmark：**按 Batch Size、序列 Shape、Dtype、Mask 与硬件报告延迟和峰值显存。
  10. **Task Quality：**确认更快 Kernel 或新的 Attention Pattern 没有破坏目标任务指标。

不要用随机 Heatmap 证明模型学会某种语言关系，不要在缺少张量形状和硬件信息时报告 Kernel 性能，也不要在没有转换测试的情况下跨 API 搬运 Mask 约定。

常见问题

Attention 就是整个 Transformer 吗？

不是。Transformer Block 还包含学习投影、输出混合、残差连接、归一化、前馈层、位置机制，以及架构专属的 Routing 或 Gating。Attention 是这个更大系统中的 Token 交互机制。

为什么 Self-Attention 训练能并行，生成仍然是串行的？

Teacher Forcing 训练时完整目标序列已经存在，Causal Mask 负责阻止未来信息泄漏，因此所有位置可以一起计算。自回归生成时，位置 t+1 依赖实际采样或选出的 Token t，所以输出位置仍逐个产生，即使每一步内部使用并行矩阵运算。

多头注意力的 Head 数应该怎么选？

Head 数是模型与 Kernel 的共同设计参数。它必须满足实现的整除和布局约束，也会改变 Head Dimension、显存行为、并行度与学习容量。应在相同训练预算和 Serving Workload 下比较候选架构，而不是假设 Head 越多越好。

FlashAttention 会让稠密注意力变成线性时间吗？

不会。它避免在高带宽显存中完整物化 Score Matrix，并减少 IO，但精确稠密 Attention 仍需计算 Query-Key 成对关系。它改善内存与实际耗时，不会把算术复杂度普遍变成线性。

后续应该按什么顺序学习？

先用 注意力机制术语页巩固定义，再阅读 Transformer 架构理解完整 Block，最后通过 KV Cache 指南学习自回归推理状态。历史与 API 事实应回到 Bahdanau Attention 论文、Transformer 论文和当前框架文档核验。