BiFormer：基于双级路由注意力的视觉Transformer优化实践

1次阅读

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

视觉 Transformer（ViT）在计算机视觉任务中表现出色，但在处理高分辨率图像时面临显著的计算瓶颈。传统 ViT 采用全局注意力机制，其计算复杂度与输入图像的分辨率呈平方关系（O(n²)）。例如，对于一个 224×224 的图像，注意力矩阵的大小将达到 50176×50176，这对计算资源和内存提出了极高要求。

内存占用高：全局注意力需要存储巨大的中间矩阵，极易导致 GPU 内存溢出
计算效率低：大量计算资源被浪费在不重要的背景区域上
长距离依赖弱：简单下采样会损失细粒度特征，影响小目标检测性能

对比当前主流的高效 ViT 变体，BiFormer 在计算效率和模型性能间取得了更好平衡：

Swin Transformer
优点：通过局部窗口和移位窗口降低计算量
限制：固定窗口大小无法自适应内容
计算复杂度：O(4hwC² + 2M²hwC)
PVT 系列
优点：金字塔结构保留多尺度特征
限制：空间缩减导致信息损失
计算复杂度：O(hwC² + (hw)²/s²)
BiFormer 创新点
动态路由：根据内容重要性分配计算资源
双粒度注意力：粗粒度筛选 + 细粒度计算
计算复杂度：O(hwC² + hwK + K²C)（K 为路由区域数）

通过轻量级决策网络生成区域重要性分数，TopK 筛选关键区域：

class RegionRouter(nn.Module):
    def __init__(self, dim, num_regions=64):
        super().__init__()
        self.num_regions = num_regions
        self.scorer = nn.Sequential(nn.Conv2d(dim, dim//4, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(dim//4, 1, 1)
        )

    def forward(self, x):
        # x: [B,C,H,W]
        scores = self.scorer(x).flatten(2)  # [B,1,HW]
        _, indices = scores.topk(self.num_regions, dim=-1)  # [B,1,K]
        return indices

在选定区域内计算精确注意力，保留局部细节：

class PixelAttention(nn.Module):
    def __init__(self, dim, head_dim=32):
        super().__init__()
        self.head_dim = head_dim
        self.scale = head_dim ** -0.5

    def forward(self, q, k, v, region_mask):
        # q/k/v: [B,H,N,C]
        # region_mask: [B,N]
        attn = (q @ k.transpose(-2,-1)) * self.scale
        attn = attn.masked_fill(~region_mask.unsqueeze(1), float('-inf'))
        attn = attn.softmax(dim=-1)
        return attn @ v

class BiFormerBlock(nn.Module):
    def __init__(self, dim, num_heads=8, num_regions=64):
        super().__init__()
        self.region_router = RegionRouter(dim, num_regions)
        self.pixel_attention = PixelAttention(dim//num_heads)

        # 初始化代码省略...

    def forward(self, x):
        B, C, H, W = x.shape
        regions = self.region_router(x)  # 获取重要区域

        # 区域特征提取
        patch_emb = x.flatten(2).transpose(1,2)  # [B,HW,C]
        region_feats = batched_index_select(patch_emb, regions)  # [B,K,C]

        # 双级注意力计算
        qkv = self.qkv_proj(region_feats)
        # 注意力计算过程省略...

        return updated_feats.reshape(B, C, H, W)

在 ImageNet-1K 上的对比实验结果：

模型	FLOPs(G)	内存(MB)	Top-1 Acc.
ViT-Base	17.6	1203	79.9%
Swin-Tiny	4.5	357	81.2%
PVT-Medium	6.7	489	81.9%
BiFormer-S	3.8	302	82.4%

关键优势：
– 相比 ViT 降低 78% 计算量
– 内存占用仅为 Swin 的 85%
– 准确率提升 1 - 2 个百分点

区域数量：一般设置为输入 token 数的 5 -10%
温度系数：softmax 温度影响路由锐度，建议 0.1-1.0
蒸馏策略：用教师模型指导路由决策

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

注意：
– 路由得分计算需保持 FP32
– 使用 torch.cuda.amp.GradScaler 防梯度下溢