从零推导(s)w-msa多头自注意力公式：新手友好的数学原理与实现指南

1次阅读

没有评论

共计 1751 个字符，预计需要花费 5 分钟才能阅读完成。

自注意力机制的核心是计算查询向量（query）与键向量（key）的相似度。我们先从最简单的标量点积开始：

给定两个向量 $q_i \in \mathbb{R}^d$ 和 $k_j \in \mathbb{R}^d$，其点积表示为：
$$ \text{score}(q_i, k_j) = q_i \cdot k_j^T $$

这个分数越大，表示两个向量的相关性越强。在实际实现中，我们会同时计算所有位置的注意力分数：

$$ A = QK^T \quad \text{其中} \quad Q,K \in \mathbb{R}^{n \times d} $$

输入投影 ：首先将输入 $X \in \mathbb{R}^{n \times d_{model}}$ 投影到查询、键、值空间：
$$ Q = XW_Q, \quad K = XW_K, \quad V = XW_V $$
缩放点积 ：为防止点积结果过大导致 softmax 梯度消失，加入缩放因子：
$$ A = \frac{QK^T}{\sqrt{d_k}} $$
softmax 归一化 ：得到注意力权重矩阵：
$$ \text{Attn} = \text{softmax}(A) $$
值加权求和 ：最终输出是值的加权和：
$$ \text{Output} = \text{Attn} \cdot V $$

标准自注意力的计算复杂度是 $O(n^2)$，当序列长度 $n$ 很大时（如 2048），这会带来巨大计算负担。滑动窗口注意力通过限制每个位置只关注局部邻域来降低复杂度：

窗口划分 ：将输入序列划分为不重叠的窗口（通常 $w=7$ 或 $w=8$）
局部计算 ：只在每个窗口内计算自注意力
复杂度分析 ：计算量从 $O(n^2)$ 降低到 $O(n \times w^2)$

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwinMSA(nn.Module):
    def __init__(self, dim, num_heads=8, window_size=7):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.window_size = window_size

        # 投影矩阵
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
        q, k, v = qkv.unbind(2)  # [B, N, num_heads, head_dim]

        # 划分窗口
        q = self.window_partition(q, self.window_size)
        k = self.window_partition(k, self.window_size)
        v = self.window_partition(v, self.window_size)

        # 缩放点积注意力
        attn = (q @ k.transpose(-2, -1)) * (1.0 / (C ** 0.5))
        if mask is not None:
            attn = attn + mask
        attn = attn.softmax(dim=-1)
        x = (attn @ v)

        # 合并窗口
        x = self.window_reverse(x, self.window_size, N)
        x = self.proj(x)
        return x