题目内容

标准的 $Transformer$ 自注意力计算为：

对于序列长度为 $N$、维度为 $d_k$ 的输入，其时间复杂度和空间复杂度分别为 $O(N^2 \cdot d_k)$ 和 $O(N^2)$。

近期研究表明，大多数任务中的依赖关系具有局部性，远距离 token 通常可以通过多层堆叠间接完成交互。因此，可以通过限制每个 token 的注意力范围来显著降低计算复杂度。

滑动窗口注意力机制是一种典型方法，其原理如下：

1. 为每个 $token$ 设定固定大小为 $w$ 的对称窗口。
2. 位置 $i$ 的 $query$ 只与满足 $|i-j| \le w$ 的位置 $j$ 的 $key$ 计算注意力，其中 $w \ll N$。
3. 为保证数值稳定性，在计算 $Softmax$ 前减去当前窗口内的最大值：

以 $(N=8,w=1)$ 为例，滑动窗口注意力可视化如下：

位置：$0$ $1$ $2$ $3$ $4$ $5$ $6$ $7$

$Q0: [x] [x] [] [] [] [] [] [] \rightarrow$ 窗口 $[0,1]$，左边界限制为 $0$

$Q1: [x] [x] [x] [] [] [] [] [] \rightarrow$ 窗口 $[0,1,2]$

$Q2: [] [x] [x] [x] [] [] [] [] \rightarrow$ 窗口 $[1,2,3]$

$Q3: [] [] [x] [x] [x] [] [] [] \rightarrow$ 窗口 $[2,3,4]$

$Q4: [] [] [] [x] [x] [x] [] [] \rightarrow$ 窗口 $[3,4,5]$

$Q5: [] [] [] [] [x] [x] [x] [] \rightarrow$ 窗口 $[4,5,6]$

$Q6: [] [] [] [] [] [x] [x] [x] \rightarrow$ 窗口 $[5,6,7]$

$Q7: [] [] [] [] [] [] [x] [x] \rightarrow$ 窗口 $[6,7]$，右边界限制为 $7$

其中 $[x] =$ 计算注意力，$[] =$ 掩码为 $0$（不计算）。

输入描述

第一行输入四个整数：

$batch_size$, $seq_len$, $d_k$, $window_size$

后续包含 $batch_size$ 组数据，每组依次包含：

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Query 矩阵。

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Key 矩阵。

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Value 矩阵。

约束

$batch_size \ge 1$

$seq_len \ge 1$

$d_k \ge 1$

$1\le window_size \le seq_len$

输出描述

输出 $batch_size$ 组结果，每组包含 $seq_len$ 行，每行 $d_k$ 个浮点数，表示该位置的输出向量。

组间用一个空行分隔，最后一组后无空行。

最终结果保留 $2$ 位小数，四舍五入。

注意力打分及 $Softmax$ 权重在题目说明中保留 $6$ 位小数展示，实际计算过程中无需提前进行舍入。

如果输入违反约束，输出 $0$。

样例1

输入

1 4 4 1
0.1 0.2 0.3 0.4
0.2 0.3 0.4 0.5
0.3 0.4 0.5 0.6
0.4 0.5 0.6 0.7
0.1 0.0 0.1 0.0
0.0 0.1 0.0 0.1
0.1 0.1 0.0 0.0
0.0 0.0 0.1 0.1
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.50 0.50 0.00 0.00
0.33 0.34 0.33 0.00
0.00 0.33 0.33 0.34
0.00 0.00 0.50 0.50

说明

输入 $batch_size=1$，$seq_len=4$，$d_k=4$，$window_size=1$。

缩放因子 $scale=\sqrt{d_k}=2$。

根据窗口大小 $window_size=1$，注意力打分为：

位置 $0$ 的权重：

$score[0]=[Q[0]K[0],Q[0]K[1]]/scale=[0.020000,0.030000]$

归一化后为：

$[0.497500,0.502500]$

位置 $1$ 的权重：

$score[1]=[Q[1]K[0],Q[1]K[1],Q[1]K[2]]/scale=[0.030000,0.040000,0.025000]$

归一化后为：

$[0.332772,0.336116,0.331112]$

位置 $2$ 的权重：

$score[2]=[Q[2]K[1],Q[2]K[2],Q[2]K[3]]/scale=[0.050000,0.035000,0.055000]$

归一化后为：

$[0.334434,0.329455,0.336111]$

位置 $3$ 的权重：

$score[3]=[Q[3]K[2],Q[3]K[3]]/scale=[0.045000,0.060000]$

归一化后为：

$[0.496250,0.503750]$

上述归一化权重保留 $6$ 位小数精度。

最后使用归一化后的权重对对应位置的 $V$ 进行加权求和，并将最终结果保留 $2$ 位小数：

$output[0]=0.497500\times V[0]+0.502500\times V[1]$

$=[0.50,0.50,0.00,0.00]$

$output[1]=0.332772\times V[0]+0.336116\times V[1]+0.331112\times V[2]$

$=[0.33,0.34,0.33,0.00]$

$output[2]=0.334434\times V[1]+0.329455\times V[2]+0.336111\times V[3]$

$=[0.00,0.33,0.33,0.34]$

$output[3]=0.496250\times V[2]+0.503750\times V[3]$

$=[0.00,0.00,0.50,0.50]$

样例2

输入

1 2 2 0
0.3 0.4
0.4 0.5
0.5 0.6
0.6 0.7
0.1 0.0
0.0 0.1

输出

0

说明

输入 $batch_size=1$，$seq_len=2$，$d_k=2$，$window_size=0$。

由于 $window_size=0$ 违反输入范围要求，因此输出 $0$。

提示

仅能使用编程语言内置的函数。