
注意力机制详解:从 Softmax 到 Flash Attention
1. 注意力机制的数学基础 1.1 从信息检索到注意力 注意力机制源于信息检索的直觉:给定查询 Query,在键值对 (Key-Value) 集合中检索相关信息。注意力是软检索——不是返回最匹配的一项,而是对所有项加权求和。 1.2 Scaled Dot-Product Attention 推导 给定 $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{m \times d_k}$, $V \in \mathbb{R}^{m \times d_v}$: $$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$ 逐步分解: 相似度计算:$S = QK^T \in \mathbb{R}^{n \times m}$,每个元素 $S_{ij} = \sum_{l=1}^{d_k} Q_{il} K_{jl}$ 是点积相似度 缩放:$\hat{S}{ij} = S{ij} / \sqrt{d_k}$,控制方差 归一化:$A_{ij} = \frac{e^{\hat{S}{ij}}}{\sum{j’} e^{\hat{S}_{ij’}}}$,每行 softmax 加权聚合:$O = AV \in \mathbb{R}^{n \times d_v}$ 方差分析:假设 $Q, K$ 各分量独立同分布,均值为 0,方差为 1,则 $S_{ij} = \sum_{l=1}^{d_k} Q_{il}K_{jl}$ 的均值为 0,方差为 $d_k$。当 $d_k=64$ 时,$S_{ij}$ 标准差为 8,softmax 输入范围 $[-24, 24]$,梯度几乎为零。除以 $\sqrt{d_k}=8$ 后方差恢复为 1。 ...