超分之SPIN - 技术栈

Lightweight image super-resolution with superpixel token interaction[C]

利用超像素token交互实现轻量级图像超分辨率

Zhang A, Ren W, Liu Y, et al.

Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 12728-12737.

文章目录

- - 摘要
  - [1. 引言](#1. 引言)
  - [2. 相关工作](#2. 相关工作)
  - - [2.1 D-CNN SR](#2.1 D-CNN SR)
    - [2.2 Lightweight SR](#2.2 Lightweight SR)
    - [2.3 Pixel Clustering (像素聚类)](#2.3 Pixel Clustering (像素聚类))
  - [3. 方法论](#3. 方法论)
  - - 超像素聚合模块（SPA）
    - [3.2 超像素交叉注意力模块（SPCA）](#3.2 超像素交叉注意力模块（SPCA）)
    - [3.3 超像素内注意力模块（ISPA）](#3.3 超像素内注意力模块（ISPA）)
  - [4. 实验](#4. 实验)
  - - [4.1 数据集](#4.1 数据集)
    - [4.2 实验细节](#4.2 实验细节)
    - [4.3 与轻量化模型比较](#4.3 与轻量化模型比较)
  - [5. 消融实验](#5. 消融实验)
  - [6. 总结](#6. 总结)

摘要

首先，指出了 Transformer-SISR实现了SOTA。但是应用于整个图像时，自注意力机制的计算成本很高。
- 目前的做法：将输入的LR图像分成小块，然后分别处理，然后融合以生成HR图像。
- 缺点：将LR图像规则块划分，过于粗糙且缺乏可解释性，导致注意操作期间出现伪影和非相似结构干扰。
然后，提出了本文的模型网络---SPIN(super token interaction network)
- SPIN原理：使用超像素将局部相似像素聚类以形成可解释的局部区域，并利用超像素内注意力来实现局部信息交互。
- SPIN的可解释性：只有相似的区域相互补充，而不相似的区域被排除在外。
最后，通过大量的实验，说明了SPIN在准确性和轻量级方面实现了SOTA。

1. 引言

首先介绍了SISR：
- 目的：从给定的 LR 图像生成高分辨率 (HR) 图像
- 应用：医学成像、监控和数码摄影。
介绍CNN-SR：
- 背景：Dong的SRCNN是开创性工作，之后CNN-SR使用更深、更复杂的架构来实现更好的性能。
- 缺点：更深、更复杂的网络架构，会增加计算资源和提高成本，从而可能会限制其应用场景。
介绍了Attention-SR：
- 特点：attention允许网络选择性地关注输入的相关区域，从而提高SR输出的质量，如SwinIR和ESRT。
- 缺点1：Attention具有较高的计算复杂度和内存消耗，这需要将大图像分成小块进行单独处理，这样会导致连续结构的分裂，从而阻碍在其他区域使用类似信息来增强图像细节。
- 缺点2：每个补丁内应用的局部注意机制在计算中涉及不相关的区域，从而导致不良的推断。
基于上述问题，提出本文的SPIN：
1. 首先，对输入图像的像素进行基于CNN的浅层特征提取，并进行局部聚类将相邻像素分组为超像素。
2. 然后，通过基于相似性的超像素聚类获得局部区域，并分别对它们进行局部特征提取。
  - 特点：以前使用固定形状块划分的方法，仅用于提高并行计算效率；SPIN划分策略更具解释性，可以更灵活、自适应地划分输入图像，并防止连续结构分裂。
3. 然后，引入超像素交叉注意模块（SPCA），通过超像素的代理实现长距离信息交互。
4. 此外，本文设计了一种应用于超像素像素的超像素内注意（ISPA）机制，将原有的注意力操作仅在常规图像区域扩展。
  - 特点：ISPA机制确保了局部注意力机制的信息交互发生在相似的区域，消除了干扰和无关计算。
- SPCA与ISPA相互交织，并在局部和全局特征提取中相互配合。
总结：
1. 提出了一种新颖的超分辨率模型，将超像素聚类与 Transformer 结构相结合，从而形成一个更易于解释的框架。
2. 提出了在超像素内和超像素之间运行的超像素内注意 (ISPA) 和超像素交叉注意 (SPCA) 模块，能够在不规则区域进行计算，同时保持捕获长距离依赖关系的能力。
3. 实验表明，与最先进的轻量级 SR 方法相比，SPIN实现了更好的 SR 重建性能。

2. 相关工作

2.1 D-CNN SR

介绍了主流的CNN-SR模型：
- SRCNN（开山之作：先将LR图像进行上采样到目标尺寸，然后使用三层CNN实现SR技术）
- VDSR（使用很深的网络结构和残差连接）
以及Attention-SR模型：
- RCAN（使用通道注意力机制）
- CSAR（使用通道注意力和空间注意力相结合）
- SwinIR（将Swin Transformer运用到SR任务， 参数量很大）

2.2 Lightweight SR

经典模型：
- FSRCNN（首次直接使用LR图像输入，最后使用反卷层上采样）
- ESPCN（使用亚像素卷积实现上采样重建）
- CARN（使用组卷积核级联机制来提高效率，但损害了其性能）
- IMDN（使用三步蒸馏来提取特征，并使用切片操作来划分提取的特征，但灵活性不足）
- BSRN（设计了深度可分离卷积来降低模型复杂度，并利用注意力机制来提高 SR 重建性能）
- ELAN（采用移位卷积和组自注意力）
缺点： Lightweight SR轻量且高效，但是SR重建的质量不高。

2.3 Pixel Clustering (像素聚类)

特点：像素聚类的一种常见方法是使用 CNN 生成将相似像素分组在一起的像素级嵌入。
典型方法：
- Affinity Derivation and Graph Merge for Instance Segmentation：该网络学习逐像素的亲和力来聚类像素
- Learning a convolutional neural network for non-uniform motion blur removal：该网络学习像素级表示来聚类图像块
- Weakly supervised object localization with pro-gressive domain adaptation.：使用 CNN 特征和稀疏标记方案将像素聚类到对象区域。
特点：像素聚类的方法充分利用了 CNN 和聚类算法的强大功能，使图像处理任务中的像素聚类更加准确和高效。
图卷积网络GCN（Graph Convolutional Network）：
- GCN 能够通过构建图像的图形表示来建模图像中像素之间的依赖关系，其中每个像素都是一个节点，而边缘表示像素之间的关系。
论文：
- Semi-supervised hyperspectral image classification with graph clustering convolutional networks：基于 GCN 的高光谱图像分类框架，该框架使用两种聚类策略来利用多跳相关性
第一个聚类策略根据相似像素的光谱相似性对它们进行分组，
第二个聚类策略根据像素的空间相邻性对它们进行分组。
聚类方法在SR任务的不足：尽管像素聚类在各种图像处理任务中都表现出良好的效果，但它尚未在超分辨率应用中得到有效应用。

3. 方法论

SPIN模型架构：
编码器模块
- 3×3的Conv卷积层，用于将LR图像嵌入到高维特征空间。
- x e m b = f e n c o d e r ( I L R ) ( 1 ) x_{emb} = f_{encoder}(I_{LR})\ \ \ \ \ \ \ \ (1) xemb=fencoder(ILR) (1)
SPI模块（Super-Pixel Interaction）：
- 堆叠k个SPI块，用于提取包含输入图像的丰富的低级和高级信息的深层特征。
- SPI由四部分组成：SPA（超像素聚合）、SPCA（超像素交叉注意力）、ISPA（超像素内注意力）、局部注意力。
- s i = f S P A ( x i − 1 ) , s_i = f_{SPA}(x_{i - 1}), si=fSPA(xi−1),
- x i = x i − 1 + f l o c a l ( f I S P A ( f S P C A ( x i − 1 ， s i ) ) ) x_i = x_{i -1} + f_{local}(f_{ISPA}(f_{SPCA}(x_{i-1， s_i}))) xi=xi−1+flocal(fISPA(fSPCA(xi−1，si)))
- s i s_i si：第i个SPI块中超像素的特征。
- f S P C A ( ⋅ ) f_{SPCA}(·) fSPCA(⋅)：用于捕获长距离像素之间的依赖关系和交互。
- f I S P A ( ⋅ ) f_{ISPA}(·) fISPA(⋅)：用于捕获每个超像素内像素的依赖关系和交互。
- f l o c a l ( ⋅ ) f_{local}(·) flocal(⋅)：用于增强局部区域内像素之间的交互。
- 残差连接：用于简化整个训练过程。
全局残差模块
- 由3×3的Conv卷积层 + Pixel Shuffle +3×3的Conv卷积层构成，
- 来获取全局的残差信息。
上采样模块
- 全局的残差信息被添加到 I L R I_{LR} ILR的上采样图像中，以解析高分辨率图像 I S R I_{SR} ISR。

超像素聚合模块（SPA）

SPA功能：超像素划分可以在感知上将相似的像素分组在一起，从而可以描绘出更精确的边界，降低产生模糊和不准确边界的风险。
SSN中的soft K-means-based 超像素算法：
- 给定视觉标记 x ∈ R N × C x\in R^{N×C} x∈RN×C（N = H×W，表示视觉标记的数量），每个标记 x ( i ) ∈ R C x(i)\in R^C x(i)∈RC属于M个超像素 s ∈ R M × C s\in R^{M×C} s∈RM×C。因此需要计算视觉标记与超像素标记之间的联系。
超像素聚合过程是一个类似于期望最大化的过程，总共包含 T 次迭代。
- 首先，通过规则网格中的标记进行平均采样，来初始化超像素标记 S 0 S^0 S0，成为Patchify。假设网格大小为 H s × W s H_s×W_s Hs×Ws ，则超像素标记的数量为 M = H H s × W W s M = \frac {H}{H_s} × \frac {W}{W_s} M=HsH×WsW 。对于第 t 次迭代，关联图如下：
- A t ( i j ) = e − ∣ ∣ x ( i ) − s t − 1 ( j ) ∣ ∣ 2 2 ( 3 ) A^t(ij) = e^{-||x(i) - s^{t-1}(j)||^2_2}\ \ \ \ \ \ (3) At(ij)=e−∣∣x(i)−st−1(j)∣∣22 (3)
- A t ( i j ) ∈ R N × M A^t(ij)\in R^{N × M} At(ij)∈RN×M：表示第i行第j列的关联图值。
- 超像素聚合仅计算从每个标记到周围超像素的关联图，这保证了超像素的局部性，从而使其在计算和内存方面也非常高效
之后得到超像素 S t S^t St 作为视觉标记的加权和，公式为：
s t ( j ) = 1 Z t ( j ) ∑ i A t ( i j ) x ( i ) ( 4 ) s^t(j) = \frac{1}{Z^t(j)}\sum_iA^t(ij)x(i) \ \ \ \ \ \ \ \ (4) st(j)=Zt(j)1i∑At(ij)x(i) (4)
- Z t ( j ) = ∑ i A t ( i j ) Z^t(j) = \sum_iA^t(ij) Zt(j)=∑iAt(ij)：沿列的归一化项。
经过T次迭代，能得到最终的关联图 A T A_T AT。

3.2 超像素交叉注意力模块（SPCA）

目的：SPA仅能捕获局部中像素的局部性和互连性，因此可能缺乏捕获超分辨率的远程依赖关系的能力。
方法：使用自注意力范式通过超像素的替代来增强远程通信。

自注意力机制具体实现：

给定超像素特征s ∈ R M × C \in R^{M×C} ∈RM×C，其中M表示超像素的数量；展平的像素特征 x ∈ R H W × C x \in R^{HW×C} x∈RHW×C.
Q s = s W q s , K x = x K k , V x = x V k ( 5 ) Q^s = sW^s_q,\ \ K^x = xK^k, \ \ V^x = xV^k\ \ \ \ \ \ \ \ (5) Qs=sWqs, Kx=xKk, Vx=xVk (5)
其中：
- query： Q s ∈ R M × D Q^s \in R^{M×D} Qs∈RM×D，key： K s ∈ R M × D K^s \in R^{M×D} Ks∈RM×D、value： Q s ∈ R M × D Q^s \in R^{M×D} Qs∈RM×D
- 权重矩阵W： W q s ∈ R C × D W^s_q \in R^{C×D} Wqs∈RC×D， W k s ∈ R C × D W^s_k \in R^{C×D} Wks∈RC×D， W v s ∈ R C × C W^s_v \in R^{C×C} Wvs∈RC×C
自注意力机制计算：
S u = s o f t m a x ( Q s ( K x ) T / D ) V x ( 6 ) S_u = softmax(Q^s(K^x)^T/\sqrt{D})V^x \ \ \ \ \ \ \ \ \ \ \ \ \ (6) Su=softmax(Qs(Kx)T/D )Vx (6)
其中：
- D \sqrt{D} D ：避免梯度消失的缩放因子。
交叉注意力机制具体实现：
- 上述的注意力机制，将像素传播到超像素 ，因此还需要进行交叉注意力计算，将将更新后的超像素特征映射回像素级 ，以实现像素之间的信息传播，具体实现如下：
  - query：利用另一个权重矩阵 W q x W^x_q Wqx从像素特征中获取query（查询） Q x Q^x Qx，
  - key：使用超像素特征做为key(值) Q s Q^s Qs，
  - value：将更新后的超像素特征作为value（值） S u S_u Su，
  - 将 Q x Q^x Qx、 Q s Q^s Qs、 S u S_u Su进行交叉注意力计算，将更新后的超像素特征映射回像素级。
Reshape层具体实现：
1. 采用了包含LN（层归一化）的FFN（逐项前馈网络），
2. 利用特征门控来调节输入特征和通道注意力，来提取全局信息。
3. 使用两个全连接层和GELU激活函数。

3.3 超像素内注意力模块（ISPA）

目的：利用同一超像素内相似像素的互补性，来提高超分辨率图像质量。
缺点：不同的超像素可能包含不同数量的像素（超像素的大小不一），这使得并行处理变得困难，也会导致意外的内存消耗。
具体做法（使用关联图并选择与每个超像素最相似的前N个像素）：
1. 假设一个超像素的关联像素为 f = { x ( i ) } N ∈ R N × C f=\{x(i)\}_N \in R^{N×C} f={x(i)}N∈RN×C，（N：选定像素的数量）。
2. 使用公式5和6的注意力机制来进行超像素内注意力计算，其中：
  - W q f 、 W k f 、 W v f W^f_q、W^f_k、W^f_v Wqf、Wkf、Wvf：query、key、value的权重矩阵。
3. Scatter：利用在前 N 个选择过程中生成的索引将细化的像素特征分散回图像中的各自位置。（可能会导致一些"被忽略"的像素，即这些像素不包含在任何超像素中。）
4. Feature Intergate：对于那些"被忽略"的像素，我们利用value投影 W v f W^f_v Wvf对其进行投影，以获得更新的特征，然后将其与通过超像素内交互更新的像素进行集成。
5. Reshape：在ISPA模块后采用相同的FNN。

4. 实验

4.1 数据集

训练集与验证集：
- DIV2K（900张HR图像）
  - 前800张用于训练
  - 后100张用于验证
测试集：
- Set5、Set14、BSDS100、Urban100 和 Manga109。、

4.2 实验细节

初始学习率： 5 e − 4 5e-4 5e−4
训练轮次：1000 epoch
优化器：Adam（ β 1 = 0.9 , β 2 = 0.999 \beta_1=0.9, \beta_2=0.999 β1=0.9,β2=0.999）
数据增强：随机旋转（90°、180°、270°）、水平翻转。
评价指标：PSNR、SSIM（将RGB图像转换为YCbCr格式后，测量Y通道上的指标）

4.3 与轻量化模型比较

定量比较（Quantitative comparison）：
- 不仅指标都实现了最佳或第二，参数量也比现有的基于Transformer的方法要少。
- 原因：上面的模型总是将图像分成规则的块，这可能会破坏输入图像中的对象、边界等 。SPIN利用超像素来实现 Transformer 的可解释和连续区域划分 。
定性比较（Quantitative comparison）：
- 优势：只要 LR 图像中存在相应的非局部信息，所提出的 SPIN 就可以有效恢复严重受损的纹理。相反，缺乏非局部注意的深度 SISR 模型无法准确重建受损纹理。

5. 消融实验

各个模块：证明了ISPA 模块和 SPCA 模块在SPIN中发挥着重要作用
SPI 块的个数：随着块数的增加，SPIN网络性能也得到了改善。然而，当块数大于 8 时，性能开始下降。
- 原因：可能是过度参数化的网络对训练数据进行了过度拟合，导致在其他基准上泛化能力较差。

6. 总结

本文提出了一种超级标记交互网络 (Super Token Interaction Network，SPIN) 的新方法，该方法利用超像素将局部相似像素分组为可解释的局部区域。
SPIN采用超像素内注意力模块 （intra-superpixel attention，ISPA）来促进不规则局部超像素区域内的局部信息交互，超像素交叉注意力模块（superpixel cross-attention，SPCA）通过超像素的替代来促进远程信息交互。
大量实验表明，SPIN在准确性和轻量级 方面实现了SOTA，并且SPIN为处理具有可解释区域划分的整个图像的挑战提供了一个有希望的解决方案。。