图像超分辨率Transformer实战：如何通过activating more pixels提升细节重建能力

1次阅读

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

在图像超分辨率任务中，基于 Transformer 的模型（如 SwinIR）虽然能够有效建模长距离依赖关系，但在重建高频细节（如边缘、纹理）时仍存在明显的模糊现象。根本原因在于标准自注意力机制（self-attention）中，每个查询向量（query vector）仅与有限的键值对（key-value pairs）交互，导致大量像素信息未被充分激活（activated）。这种现象在图像平坦区域影响较小，但在需要精细重建的区域会损失高频细节。

传统 SwinIR 使用固定窗口划分（fixed window partitioning），导致跨窗口的像素交互受限。改进方案采用：

重叠位置编码（overlapping positional encoding）：在窗口边缘引入 25% 的重叠区域，增强边界像素的交互能力
可学习尺度因子：对位置编码施加可训练的缩放系数 $\alpha$，公式如下：
$$PE_{new} = \alpha \cdot PE_{original} + \beta$$

标准注意力矩阵 $A$ 的计算为：
$$A = softmax(\frac{QK^T}{\sqrt{d_k}})$$

改进方案通过两步实现动态稀疏化：

重要性采样：根据特征图梯度选择 top- k 重要像素
局部增强：对选中像素应用 3 ×3 深度卷积（depth-wise conv）增强局部关联

class DynamicAttention(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.scale = (dim // num_heads) ** -0.5
        self.to_qkv = nn.Linear(dim, dim*3)
        # 内存优化：使用梯度检查点
        self.local_conv = checkpoint_sequential(
            nn.Sequential(nn.Conv2d(dim, dim, 3, padding=1, groups=dim),
                nn.GELU()), chunks=2
        )

    def forward(self, x):
        B, C, H, W = x.shape
        qkv = self.to_qkv(x.flatten(2)).chunk(3, dim=1)
        q, k, v = map(lambda t: t.view(B, -1, H*W).transpose(1, 2), qkv)

        # 动态稀疏注意力
        attn = (q @ k.transpose(-2, -1)) * self.scale
        mask = torch.topk(attn, k=int(H*W*0.3), dim=-1)[1]  # 激活 30% 像素
        sparse_attn = attn.scatter(-1, mask, -float('inf'))
        sparse_attn = sparse_attn.softmax(dim=-1)

        # 局部特征增强
        local_feat = self.local_conv(x)
        out = (sparse_attn @ v) + local_feat.flatten(2).transpose(1,2)
        return out.view(B, C, H, W)

在 DIV2K 验证集上的测试结果：