算法步骤

1. 输入转为 NumPy 数组

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

   • 记：

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

2. 线性变换得到 Q, K, V

   • 使用矩阵乘法 + 广播加偏置：

     Q = X_np @ W_Q_np + b_Q_np
     K = X_np @ W_K_np + b_K_np
     V = X_np @ W_V_np + b_V_np
     

   • 由于 b_Q_np 形状为 (d_model,)，NumPy 会自动在第 0 维做广播，加到每一行上。

3. 按 head 切分 Q, K, V

   • 先计算：

     T, d_model = X_np.shape
     d_k = d_model // h
     

   • 将 Q 重塑为 (T, h, d_k)，再转置为 (h, T, d_k)：

     Q_heads = Q.reshape(T, h, d_k).transpose(1, 0, 2)
     K_heads = K.reshape(T, h, d_k).transpose(1, 0, 2)
     V_heads = V.reshape(T, h, d_k).transpose(1, 0, 2)
     

   • 这样第 0 维就是 head 维度，可以批量计算每个 head 的注意力。

4. 计算 Scaled Dot-Product Attention

   • 对所有 head 一次完成：

     1. 打分矩阵 scores

        scale = 1.0 / math.sqrt(d_k)
        scores = np.matmul(Q_heads, K_heads.transpose(0, 2, 1)) * scale
        

        • Q_heads 形状 (h, T, d_k)
        • K_heads.transpose(0, 2, 1) 形状 (h, d_k, T)
        • scores 形状 (h, T, T)，每个 head 一个 T × T 的分数矩阵。

     2. 对最后一维做数值稳定的 softmax

        scores_max = np.max(scores, axis=-1, keepdims=True)
        exp_scores = np.exp(scores - scores_max)
        sum_exp = np.sum(exp_scores, axis=-1, keepdims=True)
        sum_exp = np.where(sum_exp == 0.0, 1.0, sum_exp)  # 防止除零
        attn_weights = exp_scores / sum_exp
        

        • attn_weights 形状依然 (h, T, T)，每行都是一个概率分布。

5. 用注意力权重加权 V 得到各 head 输出

   • 对所有 head 批量计算：

     H_heads = np.matmul(attn_weights, V_heads)  # (h, T, d_k)
     

   • 每个 head 一个 T × d_k 的输出。

6. 拼接所有 heads 得到 H

   • 把 (h, T, d_k) 变回 (T, d_model)：

     H = H_heads.transpose(1, 0, 2).reshape(T, d_model)
     

   • 先把维度换成 (T, h, d_k)，再把后两维合并。

7. 输出层线性映射得到 O

   • 计算：

     O = H @ W_O_np + b_O_np
     

   • 形状为 (T, d_model)。

   • 最后用 .tolist() 转回 Python 嵌套列表作为最终返回值。

Python 代码实现

import math
from typing import List
import numpy as np


class Solution:
    def multi_head_attention(
        self,
        X: List[List[float]],
        W_Q: List[List[float]],
        b_Q: List[float],
        W_K: List[List[float]],
        b_K: List[float],
        W_V: List[List[float]],
        b_V: List[float],
        h: int,
        W_O: List[List[float]],
        b_O: List[float],
    ) -> List[List[float]]:
        """使用 NumPy 实现简化版 Multi-Head Attention，并返回最终输出矩阵 O。"""

        # 将所有输入转为 NumPy 数组，统一为 float 类型
        X_np = np.asarray(X, dtype=float)
        W_Q_np = np.asarray(W_Q, dtype=float)
        b_Q_np = np.asarray(b_Q, dtype=float)
        W_K_np = np.asarray(W_K, dtype=float)
        b_K_np = np.asarray(b_K, dtype=float)
        W_V_np = np.asarray(W_V, dtype=float)
        b_V_np = np.asarray(b_V, dtype=float)
        W_O_np = np.asarray(W_O, dtype=float)
        b_O_np = np.asarray(b_O, dtype=float)

        # 基本维度信息
        T, d_model = X_np.shape
        d_k = d_model // h

        # ============ 第一步：线性变换得到 Q, K, V ============
        # 利用 NumPy 广播机制自动加偏置
        Q = X_np @ W_Q_np + b_Q_np      # 形状 (T, d_model)
        K = X_np @ W_K_np + b_K_np      # 形状 (T, d_model)
        V = X_np @ W_V_np + b_V_np      # 形状 (T, d_model)

        # ============ 第二步：按 head 切分成多头形式 ============
        # 先 reshape 成 (T, h, d_k)，再转置成 (h, T, d_k)
        Q_heads = Q.reshape(T, h, d_k).transpose(1, 0, 2)  # (h, T, d_k)
        K_heads = K.reshape(T, h, d_k).transpose(1, 0, 2)  # (h, T, d_k)
        V_heads = V.reshape(T, h, d_k).transpose(1, 0, 2)  # (h, T, d_k)

        # ============ 第三步：计算 Scaled Dot-Product Attention ============
        # scores = Q_i @ K_i^T / sqrt(d_k)，对所有 head 一次性计算
        scale = 1.0 / math.sqrt(d_k)
        scores = np.matmul(Q_heads, K_heads.transpose(0, 2, 1)) * scale  # (h, T, T)

        # 数值稳定的 softmax：沿最后一维做 softmax
        scores_max = np.max(scores, axis=-1, keepdims=True)  # 每行减去最大值
        exp_scores = np.exp(scores - scores_max)
        sum_exp = np.sum(exp_scores, axis=-1, keepdims=True)

        # 防止极端情况下出现除以 0 的情况
        sum_exp = np.where(sum_exp == 0.0, 1.0, sum_exp)
        attn_weights = exp_scores / sum_exp  # (h, T, T)

        # ============ 第四步：用注意力权重加权 V，得到每个 head 的输出 ============
        # H_heads = attn_weights @ V_heads，形状 (h, T, d_k)
        H_heads = np.matmul(attn_weights, V_heads)

        # ============ 第五步：拼接所有 head，恢复到 (T, d_model) ============
        # 先转置成 (T, h, d_k)，再 reshape 合并后两维
        H = H_heads.transpose(1, 0, 2).reshape(T, d_model)

        # ============ 第六步：输出层线性映射 ============
        O = H @ W_O_np + b_O_np  # 形状 (T, d_model)

        # 返回 Python 原生嵌套列表
        return O.tolist()

题目描述

实现一个简化版的 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$ 可为任意实数