Multi-Head Self-Attention

分类: 深度学习基础

定义

多头自注意力是 Transformer 架构的核心组件，将输入投影到多个子空间中分别计算注意力，然后拼接输出，使模型能同时关注不同表示子空间的信息。

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O

\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V) = \text{Softmax}\left(\frac{Q W_i^Q (K W_i^K)^\top}{\sqrt{d_k}}\right) V W_i^V

$h$ : 注意力头数

$d_k = d_{\text{model}} / h$ : 每个头的维度

$W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_k}$ : 投影矩阵

计算复杂度: $O(N^2 \cdot d)$ ， $N$ 为序列长度

多头机制允许模型在不同位置关注不同的语义信息

计算复杂度与序列长度 $N$ 呈二次关系，是 Transformer 效率瓶颈的主要来源

Token Pruning、FlashAttention 等方法旨在缓解这一计算瓶颈

Vaswani et al. (2017): Attention Is All You Need，原始 Transformer

ViT: 将 MHSA 应用于视觉任务

Col-Ln: 利用 MHSA 的注意力矩阵列向分布来度量 token 重要性