注意力机制简述

向量化

将不同的数据类型，通过Embedding模型（对于CV来说，也就是说feature extractor、backbone）转换为token。

对向量进行归一化后，可以进行点积（inner product）。点积可以看作是“语义相似度”的一种度量方法。我们可以通过计算点积来衡量不同词之间的“关注程度”

自注意力

通过自己乘自己的转置，让每个Token根据上下文动态调整自己的已于表示

注意力机制的Q、K、V

$Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$

1
2
3

scores = torch.matmul(Q,K.transpose(-2,-1)) / sqrt_dim
weights = torch.softmax(scores, dim=-1)
output = torch.matmul(weights, V)

查询（Query）

指当前需要处理的信息。模型根据查询向量在输入序列中查找相关信息
索引向量

键（Key）

指来自输入序列的一组表示。用于根据查询向量计算注意力权重。注意力权重反映了不同位置的输入数据与查询的相关性。
求权向量

值（Value）

指来自输入序列的一组表示。用于根据注意力权重计算加权和，得到最终的注意力输出向量，其包含了查询最相关的输入信息

注意力机制通过计算查询向量与各个键向量之间的相似性，为每个值向量分配一个权重。然后，将加权的值（注意力分数）相加，得到一个输入序列最相关信息的输出向量。这个输出向量的形状和查询向量相同，将用于下一步模型计算的输入

具体过程

得到 X ：我们将 小明/喜欢/吃/苹果/吗 这句话划分为5个6维的词向量，也就是说这句话可以化成一个 5 $\times$ 6 的矩阵 X
得到 Q、K、V：我们定义三个矩阵，分别是W_Q、W_K、W_V，大小均为 6 $\times$ 4（列数可以自定义）。将 X 与这三个矩阵相乘我们得到 Q、K、V三个矩阵（5 $\times$ 4）
- 具体来说，在这个例子中：
  Q代表语义意图，是一个问题。比如 小明 对应的 Q 是：想知道谁是主语， 吃 对应的 Q 是想知道谁是宾语。
  K代表能提供的信息类型，是高度抽象的。比如 小明 表示自己是主语， 苹果 表示自己是可食物体。
  V代表真正携带的信息内容。比如 小明 就是小明这个具体的人，吃就代表这个具体的动作。
计算相关性矩阵$QK^T$：得到是一个未归一化的得分矩阵，形状是5 $\times$ 5。在乘法运算中，Q的每一行是一个词的Query，K的每一列是一个词的Key。如果Query_i关注Key_j，那么得分矩阵的对应[i][j]的值理应该大。（简单来说，就是Q从K中找到想要的答案）
相关性矩阵$QK^T$除以$\sqrt{d_k}$后，使用softmax归一化。$\sqrt{d_k}$是一个缩放因子，除以它的目的是缩小数值范围，让Softmax函数在一个比较平缓的区域工作，从而减轻梯度消失问题
与V矩阵相乘得到最终输出。这里将归一化后的相关性矩阵的每一行乘以V矩阵的每一列。V矩阵中每一列代表一个句子。此时每个词的输出是所有词的Value的加权和，所以最终每个词输出的向量都不仅仅是自己的含义，还融合了他“所关注的那些词的信息”

通过上述方法，我们可以发现将输入序列的顺序完全大陆那后，得到的结果相同，具有置换不变性。本质上是一个全连接的GNN。为此我们必须引入位置编码。

多头注意力机制

把Q、K、V向量先线性映射成多个子空间，然后在每个子空间独立计算注意力，最后再把各个头的输出拼接起来。

Flash Attention

我们将传统的全局计算输出单元展开如下

$O_i = \frac{e^{S_{i1}}V_1 + e^{S_{i2}}V_2 + \dots + e^{S_{i4096}}V_{4096}}{e^{S_{i1}} + e^{S_{i2}} + \dots + e^{S_{i4096}}}$

其中S{ij}是相关性矩阵的值。我们发现每个$S\{ij} $都是用完就抛的，所以我们可以分块到SRAM中直接算出最终结果。具体来说，外层循环时Q的部分行，内层循环扫过所有的K和V（滚动列），将最终结果写回HBM

将Q、kK、V划分为大小$B_c \times d$ 和 $B_r \times d$的块
对每块计算局部注意力并累加，通过在线softmax校正避免数值溢出。

q = torch.randn(2,16,4096,64,device="cuda") # [batch, heads, seq_len. dim]
k = torch.randn_like(q)
v = torch.randn_like(q)

output = flash_attention(q, k, v, dropout_p=0.1, softmax_scale=1.0)

Fused Attention

将Attention 机制的算子序列（MatMul $\rightarrow$ Scale $\rightarrow$ Softmax $\rightarrow$ Dropout $\rightarrow$ MatMul）融合到了一个单一的CUDA中