解题思路

题目要求实现一种新的注意力机制 $TG\text{-}SA$，整体流程可以直接按题意模拟，核心涉及两个算法步骤：

1. 矩阵乘法计算注意力原始分数
2. 按行进行双门控分段映射与 $Top\text{-}k$ 保留

设输入矩阵为 $Q,K,V \in \mathbb{R}^{L \times d}$。

题目内容

在仅使用 $numpy/pandas/scikit-learn$ 的前提下，实现一种新型注意力机制 $TG-SA$，并完成给定序列的加权求和。与经典 $\mathrm{Q} \cdot \mathrm{~K}^{\mathrm{T}} / \sqrt{\mathrm{d}} \text { Softmax }$ 不同，$TG-SA$ 通过双门控转折 $+Top-k$ 保留得到更稀疏、可解释的权重。

给定单头 $Query/Key/Value$ 张量

$Q, K, V \in \mathbb{R}^{L \times d}$ (长度 $L$、隐维 $d$)，以及超参

$k\_top$：每行保留前 $k\_ top$ 最大注意力得分

$α$：第 $1$ 门限系数 $(0-1)$

$β$：第 $2$ 门限系数 $(0-1)$，且 $β>α$

默认 $α=0.2,β=0.6;k\_top≤L$。

计算过程：

$1.$ 初始打分：$S=QK^T\sqrt{\mathrm{d}}(L×L)$

$2.$ 双门控折线映射

其中 $s_{\max }=\max _{j} S_{i j}$ 为行内最大。

$3.$ $Top-k$ 稀疏

每行仅保留 $k\_top$ 个 $\tilde{S}_{i j}$ 最大位置；其余设 $0$。若行长 $<k\_top$，则全保留。

$4.$ 行归一化

若行全 $0$，则改为均匀分布；否则归一化到和为 $1$ ：$A_{i j}=\tilde{S}_{i j} / \sum_{j} \tilde{S}_{i j}$。

$5.$ 输出

$O=AV (L×d)$

同时返回注意力权重矩阵 $A$ 。

输入描述

单行 JSON

{

"Q":[[...],...], //L×d

"K":[[...],...],

"V":[[...],...],

"k_top":2, //2≤k_top≤L

"alpha":0.2, //可缺省

"beta":0.6 //可缺省，且＞alpha

}

约束：$2≤L≤6,2≤d≤4$；所有元素为浮点。

输出描述

仅一行 JSON

{

"A":[[a11,...],...], //L×L，6位小数，便用round(x,6)即可

"O":[[o11,...],...] //L×d，6位小数，便用round(x,6)即可

}

补充说明

为确保通过测试用例，仅允许便用 $numpy/pandas/sckit-learn$。

样例1

输入

{"Q":[[1,0],[0,1]],"K":[[1,0],[0,1]],"V":[[1,2],[3,4]],"k_top":1}

输出

{"A":[[1.0,0.0],[0.0,1.0]],"O":[[1.0,2.0],[3.0,4.0]]}