题目内容

实现 Multi-Head Attention 算法。

给定输入序列 $X$（$L \times d$）、头数 $h$，以及投影矩阵 $W_q, W_k, W_v, W_o$（各 $d \times d$），按以下步骤计算输出：

1. 线性投影：$Q = XW_q$，$K = XW_k$，$V = XW_v$。
2. 拆分多头：将 $Q, K, V$ 拆分为 $h$ 个头，每个头维度为 $d_k = d / h$。
3. 并行注意力：每个头独立计算 $\mathrm{softmax}(QK^{T} / \sqrt{d_k})V$。
4. 拼接输出：$\mathrm{output} = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_h) \cdot W_o$。

输入描述

第一行包含三个整数 $L, d, h$（$1 \le L \le 8$，$2 \le d \le 8$，$1 \le h \le d$，$d \bmod h = 0$），分别表示序列长度、模型维度、头数。

接下来 $L$ 行，每行 $d$ 个浮点数，表示输入序列 $X$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_q$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_k$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_v$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_o$。

输出描述

输出 $L$ 行，每行 $d$ 个浮点数，表示 Multi-Head Attention 的输出矩阵，保留 $4$ 位小数。

样例 $1$

输入

2 4 2
1.0 0.0 1.0 0.0
0.0 1.0 0.0 1.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0

输出

0.6698 0.3302 0.6698 0.3302
0.3302 0.6698 0.3302 0.6698

说明

$d_{model} = 4$，$h = 2$，每头维度 $d_k = 2$。$W_q = W_k = W_v = W_o = I$（单位矩阵）。

拆分为 $2$ 个头，每头处理 $2$ 维。各头独立计算 Attention 后拼接，经输出投影得到最终结果。

四个投影矩阵按样例解释写成单位矩阵。