大语言模型原理（五）：Transformer 多头注意力、残差与归一化

1import numpy as np  # 只依赖 numpy，不用任何深度学习框架2 3# ========== 上一篇算出来的输入矩阵：词嵌入 + 位置编码 ==========4X = np.array([[3.0000, 2.0000],      # 女生5              [2.8415, 2.5403],      # 撩6              [2.9093, 3.5839]])     # 刘海7d_model = X.shape[1]                 # 模型维度，本例为 28n_heads = 3                          # 注意力头的个数9 10 11# ========== 基础算子 ==========12def softmax(scores):                                 # 按行做 softmax13    m = scores.max(axis=1, keepdims=True)            # 取每行最大值14    e = np.exp(scores - m)                           # 减最大值再取指数，防止 e^x 溢出，结果不变15    return e / e.sum(axis=1, keepdims=True)          # 除以行和，每行归一化成和为 1 的权重16 17 18def relu(x):                                         # ReLU 激活函数19    return np.maximum(x, 0.0)                        # 大于 0 的保留，小于 0 的置零，引入非线性20 21 22def layer_norm(x, gamma, beta, eps=1e-9):            # 层归一化，逐行独立计算23    mu = x.mean(axis=1, keepdims=True)               # 每行的均值 μ24    var = x.var(axis=1, keepdims=True)               # 每行的方差 σ²25    x_hat = (x - mu) / np.sqrt(var + eps)            # 标准化：均值拉到 0，方差拉到 126    return gamma * x_hat + beta                      # 再做可学的缩放 γ 与平移 β27 28 29# ========== 缩放点积注意力：单个头 ==========30def attention(X, W_Q, W_K, W_V):                     # 一个注意力头的完整计算31    Q = X @ W_Q                                      # Q：每个 token 想找什么信息32    K = X @ W_K                                      # K：每个 token 自己带着什么信息33    V = X @ W_V                                      # V：每个 token 实际要传出去的内容34    d_k = K.shape[1]                                 # 键向量的维度35    S = (Q @ K.T) / np.sqrt(d_k)                     # 点积打分再除以 √d_k，避免 softmax 饱和36    A = softmax(S)                                   # 分数转成每行和为 1 的注意力权重37    return A, A @ V                                  # 返回权重矩阵和加权求和的结果38 39 40# ========== 三个头各自的参数矩阵（训练前随机初始化） ==========41head_params = [42    (np.array([[-0.6, 0.9], [0.8, -0.5]]),           # head1 的 W_Q43     np.array([[0.6, 0.1], [-0.1, 1.4]]),            # head1 的 W_K44     np.array([[1.0, -0.5], [-0.4, 1.2]])),          # head1 的 W_V，与上一篇一致45    (np.array([[0.8, 1.0], [-1.2, -0.5]]),           # head2 的 W_Q46     np.array([[-0.6, -1.2], [-0.7, -0.2]]),         # head2 的 W_K47     np.array([[0.3, 1.1], [0.9, -0.2]])),           # head2 的 W_V48    (np.array([[-0.9, 0.3], [1.1, -1.3]]),           # head3 的 W_Q49     np.array([[0.3, -0.8], [0.7, 0.9]]),            # head3 的 W_K50     np.array([[-0.5, 0.7], [1.2, 0.4]])),           # head3 的 W_V51]52 53W_O = np.array([[0.20, -0.10],                       # 输出投影矩阵 W_O，把拼接后的 6 维压回 2 维54                [0.10, 0.30],55                [-0.20, 0.20],56                [0.30, 0.05],57                [0.10, -0.20],58                [0.15, 0.25]])59 60 61def multi_head_attention(X, head_params, W_O):       # 多头注意力62    heads, weights = [], []                          # 分别收集每个头的输出和注意力权重63    for W_Q, W_K, W_V in head_params:                # 逐个头计算，各头之间互不干扰64        A, head = attention(X, W_Q, W_K, W_V)        # 该头的注意力权重与输出65        weights.append(A)                            # 记下权重，便于观察各头关注点66        heads.append(head)                           # 记下输出，稍后拼接67    concat = np.concatenate(heads, axis=1)           # 沿维度方向横向拼接，3×2 × 3 → 3×668    return weights, concat, concat @ W_O             # 乘 W_O 混合各头信息并投影回 3×269 70 71# ========== 前馈网络的参数 ==========72W1 = np.array([[0.9, -0.7, 0.4, 1.2],                # 第一次线性变换，把 2 维升到 4 维73               [0.5, 0.8, -1.1, 0.3]])74b1 = np.array([0.1, 0.2, -0.3, 0.0])                 # 第一次线性变换的偏置75W2 = np.array([[0.6, -0.4],                          # 第二次线性变换，把 4 维降回 2 维76               [0.2, 0.9],77               [-0.5, 0.3],78               [0.7, 0.5]])79b2 = np.array([0.05, -0.10])                         # 第二次线性变换的偏置80 81 82def feed_forward(x, W1, b1, W2, b2):                 # 前馈网络，对每一行独立作用83    return relu(x @ W1 + b1) @ W2 + b2               # 升维 → ReLU → 降维84 85 86# ========== 一个完整的 Transformer Encoder 层 ==========87gamma = np.ones(d_model)                             # LayerNorm 的缩放参数，本例取全 188beta = np.zeros(d_model)                             # LayerNorm 的平移参数，本例取全 089 90 91def encoder_layer(X):                                # 完整一层：注意力 → 残差 → 归一化 → FFN → 残差 → 归一化92    weights, concat, mha = multi_head_attention(X, head_params, W_O)  # 多头注意力输出93    residual1 = X + mha                              # 第一次残差：保住进注意力之前的 X94    X1 = layer_norm(residual1, gamma, beta)          # 第一次层归一化，逐行拉成均值 0 方差 195    ffn = feed_forward(X1, W1, b1, W2, b2)           # 每个 token 独立过前馈网络96    residual2 = X1 + ffn                             # 第二次残差：保住进 FFN 之前的 X197    out = layer_norm(residual2, gamma, beta)         # 第二次层归一化，得到这一层的最终输出98    return weights, concat, mha, residual1, X1, ffn, residual2, out99 100 101weights, concat, mha, r1, X1, ffn, r2, out = encoder_layer(X)  # 跑一遍完整的一层102 103np.set_printoptions(precision=4, suppress=True)      # 打印保留 4 位小数，关掉科学计数法104for i, A in enumerate(weights):                      # 逐个打印每个头的注意力权重105    print(f"head{i + 1} 的注意力权重 A =\n{A}\n")106for name, mat in [("Concat", concat), ("MultiHead", mha),   # 依次打印每一步的中间矩阵107                  ("X + MultiHead", r1), ("LayerNorm1", X1),108                  ("FFN", ffn), ("X1 + FFN", r2), ("Output", out)]:109    print(f"{name} =\n{mat}\n")                      # 输出与正文里手算的数字一致