共计 4651 个字符，预计需要花费 12 分钟才能阅读完成。
背景介绍
自注意力机制（Self-Attention）是现代自然语言处理（NLP）中的核心技术之一，尤其在 Transformer 架构中扮演着关键角色。它能够捕捉输入序列中各个元素之间的依赖关系，无论它们之间的距离有多远。这种能力使得 Transformer 在处理长距离依赖问题时表现出色。而多头自注意力（Multi-Head Self-Attention, MSA）则进一步扩展了这一机制，通过并行计算多个注意力头，模型能够从不同的子空间学习不同的表示，从而提升模型的表达能力。
数学推导
1. 基础注意力机制
自注意力的核心思想是通过计算查询（Query）、键（Key）和值（Value）之间的相似度来分配注意力权重。给定输入序列 $X \in \mathbb{R}^{n \times d}$，其中 $n$ 是序列长度，$d$ 是特征维度，自注意力机制的计算过程如下：
-
线性变换：将输入 $X$ 通过三个不同的线性变换矩阵 $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$ 分别映射到查询 $Q$、键 $K$ 和值 $V$：
$$ Q = X W_Q, \quad K = X W_K, \quad V = X W_V $$ -
注意力分数计算：通过点积计算查询和键之间的相似度，然后除以 $\sqrt{d_k}$ 进行缩放（防止梯度消失或爆炸），再通过 softmax 函数归一化得到注意力权重：
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V $$
2. 多头注意力机制
多头注意力将上述过程并行化，通过多个注意力头（通常为 $h$ 个）来捕捉不同的特征表示。每个头的计算过程如下：
-
分割线性变换：将输入 $X$ 分别映射到 $h$ 个查询、键和值子空间，每个子空间的维度为 $d_k = d / h$：
$$ Q_i = X W_{Q_i}, \quad K_i = X W_{K_i}, \quad V_i = X W_{V_i} $$ -
并行计算注意力：每个头独立计算注意力，然后将结果拼接起来：
$$ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) $$
$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W_O $$
其中 $W_O \in \mathbb{R}^{d \times d}$ 是输出线性变换矩阵。
3. (s)w-msa 的数学原理
(s)w-msa（Shifted Window Multi-Head Self-Attention）是 Swin Transformer 中提出的改进版多头注意力机制。它在计算注意力时引入了窗口划分和移位窗口的机制，从而在保持计算效率的同时捕捉全局上下文信息。其核心步骤如下：
- 窗口划分：将输入特征图划分为若干个不重叠的窗口，每个窗口内的像素只计算局部注意力。
- 移位窗口：在下一层中，窗口位置进行一定偏移，使得不同层的窗口能够覆盖不同的局部区域，从而间接实现全局信息的交互。
数学上，(s)w-msa 的计算过程与标准多头注意力类似，但注意力掩码会随着窗口的移位而变化，以限制注意力的计算范围。
代码实现
以下是一个基于 PyTorch 的 (s)w-msa 实现代码，包含详细注释和性能优化技巧：
import torch
import torch.nn as nn
import torch.nn.functional as F
class ShiftedWindowMSA(nn.Module):
def __init__(self, dim, window_size, num_heads, shift_size=0):
super().__init__()
self.dim = dim
self.window_size = window_size
self.num_heads = num_heads
self.shift_size = shift_size
self.scale = (dim // num_heads) ** -0.5
# 线性变换矩阵
self.qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
def forward(self, x):
B, H, W, C = x.shape
# 生成移位掩码
if self.shift_size > 0:
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
mask = self._create_mask(H, W)
else:
shifted_x = x
mask = None
# 划分窗口
x_windows = self._window_partition(shifted_x)
x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
# 计算 QKV
qkv = self.qkv(x_windows).view(-1, self.window_size * self.window_size, 3, self.num_heads, C // self.num_heads)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
# 计算注意力
attn = (q @ k.transpose(-2, -1)) * self.scale
if mask is not None:
attn = attn.view(-1, self.num_heads, self.window_size * self.window_size, self.window_size * self.window_size)
attn = attn + mask.unsqueeze(1)
attn = attn.view(-1, self.num_heads, self.window_size * self.window_size, self.window_size * self.window_size)
attn = F.softmax(attn, dim=-1)
out = (attn @ v).transpose(1, 2).reshape(-1, self.window_size * self.window_size, C)
# 合并窗口
out = self._window_reverse(out, H, W)
# 反向移位
if self.shift_size > 0:
out = torch.roll(out, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
# 输出投影
out = self.proj(out)
return out
def _window_partition(self, x):
B, H, W, C = x.shape
x = x.view(B, H // self.window_size, self.window_size, W // self.window_size, self.window_size, C)
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.window_size, self.window_size, C)
return windows
def _window_reverse(self, windows, H, W):
B = int(windows.shape[0] / (H * W / self.window_size / self.window_size))
x = windows.view(B, H // self.window_size, W // self.window_size, self.window_size, self.window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return x
def _create_mask(self, H, W):
img_mask = torch.zeros((1, H, W, 1))
h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None))
cnt = 0
for h in h_slices:
for w in w_slices:
img_mask[:, h, w, :] = cnt
cnt += 1
mask_windows = self._window_partition(img_mask)
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
return attn_mask性能分析
在实现多头自注意力时，性能优化是一个重要考虑因素。以下是几种常见实现方式的性能对比：
- 标准实现 ：直接计算全序列的注意力，复杂度为 $O(n^2)$，适用于短序列。
- 窗口划分 ：将序列划分为多个窗口，只在窗口内计算注意力，复杂度降低到 $O(n \cdot w^2)$，其中 $w$ 是窗口大小。
- 移位窗口 ：通过窗口移位实现不同层间的信息交互，复杂度与窗口划分相同，但能捕捉更全局的信息。
实验表明，窗口划分和移位窗口的实现在大规模数据集上能够显著减少计算量，同时保持模型的性能。
避坑指南
在实际实现多头自注意力时，开发者可能会遇到以下常见问题：
- 维度不匹配 ：在多头分割时，确保特征维度 $d$ 能够被头数 $h$ 整除，否则会导致维度不匹配的错误。
- 梯度消失或爆炸 ：注意在计算注意力分数时进行缩放（除以 $\sqrt{d_k}$），否则可能导致梯度问题。
- 掩码处理不当 ：在移位窗口的实现中，掩码的生成和处理是关键，错误的掩码会导致注意力计算的范围不符合预期。
- 计算效率低 ：避免在循环中逐头计算注意力，尽量使用矩阵运算并行化计算。
进阶思考
多头自注意力机制不仅适用于 NLP 领域，还可以扩展到其他领域：
- 计算机视觉 ：通过将图像划分为 patch 序列，自注意力机制可以用于图像分类、目标检测等任务。
- 时间序列预测 ：自注意力能够捕捉时间序列中的长距离依赖关系，适用于股票预测、天气预测等场景。
- 多模态学习 ：结合文本、图像、音频等多种模态的数据，自注意力机制可以实现跨模态的信息交互。
启发问题
- 如何在保持模型性能的同时进一步降低多头自注意力的计算复杂度？
- 除了窗口划分和移位窗口，还有哪些方法可以提升自注意力机制在长序列上的效率？
- 如何将多头自注意力机制与其他神经网络结构（如卷积、循环网络）结合，以发挥各自的优势？
通过本文的讲解和代码实现，相信读者已经对 (s)w-msa 多头自注意力机制有了深入的理解。在实际应用中，可以根据具体任务的需求灵活调整窗口大小、头数等超参数，以达到最佳的性能和效率平衡。
