题目内容

实现 Transformer Encoder Block 的前向计算，结构是 Post-Norm，不使用 Dropout，所有位置之间都可以做注意力。

输入 $X$ 的形状是 $B \times T \times d_{model}$。用 $W_Q,W_K,W_V$ 把 $X$ 投影成 $Q,K,V$，再拆成 $h$ 个头，每个头的维度 $d_k=d_{model}/h$。每个头计算缩放点积注意力，缩放因子是 $\sqrt{d_k}$，然后拼回去，乘 $W_O$。

第一层残差和 LayerNorm：$\mathrm{LN}(X+\mathrm{Attn})$。LayerNorm 做在最后一维上，方差除以 $d_{model}$，再加上 $\mathrm{eps}$。$\mathrm{gamma}$、$\mathrm{beta}$ 按最后一维相乘相加。

前馈网络是 $\mathrm{ReLU}(X_1 W_1+b_1)W_2+b_2$。$W_1$ 的形状是 $d_{model}\times d_{ff}$，$W_2$ 的形状是 $d_{ff}\times d_{model}$。再做一次残差和 LayerNorm，得到 $Y$。

实现 class Solution 的 forward(self, X, W_Q, W_K, W_V, W_O, h, W_1, W_2, b_1, b_2, gamma_1, gamma_2, beta_1, beta_2, eps)。注意力投影和输出投影没有额外偏置。

输入描述

第一行五个整数 $B,T,d_{model},h,d_{ff}$。

接着按下面的顺序给出浮点数，每个矩阵按行排列：

$X$ 共 $B \times T$ 行，每行 $d_{model}$ 个数。

$W_Q,W_K,W_V,W_O$ 各 $d_{model}$ 行，每行 $d_{model}$ 个数。

$W_1$ 共 $d_{model}$ 行，每行 $d_{ff}$ 个数。下一行是 $b_1$。

$W_2$ 共 $d_{ff}$ 行，每行 $d_{model}$ 个数。下一行是 $b_2$。

然后四行依次是 $\mathrm{gamma}_1,\mathrm{beta}_1,\mathrm{gamma}_2,\mathrm{beta}_2$，每行 $d_{model}$ 个数。

最后一行一个浮点数 $\mathrm{eps}$。

约束

$1 \le B \le 4$

$1 \le T \le 48$

$1 \le d_{model} \le 32$，$h$ 整除 $d_{model}$

$1 \le d_{ff} \le 64$

$10^{-8} \le \mathrm{eps} \le 10^{-3}$

输入和权重的绝对值不超过 $5$。

输出描述

按 batch、再按时间步输出 $Y$，一共 $B \times T$ 行，每行 $d_{model}$ 个浮点数，保留 $6$ 位小数。

样例1

输入

1 2 2 2 2
1.00 2.00
3.00 0.00
1.00 0.00
0.00 1.00
1.00 0.00
0.00 1.00
1.00 0.00
0.00 1.00
1.00 0.00
0.00 1.00
1.00 0.00
0.00 1.00
0.00 0.00
1.00 0.00
0.00 1.00
0.00 0.00
1.00 1.00
0.00 0.00
1.00 1.00
0.00 0.00
1e-05

输出

-0.999998 0.999998
0.999998 -0.999998

说明

$B=1,T=2,d_{model}=2,h=2,d_{ff}=2$。四个注意力矩阵和两个前馈矩阵都是单位矩阵，偏置是 $0$，$\mathrm{gamma}$ 是 $1$，$\mathrm{beta}$ 是 $0$。$X$ 是 $[[1,2],[3,0]]$。每个头的维度是 $1$。注意力输出再走两次 Post-Norm 之后，得到上面两行。

样例2

输入

1 1 1 1 1
2.00
1.00
1.00
1.00
1.00
1.00
0.00
1.00
0.00
1.00
0.00
1.00
0.00
1e-05

输出

0.000000

说明

只有一个数。LayerNorm 时方差是 $0$，归一化结果是 $0$，后面的 ReLU 和第二次 LayerNorm 仍然是 $0$。