解题思路

本题要求手写「带因果掩码」的多头自注意力（Decoder常用），输入为：

• num_heads
• X：形状 [batch_size, seq_len, d_model]
• Q, K, V, W_O：均为形状 [d_model, d_model] 的线性投影矩阵（对应 $W_Q, W_K, W_V, W_O$）

整体流程（Scaled Dot-Product Attention + 多头并行）：

题目内容

在Transformer模型中，Multi-Head Self-Attention是核心组件，用于捕捉序列中的依赖关系。你需要从头实现一个Masked Multi-Head Self-Attention函数，支持自注意力（即queries、keys和values来自同一输入序列），并处理编码（mask）以防止未来位置的信息泄露（常见于Decoder中）。

具体要求：

• 支持多头注意力：将注意力机制并行分成多个"头"，每个头学习不同的注意力模式，增强模型对多维度特征的捕捉能力。
• 计算过程：
  1. 生成Q、K、V矩阵 对输入序列X（维度：[batch_size, seq_len, d_model]）通过3个线性层分别生成查询（Query, Q）、键（Key, K）、值（Value, V）矩阵：（$Q = X \cdot W_Q$，$K = X \cdot W_K$，$V = X \cdot W_V$），其中 $W_Q, W_K, W_V \in \mathbb{R}^{d_{model} \times d_{model}}$。
  2. 将Q、K、V拆分为多个头 将Q、K、V分割为num_heads个并行的子矩阵（每个头的维度为d_k = d_model / num_heads）。 分割后维度为[batch_size, num_heads, seq_len, d_k]。
  3. 对于每个头，计算注意力分数：attention_scores = ( $Q \cdot K^T$ ) / sqrt(d_k)。
  4. 提供mask（一个(batch_size, seq_len, seq_len)的布尔数组，其中True表示需要掩码的位置），则将masked位置的注意力分数设置为负无穷（-inf），以确保softmax后为0。掩码后的分数为masked_scores。
  5. 对掩码后的分数应用softmax得到注意力权重。 softmax_scores=softmax(masked_scores)。
  6. 计算注意力输出：attention=softmax_scores · V。
  7. 拼接多头输出，并通过一个线性投影得到最终结果。 $output = concat(attention_1, ..., attention_{num_heads}) · W_O$ ，其中 $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ 是可学习参数，输出维度为 [batch_size, seq_len, d_model].

注意： 1、需处理批次（batch_size > 1）和变长序列。

2、输入参数以分号分隔。第一个参数为多头数量num_heads； 第二个参数为Q矩阵；第三个参数为K矩阵；第四个参数为V矩阵；第五个参数为 $W_O$。

3、输出为List，需要将np.ndarray转为List

输入描述

以";"分隔，分别为 num_heads, X, Q、K、V，$W_O$

输出描述

输出为最终结果 $output$，输出保留两位有效小数，并且为 List。

样例1

输入

2;[[[ 1.92, 1.48], [0.67, -1.23], [0.35, -0.68]], [[-1.11, 0.09], [-0.3, -0.39], [-0.59, -0.06]]];[[1.0, 2.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]]

输出

[[[14.64, 14.64], [-5.36, -5.36], [-4.44, -4.44]], [[-2.79, -2.79], [-3.04, -3.04], [-2.79, -2.79]]]

样例2

输入

2;[[[ 1.92, 1.48], [0.67, -1.23], [0.35, -0.68]], [[-1.11, 0.09], [-0.3, -0.39], [-0.59, -0.06]]];[[1.0,1.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]];[[1.0, 1.0], [2.0, 2.0]]

输出

[[[14.64, 14.64], [-5.37, -5.37], [-4.62, -4.62]], [[-2.79, -2.79], [-3.03, -3.03], [-2.77, -2.77]]]

提示

1. 手动实现softmax：exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))；softmax = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)。确保数值稳定性，减去每行最大值
2. 使用np.around(np.ndarray, 2)将输出保留2位小数
3. 通过下三角矩阵实现序列掩码mask，确保每个位置只能关注自身及之前的位置。下三角为1，上三角为0。
4. 处理-inf：可以使用np.where(mask == 0, -np.inf, attention_scores)