P4496.多头注意力
题目描述
实现一个简化版的 Multi-Head Attention，具体步骤如下：
-
将输入序列
X∈RT×dmodel
经线性变换生成
Q=XWQ+bQ,K=XWK+bK,V=XWV+bV
-
将 Q,K,V 按 head 数 h 切分为维度
dk=hdmodel
的子矩阵。
-
对第 i 个 head 计算 Scaled Dot-Product Attention：
Attentioni(Qi,Ki,Vi)=softmax(dkQiKi⊤)Vi
-
将所有 head 输出拼接：
H=concat(H1,H2,…,Hh)
-
输出层线性映射：
O=HWO+bO
输入参数
X：形状为 T×dmodel 的输入序列W_Q, W_K, W_V：形状为 dmodel×dmodel 的线性层权重b_Q, b_K, b_V：长度为 dmodel 的偏置项h：Multi-head attention 的 head 数量W_O：形状为 dmodel×dmodel 的输出层权重b_O：长度为 dmodel 的输出偏置
返回值
O：形状为 T×dmodel 的最终注意力输出矩阵
示例
输入：
X =
[[1, 2, 0, 1],
[0, 1, 1, 0]]
W_Q =
[[1, 0, 1, 0],
[0, 1, 0, 1],
[1, -1, 0, 0],
[0, 0, 1, -1]]
b_Q = [0, 0, 0, 0]
W_K =
[[1, 1, 0, 0],
[0, 1, 1, 0],
[1, 0, 0, 1],
[0, 0, 1, 1]]
b_K = [0, 0, 0, 0]
W_V =
[[1, 0, 0, 1],
[0, 1, 1, 0],
[1, 0, 1, 0],
[0, 1, 0, 1]]
b_V = [0, 0, 0, 0]
W_O =
[[1, 0, 1, 0],
[0, 1, 0, 1],
[1, 0, 0, 1],
[0, 1, 1, 0]]
b_O = [0, 0, 0, 0]
输出：
O = [[3.0, 4.7768, 2.8884, 4.8884], [3.0, 3.0, 2.0, 4.0]]
提示
-
输入范围：
−1000≤X[i,j]≤1000
-
权重范围：
−10≤WQ,;WK,;WV,;WO≤10
-
偏置范围：
−5≤bQ,;bK,;bV,;bO≤5
-
head 数要求：
1≤h≤dmodel，且 dmodelmodh=0
-
softmax 输出满足范围与归一性：
0≤softmax(zi)≤1，且 ∑isoftmax(zi)=1
-
最终输出矩阵 O 可为任意实数
