点击下方卡片，关注
「AI视界引擎」
公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
基于深度学习的单图像超分辨率（SISR）技术专注于将低分辨率（LR）图像提升为高分辨率（HR）图像。尽管取得了显著的进展，但在计算复杂度和质量方面仍存在挑战，尤其是在遥感图像处理中。
为了解决这些问题，作者提出了作者的全尺度RWKV超分辨率（OmniRWKVSR）模型，该模型提出了一种新颖的方法，将接受权重关键值（RWKV）架构与特征提取技术，如视觉RWKV空间混合（VRSM）和视觉RWKV通道混合（VRCM）相结合，旨在克服现有方法的局限性，实现卓越的超分辨率性能。该研究已证明能够提供高质量图像重建的有效解决方案。
在4倍超分辨率任务中，与MambaIR模型相比，作者实现了平均0.26%的峰值信噪比（PSNR）和0.16%的结构相似性（SSIM）的提升。
单图像超分辨率（SISR）技术在图像处理领域占据重要地位，旨在将低分辨率（LR）图像重建为高分辨率（HR）图像，以满足对图像清晰度的日益增长的需求。SISR在医学成像[1]、遥感[2]和视频监控[3]等多个领域具有广泛的应用潜力。高分辨率图像不仅提供了更详细的视觉信息，而且显著提高了决策支持的准确性，如在医学诊断和环境监测应用中[4]、[5]。
近年来，基于深度学习的超分辨率图像重建（SISR）方法在利用大量低分辨率/高分辨率（LR/HR）图像对进行模型训练后，在重建低分辨率图像方面取得了显著进展。卷积神经网络（CNNs）通常用于特征提取，通过多次非线性变换实现详细的图像恢复[6]，[7]。然而，将SISR应用于遥感图像仍存在挑战，如局部特征、复杂结构、尺度差异和噪声[8]。为此，已提出了一些高级模型。SwinIR模型利用层次化特征表示和基于移位窗口的自注意力机制来捕获长距离依赖关系[9]。SwinFIR通过整合快速傅里叶变换卷积（FFC）组件，进一步提高了全局信息效率[10]。MambaIR模型采用Mamba模型进行图像恢复，结合卷积和通道注意力机制以改善表示[11]，[12]。混合注意力 Transformer （HAT）通过选择性应用注意力优化资源分配并增强特征提取[13]。RWKV模型[14]，特别是其第六代（RWKV v6），通过整合循环神经网络（RNNs）和Transformer的优势，同时避免具有线性复杂度的全注意力机制，提高了训练稳定性和收敛速度[15]。
本研究提出了一种新颖的方法，该方法将RWKV架构与高级特征提取技术相结合，并通过消融实验进行验证，旨在克服现有基于注意力方法（如二次复杂度）的局限性。本研究的贡献主要包括以下几点：
- FFN ：ChannelMix在增强通道级特征混合和提升信息流方面优于传统的MLP（多层感知器），因此在前特征提取和图像重建质量上超越了MLP。
- 全四象限转换：为了捕捉长距离依赖关系，作者引入了一种补充机制：全四象限转换。该机制使模型能够有效地捕捉多尺度特征和空间转换。
- 加权键值(WKV)二维扫描：传统的线性模型只能捕捉一维信息，这也被称为“扫描”。尽管Spate-Space模型实现了二维扫描，但其质量和性能仍有待提升。为了解决这一问题，本文提出了加权键值(WKV)机制的二维扫描版本。
- 作者的模型在训练速度上实现了最先进的性能：与MambaIR相比，作者的模型在Set14、BSD100等流行数据集上达到了最高分数，同时训练时间减少了约15%。
近年来，注意力机制通过使模型能够有效地捕捉长距离依赖关系并专注于关键图像区域，显著提升了超分辨率图像重建（SISR）的性能。值得注意的是，基于Swin Transformer的模型SwinIR[9]在多种图像恢复任务中表现出色，包括超分辨率、去噪和JPEG伪影减少[16]，[17]。
尽管注意力机制在超分辨率图像重建（SISR）[9]、[18]中带来了显著的进步，但一个显著的局限性依然存在：自注意力操作的计算复杂度与输入大小呈二次方增长。这种二次方复杂度产生的原因是，注意力机制会计算输入序列中所有元素之间的成对交互，导致巨大的计算和内存需求，尤其是在高分辨率（HR）图像中。
图1. OmniRWKVSR的架构。
为了解决基于注意力的超分辨率图像重建（SISR）的计算难题，一些模型[9]、[10]集成了创新机制以提高效率。MambaIR模型[11]引入了一种选择性状态空间二维（SS2D）机制[19]、[20]，该机制采用具有选择性扫描策略的统计状态模型（SSM）来捕捉长距离依赖关系，同时保持与输入大小成线性关系的计算复杂度。这种设计有效地减轻了传统二次注意力机制的计算负担，使MambaIR在HR图像修复任务中具有可扩展性和高效性。SS2D机制允许MambaIR建模复杂的图像细节，而无需承担自注意力方法通常带来的巨大成本，从而在大型图像处理中平衡性能与效率。
尽管MambaIR[11]和SS2D[12]机制引入了进展，但在长时间训练、不稳定的性能指标、快速收敛和次优的特征提取能力方面仍存在挑战。RWKV v6[14]通过整合RNNS和Transformer的优势并消除对注意力机制的依赖，为这些问题提供了解决方案。RWKV v6的一个关键改进是其提高了训练稳定性和收敛速度。通过采用谨慎的初始化策略和优化的训练协议，RWKV v6减轻了与训练不稳定和过早收敛相关的问题，使得在各项任务中表现出更稳健和可靠的性能。此外，RWKV关于输入序列长度的线性计算复杂度解决了传统注意力机制中二次复杂度相关的低效问题。因此，在本文中，作者提出了这种创新方法，该方法结合了RWKV架构的优势和高级特征提取技术，以增强SIsR任务。
本节首先详细阐述了作者所提出的全场景RWKVSR模型的整体流程，然后解释了网络的基本模块：视觉RWKV残差组（VRG）。
图1展示了Omni-Scale RWKV模型的总体结构。作者用
表示低分辨率输入图像，用
表示高分辨率目标图像。最初，作者使用一个
的卷积神经网络层，作为浅层特征提取器（提取），从低分辨率输入图像中提取低级特征。
其次，作者使用多个虚拟现实生成器（VRGs）从
中提取高级特征，记为
。
在本文中，
表示一个VRG（虚拟现实生成器），
是浅层特征提取器的输出，而
则是前一个VRG的输出。
第三，作者采用一个
卷积神经网络层以及一次上采样操作（重建）来将
上采样到高分辨率，记为
。
在本文中，
表示一个
的卷积神经网络层，
是第
个视觉回归生成器（VRG）的输出，而上采样则表示一个步骤式的上采样操作。
最后，作者将均值和方差重新加回到
中，得到高质量的重构图像，记作
。
在元素级乘法运算符
下，
和
分别表示 LR 输入图像的均值和方差。
VRG是一种包含残差块集合的结构，其中VRSM和VRCM是其核心组件。VRSM促进了不同分辨率间特征图的融合，而VRCM将这些特征限制在特定通道中。VRG由多个视觉RWKV残差块（VRBs）组成，VRB的整体结构如图2所示。
VRG从上一个VRG的输出或浅层特征提取器接收输入，其输出被导向后续VRG或高质量重建过程。进入gVRSM的特征最初经过Omni-Quad Shift操作进行位移，生成一组包含邻近像素特征的位移特征图。然后，这些图通过空间混合操作中的通道线性注意力机制进行整合，该机制选择重要通道并执行2D扫描操作以重新组装特征图。结果通过残差连接反馈到输入特征图。随后，输出被导向VRCM，该模块将特征图限制在特定通道内，其输出也通过残差连接重新整合。VRCM的输出随后通过一个3x3卷积层，通过残差连接添加到输入特征图中，并传递到下一个残差块或高质量重建阶段。
WKV2DScan模块是一个两阶段的过程，它对输入张量进行处理，生成特征图和注意力权重。WKV2DScan的视觉示意图如图3所示，其中扫描操作在水平和垂直方向上进行，以获取两个方向的特征。在WKV2DScan的每个阶段，都会执行一个WKV-Scan操作。WKV-Scan从位移操作开始，以收集上下文数据，将位移后的张量定义为
。
随后进行加权求和计算，将平移后的张量与原始张量相加，得到
。
表示输入张量的时间相关移动平均，而
是偏置项。
在两次扫描过程中，
、
、
、
、
是可学习的共享权重，而
、
、
、
则分别存在于各自的扫描中，因此它们不仅相互依赖，而且在训练过程中可以从任一扫描中学习得到。
扫描后的特征图随后通过WKV6模块得到，其公式为：
作者的Omni-QuadShift利用多尺度卷积核（1×1、3×3、5×5和7×7）在不同感受野尺度上聚合信息。每个卷积核都是深度可分离和扩张的，能够捕捉局部和全局特征。通过可学习的权重将这些输出组合，以产生一个统一的表示。在这些卷积中，
捕获点特征，
提取局部空间细节，
增强中等范围依赖性，而
扩展感受野以捕捉全局上下文。此外，可学习的权重
调节每个卷积和四向移位（QShift）的贡献。
其中
代表 Head 的尺寸。
首次WKV-Scan操作应用于输入特征图
，从而得到第一组扫描特征图、转置后的输入特征图和注意力权重
，
。对于第二次操作，使用
来获得第二组
，
。最后，WKV2DScan模块输出一个包含扫描特征图和注意力权重的元组，其计算方式如下：
其中
是可学习权重，而
是VRSM的输出。
在获取输出张量
之后，作者将输出重新排列成原始的空间格式。
训练和验证集共包含来自Flickr2K[21]的2650张图像，DIV2K[22]的800张图像，以及RSSCN7[23]的2800张图像，其中训练集占95%。在Flickr2K和DIV2K数据集（DF2K）中，图像分辨率分别为
和
像素。在RSSCN7数据集中，图像分辨率为
像素。训练集中低分辨率（LR）的图像分辨率固定为
，而高分辨率（HR）图像在2倍缩放因子下的分辨率为
，在4倍缩放因子下的分辨率为
。作者的实验旨在确保所有训练的超参数和参数数量大致相同。作者的模型在广泛使用的基准数据集上进行了评估：Set5[24]、Set14[25]、BSD100[26]、Urban100[27]和Manga109[28]。作者还对作者自己的Gangcha MouseHole（GMH）遥感数据集、DF2K数据集的验证集以及公开的RSSCN7遥感数据集进行了测试。
GMH数据集是通过中国青海省的DJI Mavic 3E无人机采集的，其中包括中游海拔30米处的鼠洞325张图片，上游海拔40米处的鼠洞222张图片，以及中上游海拔80米处的草甸鼠洞60张图片。总计607张图片，每张图片分辨率为
像素，分割成
的LR块进行超分辨率任务。
RSSCN7 [23] 是一个用于场景分类的遥感图像数据集，包含2800张分辨率为400×400像素的高分辨率RGB图像。它包括七个类别：草地、森林、农田、工业/商业、河流/湖泊、住宅区和停车场，每个类别各有400张图像。
作者的模型在RGB通道和经过随机旋转、尺寸调整及翻转增强的数据上进行训练。在YCbCr空间中，Y通道的峰值信噪比（PSNR）[29]和结构相似性（SSIM）[30]作为定量测量值。作者实现的OmniRWKVSR模型包含96个通道和16个VRG块。小批量大小设置为16，使用Adam优化器进行优化。在第一阶段，使用
损失函数进行了20000次迭代的预训练，学习率设置为10^-4。作者没有使用学习率调度器，以确保作者的模型在训练过程中不依赖任何技巧来保持其性能。
图5展示了在Setl4和Urban100数据集上，针对4倍尺度因子对所提出方法进行的质量比较。该图展示了高分辨率（HR）、双三次（BICUBIC）、HAT [13]、MambaIR [11]、SwinIR [9]和OmniRWKVSR（作者的工作）在指定图像上的视觉效果。
以“img092”为例，MambaIR和SwinIR在恢复条纹纹理时均存在问题，条纹显得模糊且交错。作者的方法在实现高质量超分辨率重建方面也面临着挑战，因为测试集与训练集之间存在显著差异。然而，OmniRWKVSR在这方面表现相对较好。它能够更准确地恢复条纹特征，使得条纹比其他方法更为清晰和逼真。在确保轻量级复杂度的同时，OmniRWKVSR能够在既定条件下有效地实现更好的超分辨率重建结果。
为了全面评估OmniRWKVSR在图像超分辨率重建质量方面的有效性，作者对其与多种SOTA超分辨率模型进行了详细比较，包括SwinIR [9]、SwinFIR [10]、MambaIR [11]及其改进版本MambaIRv2 [31]。所有比较均在放大倍数
和
下进行。OmniRWKVSR通过整合局部信息处理卷积和通道注意力机制，显著提升了其表征能力。作者在训练和推理时间上对OmniRWKVSR与其他模型进行了比较。
根据表1中的数据，在缩放因子为
的情况下，OmniRWKVsR在多个关键性能指标上表现出优异的性能。在DF2K
RSSCN7验证数据集上，OmniRWKVSR实现了28.0720 dB的峰值信噪比（PSNR）和0.7192的结构相似性指数（SSIM），在所有比较的模型中位居前列。
特别是在Set5数据集上，OmniRWKVSR的PSNR和SSIM分别达到了30.1390 dB和0.8458，这表明其在保留图像细节和结构方面具有明显优势。在Urban100数据集上，OmniRWKVSR的PSNR和SSIM分别为24.5824 dB和0.7160，略低于MambaIRv2，但仍然表现出很强的竞争力。在Manga109数据集上，OmniRWKVSR的PSNR和SSIM分别为27.8930 dB和0.8533，不仅在结构相似性方面超越了所有其他模型，而且为漫画风格图像的超分辨率设定了新的基准，因为保留复杂的线条细节和色彩渐变在这里至关重要。
关于训练和推理速度，如表二所示。所有训练和推理过程均在NVIDIA A800 40G GPU上完成。
推理涉及对表一中所列数据集内的所有图像进行超分辨率处理。从时间复杂度和计算效率的角度来看，作者的模型OmniRWKVSR因其高度优化的CUDA Kernel 和线性计算复杂度，表现出卓越的性能。相比之下，类似SwinIR [9]的模型采用窗口注意力机制来提高计算效率；然而，这种机制仍以二次复杂度运行，导致训练和推理时间变慢。虽然MambaIR及其衍生模型，如SS2D和Attentive SSM [31]也展现出线性复杂度，但它们的计算效率并不及OmniRWKVSR。这种差异可以归因于这些模型中CUDA操作的优化程度较低，与作者的方法相比，导致处理时间更长。通过结合高效的CUDA操作和线性复杂度，OmniRWKVSR能够在保持高性能的同时实现更快的处理速度。
在作者的比较中，作者还考虑了浮点运算次数（FLOPs），这代表了网络在处理
像素图像时所进行的浮点运算总数。较低的FLOPs数值意味着模型在计算上更为高效。与其他模型相比，OmniRWKVSR实现了最低的FLOPs，这表明作者的模型有效地降低了计算复杂性。计算操作使用torch-operation-counter [32]进行。
在超分辨率图像重建（SISR）领域，理解模型中各个组件的影响对于优化性能至关重要。为此，作者对提出的OmniRWKVSR进行了系统性的消融研究，重点关注三个关键组件： FFN 、位移机制和WKV2D扫描模块。
最后，作者将原始的WKV Scan与新提出的WKV2D Scan进行了比较。WKV Scan采用了一维扫描方法，这可能导致无法完全捕捉图像中的空间依赖性。
相比之下，WKV2D Scan将这一方法扩展到了二维扫描机制，从而能够全面捕捉水平和垂直方向的空间依赖性。在作者的Set14数据集上的实验中，WKV2D Scan的表现显著优于WKV Scan，其PSNR达到了30.1390，SSIM为0.8458，而WKV Scan在Set14数据集上的测试结果为PSNR 30.1149，SSIM 0.8449。这表明WKV2D Scan在处理输入特征和生成准确的注意力权重方面具有明显优势。
performance. 通过广泛的实验，作者表明，在PSNR和SSIM等性能指标上，作者的模型在各种基准数据集上，包括Set5、Set14、BSD100、Urban100和Mangal09，均达到了最先进的水平。定性的结果进一步验证了该模型生成清晰无噪超分辨率图像的能力。消融研究突出了每个组件在提升模型整体性能中的重要性。
尽管作者的模型在PSNR和SSIM方面表现出优异的性能，但在可视化结果中仍存在纹理错误和缺乏清晰度等问题。此外，在GMH数据集上的表现并不令人满意。未来的工作将着重解决这些问题。
参考
[1]. Exploring Linear Attention Alternative for Single Image Super-Resolution* .
点击上方卡片，关注
「AI视界引擎」
公众号