Q、K、V 的投影与注意力计算

以包含 4 个 Query 头和 2 个 KV 头的 GQA 为例,展示各步的张量形状。

Q = XWQ K = XWK V = XWVHi = softmax(QiKgT / √dh + M)Vg
Query 头

图示可左右滑动,保留矩阵比例与文字大小。

01

投影与拆头

输入分别投影为 Q、K、V,再按每头宽度拆分。

X3 × 8×WQWKWV8 × (8 + 4 + 4)Q3 × 8K3 × 4V3 × 4每 2 列拆为一个头
02

注意力概率

Q 与 K 的匹配分数经缩放、因果遮罩和 softmax 转为概率。

Qᵢ3 × 2×Kᵍᵀ2 × 3缩放 · 因果遮罩 · softmax沿 Key 位置归一化Aᵢ3 × 3行:Query 位置列:Key 位置未来位置的概率为 0
03

汇总 Value

按注意力概率对各位置的 Value 加权求和,得到单个头的输出。

Aᵢ3 × 3×Vᵍ3 × 2=Hᵢ3 × 2各 Query 头分别计算注意力概率。共享 K/V 的头仍有各自的 A 与 H。
04

拼接与输出

各头输出沿特征维拼接,再由输出投影映射回 d 维。

[H₁ | H₂ | H₃ | H₄]3 × 8×Wₒ8 × 8=O3 × 8输出宽度为 d = 8与残差流的宽度一致
维度
B = 1(省略批次轴),T = 3,d = 8,h = 4,h_kv = 2,d_h = 2。每格对应一个元素。
对象
W 表示可学习权重,Q、K、V、H、O 表示中间张量。A 的每行是一组注意力概率,和为 1。色块深浅仅作示意。
约定
头编号从 1 开始,g = ⌈i / 2⌉。Q、K 在匹配前应用 RoPE;M 在可见位置取 0,在未来位置取 −∞。图中省略偏置与 dropout。