【AI 工业核题 D2】Multi-Head Attention 多头注意力（MHA 分头与拼接）（Multi-Head Attention (MHA)）深度实现与原理解析

题目分类：Part D · 注意力机制与 Transformer 核心组件 (Part D · Attention Mechanisms & Transformer Blocks) | 难度等级：Medium | 工业重要度：工业基石 (核心高频)

一、核心题意与背景

输入线性映射切分多头、多子空间并行计算、拼接输出残差降维，Transformer 骨架核心。

Industrial-grade implementation and mathematical foundations of Multi-Head Attention (MHA).

二、数学原理与公式推导

多子空间联合注意力表征

单头注意力会将所有特征压制在一个共同的空间中求平均；而多头注意力（MHA）通过 $h$ 个独立的线性投射矩阵，将模型容量投影到不同的低维子空间：
$$mathrm{head}_i = mathrm{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$$
每个头的维度 $d_k = D / h$。这使得不同的头能够各自专注于语法依赖、实体共指、长程因果等不同视角的模式。
最后将各个头的输出在最后一维拼接，并通过输出投影矩阵 $W^O in mathbb{R}^{D times D}$ 融合特征。

三、工业级 Python 核心实现

import numpy as np

class MultiHeadAttention:
    def __init__(self, d_model: int, num_heads: int):
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 权重参数 (为简化展示采用随机初始化)
        self.W_q = np.random.randn(d_model, d_model) * 0.02
        self.W_k = np.random.randn(d_model, d_model) * 0.02
        self.W_v = np.random.randn(d_model, d_model) * 0.02
        self.W_o = np.random.randn(d_model, d_model) * 0.02

    def forward(self, x: np.ndarray, is_causal: bool = True) -> np.ndarray:
        """
        x: (B, S, D)
        """
        B, S, D = x.shape
        # 1. 线性投影: (B, S, D) @ (D, D) -> (B, S, D)
        Q = x @ self.W_q
        K = x @ self.W_k
        V = x @ self.W_v

        # 2. 分头与维度置换: (B, S, H, D_k) -> (B, H, S, D_k)
        Q = Q.reshape(B, S, self.num_heads, self.d_k).swapaxes(1, 2)
        K = K.reshape(B, S, self.num_heads, self.d_k).swapaxes(1, 2)
        V = V.reshape(B, S, self.num_heads, self.d_k).swapaxes(1, 2)

        # 3. 批量多头缩放点积
        scores = np.matmul(Q, K.swapaxes(-1, -2)) / np.sqrt(self.d_k)
        if is_causal:
            mask = np.triu(np.ones((S, S), dtype=bool), k=1)
            scores = np.where(mask, -1e9, scores)

        scores_max = np.max(scores, axis=-1, keepdims=True)
        exp_s = np.exp(scores - scores_max)
        attn_probs = exp_s / np.sum(exp_s, axis=-1, keepdims=True)
        context = np.matmul(attn_probs, V)  # (B, H, S, D_k)

        # 4. 拼接多头并输出投影
        # 置换回 (B, S, H, D_k) 并平铺为 (B, S, D)
        context = context.swapaxes(1, 2).reshape(B, S, D)
        return context @ self.W_o

四、自动化单元测试与边界断言

import numpy as np
mha = MultiHeadAttention(d_model=16, num_heads=4)
x = np.random.randn(2, 6, 16)
out = mha.forward(x, is_causal=True)
assert out.shape == (2, 6, 16), "输出形状错误"
assert not np.isnan(out).any(), "包含 NaN"
print("✓ MultiHeadAttention 自测通过")

五、张量形状与维度变换流 (Tensor Flow)

• 中文解析：(B, S, D) -> 投影为 Q,K,V -> (B, S, H, D_k) -> swapaxes -> (B, H, S, D_k) -> Attention -> (B, H, S, D_k) -> swapaxes & reshape -> (B, S, D) -> W_o 投影 -> (B, S, D)
• 英文对齐：(B, S, D) -> 投影为 Q,K,V -> (B, S, H, D_k) -> swapaxes -> (B, H, S, D_k) -> Attention -> (B, H, S, D_k) -> swapaxes & reshape -> (B, S, D) -> W_o 投影 -> (B, S, D)

六、工业级数值稳定性避坑清单 (Checklist)

• ⚠️ 分头置换后，在恢复连续内存前必须保证 reshape 维度对齐，在 PyTorch 中需显式调用 .contiguous()
• ⚠️ 工业实现通常将 W_q, W_k, W_v 拼接为单一大矩阵 (D, 3*D) 执行单次 GEMM 提高吞吐

English Checklist:
– Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
– Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
– Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

七、考场秒记心法口诀

💡 投影拆头换轴线，点积加权分头看，换回轴线拼齐整，输出矩阵定乾坤

Master Multi-Head Attention (MHA): enforce numerical stability, check tensor shapes, and eliminate redundant memory allocations.

八、高频面试追问与答题策略

Q1：如果给定批次中序列长度极不平衡（如一个样本长 2048，另一个长 16），MHA 的 Padding 会造成怎样的算力浪费？工业界如何优化？
(EN: What are the key trade-offs and memory bottlenecks when deploying Multi-Head Attention (MHA) in high-throughput inference?)

答：Padding 会引入大量无用的全 0 甚至无意义的点积掩码计算，造成高达数倍的无效显存与矩阵乘法开销。工业界现代解决方案是使用 FlashAttention 团队提出的 Variable-length (FlashAttention-varlen) 或 Packed Sequences，将所有有效 Token 拼接成一维长张量，使用 cu_seqlens 数组记录边界，彻底消除所有 Padding 浪费。

(EN: Memory bandwidth (HBM to SRAM I/O) is the primary latency factor. Fusing element-wise operations and avoiding intermediate tensor materialization significantly outperforms naive implementations.)