3D局部窗口自注意力机制：原理剖析与高效实现指南

1次阅读

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

在 3D 数据处理（如视频分析或医学影像）中，传统的全局自注意力机制虽然建模能力强，但计算复杂度高达 $O(N^4)$。例如处理 128×128×128 的体数据时，注意力矩阵需要 $128^6=2^{42}$ 次计算——这直接导致显存爆炸和训练困难。

滑动窗口：固定大小的立方体邻域计算注意力，计算复杂度降至 $O(k^3N)$（k 为窗口边长）
轴向注意力：分别在 XYZ 三个轴向做 1D 注意力，复杂度 $O(3kN)$ 但会损失空间关联性
交叉窗口通信 ：通过移位窗口(Shifted Window) 或全局 token 引入跨窗口交互

def create_3d_windows(x, window_size=(8,8,8)):
    """
    输入: x shape [B, C, D, H, W]
    输出: windows shape [num_windows*B, window_size[0]*window_size[1]*window_size[2], C]
    """
    B, C, D, H, W = x.shape
    x = x.view(B, C, 
               D // window_size[0], window_size[0],
               H // window_size[1], window_size[1],
               W // window_size[2], window_size[2])
    windows = x.permute(0,2,4,6,3,5,7,1).contiguous()
    return windows.view(-1, window_size[0]*window_size[1]*window_size[2], C)

对于窗口内坐标为 $(i,j,k)$ 和 $(m,n,p)$ 的两个位置，其相对位置编码为：

$$\text{PE}_{rel} = \text{Linear}(\text{Concat}(\text{PE}_x(i-m), \text{PE}_y(j-n), \text{PE}_z(k-p)))$$

其中 $\text{PE}_*$ 为可学习的 1D 位置编码。

方法	复杂度	128^3 输入显存占用
全局注意力	$O(N^4)$	256GB(爆显存)
窗口注意力(k=8)	$O(512N)$	6.4GB

# 使用 DDP 时避免不必要的 all-gather
with torch.no_grad():
    local_rank = torch.distributed.get_rank()
    # 各 GPU 处理不同窗口块
    windows = windows.chunk(torch.distributed.get_world_size(), 
        dim=0
    )[local_rank]