算法步骤

1. 输入转为 NumPy 数组

   • 将所有输入参数 X, W_Q, b_Q, ... 转为 np.ndarray，统一为 float 类型，方便后续矩阵运算和广播。

   • 记：

     • X_np.shape = (T, d_model)
     • W_Q_np.shape = (d_model, d_model) 等等。

题目描述

实现一个简化版的 Multi-Head Attention，具体步骤如下：

1. 将输入序列

   $$X \in \mathbb{R}^{T \times d_{\text{model}}}$$

   经线性变换生成

   $$Q = X W_Q + b_Q,\quad K = X W_K + b_K,\quad V = X W_V + b_V$$

2. 将 $Q, K, V$ 按 head 数 $h$ 切分为维度

   $$d_k = \frac{d_{\text{model}}}{h}$$

   的子矩阵。

3. 对第 $i$ 个 head 计算 Scaled Dot-Product Attention：

   $$\text{Attention}_i(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^\top}{\sqrt{d_k}}\right) V_i$$

4. 将所有 head 输出拼接：

   $$H = \text{concat}(H_1, H_2, \dots, H_h)$$

5. 输出层线性映射：

   $$O = H W_O + b_O$$

输入参数

• X：形状为 $T \times d_{\text{model}}$ 的输入序列
• W_Q, W_K, W_V：形状为 $d_{\text{model}} \times d_{\text{model}}$ 的线性层权重
• b_Q, b_K, b_V：长度为 $d_{\text{model}}$ 的偏置项
• h：Multi-head attention 的 head 数量
• W_O：形状为 $d_{\text{model}} \times d_{\text{model}}$ 的输出层权重
• b_O：长度为 $d_{\text{model}}$ 的输出偏置

返回值

• O：形状为 $T \times d_{\text{model}}$ 的最终注意力输出矩阵

示例

输入：

X =
[[1, 2, 0, 1],
 [0, 1, 1, 0]]

W_Q =
[[1, 0, 1, 0],
 [0, 1, 0, 1],
 [1, -1, 0, 0],
 [0, 0, 1, -1]]

b_Q = [0, 0, 0, 0]

W_K =
[[1, 1, 0, 0],
 [0, 1, 1, 0],
 [1, 0, 0, 1],
 [0, 0, 1, 1]]

b_K = [0, 0, 0, 0]

W_V =
[[1, 0, 0, 1],
 [0, 1, 1, 0],
 [1, 0, 1, 0],
 [0, 1, 0, 1]]

b_V = [0, 0, 0, 0]

W_O =
[[1, 0, 1, 0],
 [0, 1, 0, 1],
 [1, 0, 0, 1],
 [0, 1, 1, 0]]

b_O = [0, 0, 0, 0]

输出：

O = [[3.0, 4.7768, 2.8884, 4.8884], [3.0, 3.0, 2.0, 4.0]]

提示

• 输入范围： $-1000 \le X[i,j] \le 1000$

• 权重范围： $-10 \le W_Q,; W_K,; W_V,; W_O \le 10$

• 偏置范围： $-5 \le b_Q,; b_K,; b_V,; b_O \le 5$

• head 数要求： $1 \le h \le d_{\text{model}}$，且 $d_{\text{model}} \bmod h = 0$

• softmax 输出满足范围与归一性： $0 \le \text{softmax}(z_i) \le 1$，且 $\sum_i \text{softmax}(z_i) = 1$

• 最终输出矩阵 $O$ 可为任意实数