Adapt or Perish: 自适应稀疏Transformer在图像修复中的实践与优化

1次阅读

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

图像修复任务需要模型对局部缺失区域进行语义合理的填充，这对模型的感受野和上下文理解能力提出了极高要求。传统 Transformer 模型（如 Vision Transformer）虽然能够通过全局自注意力机制捕获长距离依赖关系，但在实际应用中面临两个致命问题：

计算复杂度爆炸 ：标准自注意力机制的计算复杂度与输入图像尺寸呈平方关系（O(n²)）。对于 512×512 的图像，单层注意力矩阵就达到 262,144×262,144，显存直接爆满
内存墙限制 ：高分辨率图像处理需要存储中间激活值，在训练时经常遇到 CUDA out of memory 错误，即使使用梯度检查点技术也难以缓解

目前图像修复领域主要有三类技术路线，各自特点如下：

CNN-based 方法 （如 PartialConv, EdgeConnect）
优势：局部感受野设计天然适合图像修复，计算效率高
劣势：难以建模远距离语义关系，修复大区域时易出现结构扭曲
传统 Transformer 方法 （如 TransFill, MAT）
优势：全局注意力实现像素级信息融合，修复质量高
劣势：计算资源消耗大，无法处理高分辨率输入
稀疏 Transformer 变体 （如 Swin Transformer, ISTN）
优势：通过窗口 / 稀疏注意力降低计算量
劣势：固定稀疏模式可能损失关键位置信息

核心思想是让模型动态决定哪些位置需要密集计算，哪些位置可以稀疏化处理。具体实现分为三步：

重要性评分 ：对查询向量 Q 和键向量 K 计算粗略相似度得分

# PyTorch 实现
def compute_importance(q, k, topk_ratio=0.3):
    scores = torch.matmul(q, k.transpose(-2,-1))  # [B,H,N,N]
    threshold = torch.topk(scores.flatten(-2), 
                          int(topk_ratio*scores.size(-1)), 
                          dim=-1).values.min()
    return scores > threshold.unsqueeze(-1)

动态掩码生成 ：根据评分保留 top- k 重要连接，其余置零
稀疏矩阵优化 ：使用 block-sparse 格式存储注意力矩阵，减少显存占用

在解码阶段引入多尺度特征精炼单元（Attentive Feature Refinement, AFR）：

层级特征融合 ：通过跨尺度跳跃连接聚合不同粒度的特征
门控注意力 ：使用 sigmoid 门控控制信息流
残差学习 ：每个精炼单元包含 shortcut 连接保持梯度流动

class AFR(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels*2, channels, 3, padding=1)
        self.attention = nn.Sequential(nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, channels//8, 1),
            nn.ReLU(),
            nn.Conv2d(channels//8, channels, 1),
            nn.Sigmoid())

    def forward(self, x, skip):
        fused = torch.cat([x, skip], dim=1)
        x = self.conv1(fused)
        attn = self.attention(x)
        return x * attn + x

在 Places2 数据集上的实验结果（RTX 3090 显卡）：