3D自注意力机制深度解析：从数学原理到PyTorch实现

1次阅读

没有评论

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

传统 Transformer 的 2D 自注意力在处理视频、医疗影像等三维数据时会遇到两个致命问题：

计算复杂度爆炸 ：假设视频尺寸为 T×H×W，传统自注意力复杂度是 O(T²H²W²)。一个 112×112×16 的输入就会产生 3.2e10 次计算！
时空关系割裂 ：将三维数据展平为一维序列会破坏时空局部性，比如相邻帧的像素在序列中可能相隔数万位置

举个实际案例：在手术视频分析中，2D 注意力模型会消耗 48GB 显存处理 10 秒片段，而 3D 优化版本仅需 7GB。

3D 相对位置编码的核心思想是解耦时空维度。给定两个位置 i =(t,h,w) 和 j =(t’,h’,w’)，其位置关系可表示为：

$$e_{ij} = w_t \cdot \phi_t(t-t’) + w_h \cdot \phi_h(h-h’) + w_w \cdot \phi_w(w-w’)$$

其中：

$\phi_\cdot$ 是各维度的位置编码函数（常用正弦函数）
$w_\cdot$ 是可学习的维度权重

这种分解带来两个优势：

参数量从 O(T²H²W²) 降到 O(T²+H²+W²)
可以单独调整时空维度的关注强度

import torch
import einops

class RelativePosition3D(nn.Module):
    def __init__(self, max_len=32, dim=256):
        super().__init__()
        # 分别初始化三个维度的编码表
        self.emb_t = nn.Parameter(torch.randn(2*max_len+1, dim//3))
        self.emb_h = nn.Parameter(torch.randn(2*max_len+1, dim//3))
        self.emb_w = nn.Parameter(torch.randn(2*max_len+1, dim//3))

    def forward(self, q):
        b, t, h, w, c = q.shape
        # 生成相对位置索引
        idx_t = torch.arange(t).view(1,-1) - torch.arange(t).view(-1,1)  # [T,T]
        idx_h = torch.arange(h).view(1,-1) - torch.arange(h).view(-1,1)  # [H,H]
        idx_w = torch.arange(w).view(1,-1) - torch.arange(w).view(-1,1)  # [W,W]

        # 查表获取编码并拼接
        e_t = self.emb_t[idx_t + self.max_len]  # [T,T,D/3]
        e_h = self.emb_h[idx_h + self.max_len]  # [H,H,D/3]
        e_w = self.emb_w[idx_w + self.max_len]  # [W,W,D/3]

        # 使用 einops 高效组合
        e_thw = torch.einsum('tad,had,wad->thwad', e_t, e_h, e_w)
        return einops.rearrange(e_thw, 't h w a d -> a (t h w) d')

结合 FlashAttention 实现显存优化：

from flash_attn import flash_attn_qkvpacked

def block_3d_attention(q, k, v, block_size=16):
    """
    q/k/v: [B, T, H, W, C]
    分块策略：沿时间维度分块
    """
    orig_shape = q.shape
    q = einops.rearrange(q, 'b (t bt) h w c -> (b h w) bt (t c)', bt=block_size)
    k = einops.rearrange(k, 'b (t bt) h w c -> (b h w) bt (t c)', bt=block_size)
    v = einops.rearrange(v, 'b (t bt) h w c -> (b h w) bt (t c)', bt=block_size)

    out = flash_attn_qkvpacked(torch.stack([q,k,v], dim=2),
        dropout_p=0.1,
        softmax_scale=1.0
    )
    return out.reshape(orig_shape)