思路步骤

1. 缩放点积得分 $S=\frac{QK^\top}{\sqrt{d_k}}$，其中 $d_k=K$ 的列数。
2. 加掩码

题目内容

设 $d_k=d$ 。

1.缩放点积得分 $S=\frac{QK^T}{\sqrt{d_k}}$

2.加掩码(如果提供) $S←S+mask$

• 掩码元素为 $0$ 或 $-1e9$ ——等价于将被屏蔽位置的 $softmax$ 概率压至 $0$

3.行级数值稳定 $softmax$

$S=S-S.max(axis=1,keepdims=True)$

$P=exp(S)/exp(S).sum(axis=1,keepdims=True)$

4.加权求和 $O = PV$ ，结果形状 $(m,d_v)$

5.输出展开

• 按行优先顺序进行展平操作，即将 $O[0],O[1],...,O[m-1]$ 依次连接成一个一维数组

输入描述

单行 $JSON$：

{

"Q":[[...],...], //形状 $(m,d)$

"K":[[...],...], //形状 $(n,d)$

"V":[[...],...], //形状 $(n,d_v)$；本题 $d_v==d$

"mask":[[...],...] or null //形状 $(m,n)$ ；$0$ 表示可见，$-1e9$ 表示屏蔽

}

• $2≤d≤4, 1≤m,n≤3$

• 所有值为实数

• 若"$mask$"为 $null$ ，视为全 $0$

输出描述

使用 $round(x,6)$ 保留小数位即可，返回单行 $JSON$ 数组

补充说明

为了确保通过测试用例，仅允许使用 $Numpy$ 实现

样例1

输入

{"Q":[[1,0]],"K":[[1,0]],"V":[[5,5]],"mask":null}

输出

[5.0, 5.0]