Attention 阶段核心包含以下几个数学公式：
- \(Q = X W_{Q}\)
- \(K = X W_{K}\)
- \(V = X W_{V}\)
- \(Y = \text{Attention}(Q, K, V) = \text{Softmax}(\frac{Q K^{T}}{\sqrt{d_k}})V\)
- \(Z = YW_{O}\)
- \(\text{Output} = \text{LayerNorm}(X+Z)\)
Attention 计算可以分为三大部分：
- 线性投影/矩阵乘法（linear projection/gemm）：\(X\) 和 \(W_{Q}, W_{K}, W_{V}, W_{O}\) 相乘
- 注意力得分计算（attention score）
- 其他运算：layernorm
QKV
其中 gemm 的计算量和访存量如下表所示：
| operation | inference FLOPs | params | output shape |
|---|---|---|---|
| \(A[B,T,\textcolor{red}{D}] \cdot W_Q[\textcolor{red}{D},N,H]\) | \(2BTDNH\) | \(DNH\) | \(Q[B,T,D,H]\) |
| \(A[B,T,\textcolor{red}{D}] \cdot W_K[\textcolor{red}{D},K,H]\) | \(2BTDKH\) | \(DKH\) | \(K[B,T,K,H]\) |
| \(A[B,T,\textcolor{red}{D}] \cdot W_V[\textcolor{red}{D},K,H]\) | \(2BTDKH\) | \(DKH\) | \(V[B,T,K,H]\) |
| \(A[B,T,\textcolor{red}{N,H}] \cdot W_O[\textcolor{red}{N,H},D]\) | \(2BTDNH\) | \(DNH\) | \(\text{Z}[B,T,D]\) |
Attention Score
其中 attention score 的计算量如下表所示：
| operation | inference FLOPs | output shape |
|---|---|---|
| \(Q[\textcolor{blue}{B},T,\textcolor{blue}{K},G,\textcolor{red}{H}] \cdot K[\textcolor{blue}{B},S,\textcolor{blue}{K},\textcolor{red}{H}]\) | \(2BTSKGH=2BTSNH\) | \(\text{score}[B,T,S,K,G]=[B,T,S,N]\) |
| \(\text{softmax}_{S}\ L[B,T,S,K,G]\) | \(O(BTSKG)=O(BTSN)\) | |
| \(S[\textcolor{blue}{B},T,\textcolor{red}{S},\textcolor{blue}{K},G] \cdot V[\textcolor{blue}{B},\textcolor{red}{S},\textcolor{blue}{K},H]\) | \(2BTSKGH=2BTSNH\) | \(Y[B,T,K,G,H]=[B,T,N,H]\) |
根据以上推导，可以得到以下结论：
- Z 的 shape 和输入相同，都是 \([B,T,D]\)，所以可以将两者直接相加，再进行 LayerNorm 得到 Output。
- Self Attention 的计算量取决于 q 和 k/v length。
- 如果忽略 softmax 的话，总共的计算量为 \(O(BTSNH)\)。
为了方便理解，我们考虑以下 \(B=1\)、\(N=1\) 的情况，计算过程为如下公式：
\[ \begin{aligned} Y&=\underbrace{\begin{bmatrix} \alpha_{11} & \alpha_{12} & \cdots & \alpha_{1s}\\ \alpha_{21} & \alpha_{22} & \cdots & \alpha_{2s}\\ \vdots & \vdots & \ddots & \vdots \\ \alpha_{t1} & \alpha_{t2} & \cdots & \alpha_{ts} \end{bmatrix}}_{\displaystyle \alpha=\operatorname{Softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)} \begin{bmatrix} v_1\\ v_2\\ \vdots\\ v_s \end{bmatrix} \\[4pt] &=\begin{bmatrix} \alpha_{11}v_1+\alpha_{12}v_2+\cdots+\alpha_{1s}v_s\\ \alpha_{21}v_1+\alpha_{22}v_2+\cdots+\alpha_{2s}v_s\\ \vdots\\ \alpha_{t1}v_1+\alpha_{t2}v_2+\cdots+\alpha_{ts}v_s \end{bmatrix}. \end{aligned} \]
其中 \(\alpha_{ij}\) 是当前 token 和先前每一个 token 的注意力得分，通过以下方式计算出来，注意 Softmax 是按照行作用的：
\[ s_{ij}=\frac{q_i\cdot k_j}{\sqrt{d_k}},\qquad \alpha_{ij}=\frac{e^{s_{ij}}}{\sum_{t=1}^{n} e^{s_{it}}} \]
\[ \begin{align*} QK^{T} &= \begin{bmatrix} q_1 \\ q_2 \\ \vdots \\ q_t \end{bmatrix} \begin{bmatrix} k_1^T & k_2^T & \cdots & k_s^T \end{bmatrix} \\ &= \begin{bmatrix} q_1 \cdot k_1^T & q_1 \cdot k_2^T & q_1 \cdot k_3^T & \cdots & q_1 \cdot k_s^T \\ q_2 \cdot k_1^T & q_2 \cdot k_2^T & q_2 \cdot k_3^T & \cdots & q_2 \cdot k_s^T \\ \vdots & \vdots & \ddots & \vdots & \vdots \\ q_t \cdot k_1^T & q_t \cdot k_2^T & q_t \cdot k_3^T & \cdots & q_t \cdot k_s^T \end{bmatrix} \\ &\xrightarrow{\text{softmax 逐行归一化}} \begin{bmatrix} \alpha_{11} & \alpha_{12} & \cdots & \alpha_{1s}\\ \alpha_{21} & \alpha_{22} & \cdots & \alpha_{2s}\\ \vdots & \vdots & \ddots & \vdots \\ \alpha_{t1} & \alpha_{t2} & \cdots & \alpha_{ts} \end{bmatrix} \end{align*} \]
注意力得分矩阵的 shape 为 \([T, S]\)，行长度就是 query length，列长度就是 kv length，每行就是一个 token 和之前 token 的注意力打分，还需要乘上对应的 \(v\) 向量。再乘以 \(V\) 的时候收缩的维度是在行上，所以 contracting dimension 是 \(S\)。