备注好友:
方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
近年来,基于Transformer的模型通过利用其内在捕捉复杂上下文特征的能力,在图像修复领域取得了显著进展。近期,Mamba模型凭借其处理长程依赖的能力和相较于Transformer的显著计算效率,在计算机视觉领域引起了广泛关注。
然而,Mamba在上下文学习能力方面目前落后于Transformer。
为了克服这两种模型的局限性,作者提出了一种名为MatIR的Mamba-Transformer混合图像修复模型。
具体来说,MatIR交叉循环Transformer层和Mamba层的块以提取特征,从而充分利用这两种架构的优势。
在Mamba模块中,作者引入了图像修复状态空间(IRsS)模块,它沿着四个扫描路径遍历,以实现长序列数据的有效处理。
在Transformer模块中,作者将基于三角窗口的局部注意力与基于通道的全局注意力相结合,以在更广泛的图像像素范围内有效地激活注意力机制。大量的实验结果和消融研究表明了该方法的有效性。
引言
图像恢复旨在从退化或损坏的输入中恢复清晰且高质量的图像。这是计算机视觉中的一个长期问题,包括了一系列子问题,如超分辨率、图像去噪和去模糊。随着现代深度学习模型,如卷积神经网络CNNs 和 Transformer 的引入,近几年来,该领域的最先进性能得到了持续提升。去噪、去模糊和超分辨率等任务需要能够准确重建图像细节同时保留结构信息的模型。传统的基于卷积的方法通常无法捕捉到涉及大或严重退化图像的任务中至关重要的长距离依赖关系。深度学习领域的最新进展,如Transformer架构,在捕捉图像中的全局依赖性方面显示出了潜力。然而,Transformer的计算成本随序列长度的平方增长,这限制了其可扩展性,尤其是在高分辨率图像恢复任务中,这些任务以二次复杂度为代价提供了全局感受野。
近期,由于其在处理长距离依赖关系和相对于Transformer在计算效率上的显著优势,Mamba在计算机视觉领域日益突出。Mamba架构是一种新的序列模型,通过引入状态空间模型(SSM)的概念,实现了对长序列数据的有效处理。状态空间模型(SSM)利用状态空间表示来达到线性计算复杂度,并且能够在不牺牲准确性的情况下高效地处理长序列。与序列长度相关的线性计算复杂度显示了其在序列建模任务中处理长距离依赖关系的高效性和有效性。然而,现有研究显示,在上下文学习(ICL)能力方面,Mamba落后于Transformer。[14, 19, 50] 从现代深度神经网络发展的最新进展中汲取灵感。在本研究中,作者的目标是弥补这些模型的不足,同时利用它们的各自优势。作者提出了一种混合Mamba-Transformer图像恢复模型,称为MatIR。这是一个新颖的混合架构,它结合了Mamba架构的优势(在处理长序列时以内存效率著称)和Transformer(在上下文学习和信息检索方面表现卓越)[14, 19, 50]。通过结合这两种方法,MatIR旨在为各种图像恢复任务提供一个强大而高效的解决方案。[33, 45]
具体而言,1)浅层特征提取阶段采用简单的卷积层提取浅层特征。随后,2)深度学习特征提取阶段使用堆叠了Mamba层的Transformer。在Transformer层中,作者运用三角窗局部注意力(TWLA)和通道全局注意力(CGA)机制,有效激活了更广泛的图像像素范围,从而提升了该模块的性能。在Mamba层中,作者运用图像恢复状态空间(IRSS)模块,通过沿四个扫描路径进行遍历,实现对不同方向和路径的长序列数据的有效处理。在提升性能和吞吐量的同时,保持可管理的内存占用。作为MatIR的三个核心组件,TWLA和CGA分别激活了局部和全局范围内,三角窗和矩形窗中的更多输入像素,实现了更高品质的图像恢复。IRSS通过四个不同方向路径(左、上、右、下)在状态空间中创建与序列长度相关的线性计算复杂度信息,展示了在序列建模任务中处理长距离依赖的效率和有效性。最后,3)高质量图像重建阶段聚合浅层和深层特征,生成高质量的输出图像。凭借局部和全局有效的感受野、高效的内存管理和计算,MatIR成为图像恢复主干网络的新选择。
总的来说,作者的主要贡献可以概括如下:
-
作者通过广泛实验将状态空间模型和Transformer注意力机制应用于图像修复领域,从而制定了MatIR,该算法在保持性能的同时大幅提升了计算效率。作者提出了图像修复状态空间(IRSS)模块,该模块沿着四个不同的路径进行扫描和遍历,以实现对长序列数据的高效处理。
-
作者提出了一种三角形窗口局部注意力(TWLA)模块和一种通道全局注意力(CGA)模块,分别用于从局部和全局范围内、三角形窗口和矩形窗口中激活更多输入像素,以实现更高质量的图像修复。
-
通过在多个基准数据集上的全面评估,作者的方法在性能上优于其他最先进的方法,为图像恢复提供了一个强大且前景广阔的 Backbone 解决方案。
2. 相关研究工作
图像恢复。图像恢复是计算机视觉领域的一个长期问题。在过去十年中,在包括图像恢复在内的多个领域,人们投入了大量努力以提升深度学习方法的性能。由SRCNN [16]开创,深度学习通过一个简单的三层卷积神经网络(CNN)被引入到图像恢复超分辨率中。自那以后,许多研究探索了各种架构改进以提升性能。VDSR [23]实现了一个更深层的网络,而DRCN [24]提出了递归结构。EDSR [31]和RDN [69]开发了新的残差模块,以进一步改善CNN在超分辨率中的能力。然而,尽管CNN取得了成功,但其感受野本质上有限,这使得捕捉长距离依赖关系变得困难。
近年来,视觉Transformer(ViT)[17]及其变体[11, 34, 52]将自注意力机制引入图像处理领域,使模型能够学习全局关系。基于此,IPT[6]成功尝试利用基于Transformer的网络进行各种图像恢复任务。此后,开发了多种技术来提升图像恢复Transformer的性能。这些技术包括SwinIR[30]和CAT[9]实现的移位窗口自注意力、ELAN[66]的分组多尺度自注意力机制、ART[59]和OmniSR[51]的稀疏注意力,以及GRL[29]的 Anchor 定自注意力机制、多注意力机制DART和DISR[54, 55],它们都旨在扩大感受野范围,以获得更好的效果。然而,这以二次计算复杂度为代价提供了全局感受野。自注意力在序列长度方面的二次计算复杂性问题构成了挑战,尤其是在处理高分辨率图像时。
状态空间模型。近年来,由于能够处理长距离依赖关系以及相对于Transformer的高计算效率,Mamba在计算机视觉领域日益突出。Mamba架构是一种新的序列模型,通过引入状态空间模型(SSM)的概念[14, 19, 50],实现了对长序列数据的有效处理。状态空间模型(SSM)利用状态空间表示,实现线性计算复杂度,并能高效处理长序列数据而不影响准确性。与序列长度相关的线性计算复杂度显示出在序列建模任务中处理长距离依赖关系的高效性和有效性。然而,现有研究表明,Mamba在上下文学习(ICL)能力上落后于Transformer[50]。在高效计算和全局建模之间的权衡困境尚未得到根本解决[20, 33, 45]。考虑到两种当前最先进模型的局限性,作者探讨了混合Mamba-Transformer方法在图像修复中的潜力。在本论文中,基于Mamba和Transformer模型的有效性,作者提出了三个核心组件:图像修复状态空间(IRSS)模块,该模块沿着四种不同的路径扫描和遍历,以实现长序列数据的有效处理。三角窗口局部注意力(TWLA)块和通道全局注意力(CGA)块,每个块都通过注意力机制激活来自局部和全局、三角形和矩形窗口范围内更多的输入像素,以实现更高品质的图像修复。
3. 研究方法
3.1. 前言
结构化状态空间序列模型(S4)是近期出现的一类用于深度学习的序列模型,它们与循环神经网络(RNN)、卷积神经网络(CNN)和经典状态空间模型有广泛的关联。这些模型灵感来源于一个特定的连续系统,该系统通过一个隐含的状态
将一维函数或序列
映射到
。形式上,该系统可以表述为一个线性常微分方程(ODE)[14, 19, 50]:
在状态大小为
的情况下,
,
,
,
。之后,通常采用离散化过程将方程(1)集成到实际的深度学习算法中。具体来说,用
表示时间尺度参数,将连续参数
、
转换为离散参数
、
。常用的离散化方法是零阶保持法(ZOH),其定义如下:
在离散化之后,步长为
的方程(1)的离散化形式可以重写为以下循环神经网络(RNN)形式:
此外,式(3)也可以通过数学等效变换转化为以下卷积神经网络(CNN)形式:
在此,
代表输入序列的长度,
表示卷积操作,而
是一个结构化的卷积核。
通常,该模型在高效并行化训练(在整个输入序列提前可见的情况下)时使用卷积模式(式(4)),并在高效的自回归推理(输入逐个时间步可见)时切换到循环模式(式(3))。从方程中作者可以看到一个重要的性质,即模型的动态特性随时间保持恒定。这一特性被称为线性时不变性。从这个属性中,作者可以看到其模型动态与Transformer的优势,Transformer的计算成本和序列长度呈二次增长。状态空间模型(SSM)利用状态空间表示来实现线性计算复杂度,其计算复杂度与序列长度线性相关,可以高效处理长序列而不影响准确性。
3.2 全局注意力模块(Channel Global AttentionBlock)
近年来,基于Transformer的图像恢复网络在诸如SwinIR、GRL、DART和DISR(参考文献[28, 30, 54, 55])等应用中采用了最新的块设计,证明了基于通道的全局注意力(CGA)在图像恢复领域表现出色。这表明,通过定制新的块结构,将CGA应用于MatIR基于的恢复网络具有巨大的潜力。因此,作者输入数据,令输入特征图
,其中
代表通道数,
代表特征图的空间维度。沿通道维度展平输入
以获得
,其中
。对空间维度进行全局池化以聚合空间信息,为每个通道获取全局表示:
,其中
代表每个通道的全局描述。注意力机制将计算资源导向输入信号中信息最丰富的部分,从而提高模型效率。Transformer的注意力机制基于 Query (
)、键(
)和值(
)。
在通道注意力机制中,作者关注通道维度,并定义如下: Query :
键:
值:
其中,
是可学习的权重矩阵。具体来说,基于通道的全局注意力机制在特征维度上操作,使作者能够理解模型在做出特定决策时依赖哪些特征。通过分析跨通道的注意力权重,作者可以深入了解模型的推理过程,这有助于识别偏见或改进的领域。因此,作者计算通道注意力矩阵。
其中,
表示通道之间的注意力权重矩阵,属于
。输出结果是通过加权通道表示得到的:
。将通道注意力结果重新应用于原始输入特征
,通常通过点积或按通道加权完成: