k-NN图像Transformer：基于局部注意力的图像恢复方法

21390基于KNN的图像恢复局部注意力0Hunsang Lee 1，Hyesong Choi 2，Kwanghoon Sohn 1，Dongbo Min 2†01 韩国延世大学，2 韩国梨花女子大学0hslee91@yonsei.ac.kr，hyesongchoi2010@gmail.ac.kr，khsohn@yonsei.ac.kr，dbmin@ewha.ac.kr0摘要0最近的研究尝试将非局部操作与CNN或Transformer相结合，在图像恢复任务中取得了显著的性能。然而，全局相似性存在局部性不足和计算复杂度高的问题，其计算复杂度与输入分辨率的平方成正比。局部注意力机制通过引入类似卷积操作的归纳偏差来缓解这些问题。然而，局部注意力仅关注相邻位置，对于图像恢复来说，感受野不足。在本文中，我们提出了一种新的图像恢复注意力机制，称为k-NN图像Transformer（KiT），它纠正了上述限制。具体而言，KiT使用局部敏感哈希（LSH）将k个最近邻补丁分组，并通过执行成对的局部注意力将分组的补丁聚合到每个查询补丁中。通过这种方式，成对操作建立了非局部连接，同时保持了局部注意力的期望特性，即局部性的归纳偏差和对输入分辨率的线性复杂度。所提出的方法在图像去噪、去模糊和去雨等基准测试中优于现有的恢复方法。代码即将发布。01. 引言0图像恢复旨在从各种降解（如噪声、模糊、雨和压缩伪影）中恢复出清晰的图像，这对于图像分类[14,56]、目标检测[22, 46]、分割[4,10]等下游任务的性能有巨大影响。它是一个高度不适定的反问题，因为对于一个降解图像可能存在多个解。最近的恢复工作[17, 36,76]尝试利用卷积神经网络（CNNs）的表示能力建立清晰图像与降解图像之间的映射关系。然而，CNN中使用的一系列局部操作在捕捉长程依赖性方面能力较弱，对整个图像的全局信息进行思考存在一定的局限性。为了扩大感受野，已经提出了增加网络深度[51]、空洞卷积[66]和分层架构[40]等方法，但感受野仍然无法获得全局信息，因为它仅限于局部区域。最近，非局部操作再次成为非学习型恢复方法[5,15]的有力解决方案，非局部神经网络的成功使非局部自相似性计算成为图像恢复的重要线索。[16, 28, 32,37, 38, 43, 53, 77,78]。[58]的非局部自相似性能够在深度网络中捕捉长程依赖性，但与输入特征分辨率相比的二次复杂度限制了网络容量。因此，它仅在相对低分辨率的特定层的特征图中使用[16, 32, 77]。最近，VisionTransformer（ViT）[18]提出了一种将Transformer的全局注意力机制（可以看作是非局部操作）应用于视觉任务的新方法，通过将图像分割成0*本工作得到韩国国家研究基金会（NRF-2021R1A2C2011624和NRF-2021R1A2C2006703）的中期研究者计划以及2021年延世大学研究基金（2021-22-0001）的支持。† 通讯作者：dbmin@ewha.ac.kr0(a)0(b)0(c)0图1. 不同注意力方法的比较：(a) 全局注意力[18, 45,57]在全局范围内计算补丁之间的自相似性，(b) 局部注意力[33,59]在像素级别上测量单个补丁内的自相似性，(c)所提出的方法在像素级别上通过成对的局部注意力聚合相似的k个补丁。0图像恢复的目标是从各种降解（如噪声、模糊、雨和压缩伪影）中恢复出清晰的图像，这对于图像分类[14,56]、目标检测[22, 46]、分割[4,10]等下游任务的性能有巨大影响。它是一个高度不适定的反问题，因为对于一个降解图像可能存在多个解。最近的恢复工作[17, 36,76]尝试利用卷积神经网络（CNNs）的表示能力建立清晰图像与降解图像之间的映射关系。然而，CNN中使用的一系列局部操作在捕捉长程依赖性方面能力较弱，对整个图像的全局信息进行思考存在一定的局限性。为了扩大感受野，已经提出了增加网络深度[51]、空洞卷积[66]和分层架构[40]等方法，但感受野仍然无法获得全局信息，因为它仅限于局部区域。最近，非局部操作再次成为非学习型恢复方法[5,15]的有力解决方案，非局部神经网络的成功使非局部自相似性计算成为图像恢复的重要线索。[16, 28, 32, 37, 38, 43,53, 77,78]。[58]的非局部自相似性能够在深度网络中捕捉长程依赖性，但与输入特征分辨率相比的二次复杂度限制了网络容量。因此，它仅在相对低分辨率的特定层的特征图中使用[16,32, 77]。最近，VisionTransformer（ViT）[18]提出了一种将Transformer的全局注意力机制（可以看作是非局部操作）应用于视觉任务的新方法，通过将图像分割成21400一组不重叠的块，嵌入到特征空间中，并将它们输入到多个变换器层中，以建模块块之间的全局自相似性（图1（a））。ViT在图像分类任务中取得了准确性和计算复杂性之间的令人满意的平衡，但是相对于输入特征分辨率的二次复杂性使得将变换器应用于密集预测任务几乎不可行。为了克服这个限制，与ViT不同，ViT在整个网络中保持特征分辨率，一些方法[45,57]提出了一种分层架构，以利用适用于密集预测任务的多尺度特征图。然而，它们只关注捕捉全局自相似性，而在探索对于图像恢复至关重要的局部性方面，它们的能力不及CNN。在这种情况下，已经提出了许多方法来将局部性的归纳偏差引入变换器架构[30,33,59,61,63]。其中，最近的工作中考虑了局部注意力[31,33,54,59,67]，但代价是限制变换器中的感受野。这些方法提出了局部自注意模块，实现了对输入特征分辨率的线性复杂性（图1（b））。由于它们仅在局部块内约束自注意计算，因此还需要应用一种移位方法[31,33,59]来在不重叠的块之间交换信息。然而，它们只考虑了相邻的块，因此仍然具有不足的感受野。在本文中，我们提出了一种新颖的非局部图像恢复方法，称为k-NN图像变换器（KiT），它成功地在明确建立非局部连接的同时捕捉了局部性，方法是考虑k个最近邻（k-NN）块的局部注意力。为了弥补局部注意力中固有的长程依赖性的缺失，所提出的方法考虑了生成不同位置的块之间的非局部连接的k个匹配块。具体而言，KiT首先通过k-NN匹配为每个基准块搜索一组相似块，然后将基准块设置为查询，将k个匹配块设置为键和值，以应用局部的成对注意力，如图1（c）所示。这使得我们的方法能够在整个图像上应用局部注意力，同时保持与特征分辨率的线性复杂性。此外，局部性的归纳偏差增强了局部特征提取能力。如图2所示，我们的方法由一系列k-NN变换器块（KTB）组成，并采用U形分层架构，以有效地利用多尺度特征。对各种图像恢复任务的全面实验证明了所提出方法的有效性。02. 相关工作0最近，随着非局部神经网络的成功，一些方法[16,32,77]尝试将非局部操作与CNN结合起来，通过建立全局自注意力来进行图像恢复任务。然而，其昂贵的计算成本限制了特征图或网络深度的空间分辨率。为了减少计算成本，在输入特征图中使用了稀疏连接而不是全连接[28,37,38,43,53,78]。N3Net [43]和GCDN[53]以可学习的方式找到在嵌入空间中接近的k个最近邻，并将它们聚合起来进行高效计算。DAGL[38]根据图像内容动态选择每个查询的邻居数量，这些邻居具有不同的分布。IGNN [78]和CPNet[28]通过在跨尺度特征图中考虑稀疏性和跨尺度块循环性来找到k个最近邻块。然而，前面提到的方法在k-NN匹配上具有二次复杂性，严重降低了整个过程的速度。NLSN[37]通过使用局部敏感哈希（LSH）进行局部敏感的非局部稀疏注意力，将k-NN匹配过程的复杂性降低为渐进线性。但是，由于NLSN[37]在像素级别上近似了全局注意力的全连接，因此无法捕捉局部信息。0视觉Transformer。在[18]中，将最初用于自然语言处理的Transformer架构[55]应用于图像分类任务。这种方法被称为视觉Transformer（ViT），通过对图像补丁应用全局注意力来捕捉长程依赖关系，但由于输入空间分辨率的二次复杂度，不适用于密集预测。与ViT不同的是，分层架构采用逐渐降低特征分辨率的方式，更有效地进行密集预测[9, 45, 57,62]。PvT[57]使用空间降采样注意力（SPA）层构建金字塔特征图。IPT[9]和DPT[45]提出了编码器-解码器架构来恢复细粒度预测。然而，这些基于全局注意力的方法缺乏探索图像恢复中必要的局部性能力。最近，SwinTransformer[33]通过移位方法利用局部注意力进行补丁连接，在目标检测和分割任务上取得了竞争性的性能，并具有较低的复杂度。由于局部注意力模块仅在相邻元素之间生成注意力权重，计算复杂度与空间分辨率呈线性关系，并且注入了局部性的归纳偏差。Uformer[59]和SwinIR[31]采用局部注意力进行图像恢复任务，取得了令人印象深刻的结果。然而，移位方法仍然具有有限的感受野。21410图2. k-NN图像Transformer（KiT）的整体架构：（a）采用U形分层架构进行图像恢复。（b）每个阶段有两个k-NN TransformerBlocks（KTBs）和一个插值层。对于跳跃连接，编码器中第i个阶段的输出特征与解码器中第（4-i）个阶段的输出特征进行连接。（c）KTB由层归一化（LN）、k-NN局部注意力（KLA）和由深度卷积（DW）和多层感知机（MLP）组成的前馈网络（FFN）组成。0由于考虑了邻域补丁，传统的图像恢复方法在建立非局部连接时会丢失非局部连通性。相反，我们提出的方法通过对k个最近邻补丁进行成对局部注意力计算来建立非局部连通性。这使得我们能够以线性复杂度对空间分辨率进行非局部连通性建模，同时在注意力模块中捕捉局部性。03.提出的方法03.1.问题陈述和概述0众所周知，非局部自相似性在图像恢复任务中非常有效[9,16, 32, 37, 43,77]。由于相似模式在图像中全局分布，这要求能够捕捉长程依赖关系。ViT[18]将原始Transformer[55]的注意力机制直接应用于图像补丁的序列中。对于给定的输入X∈RHW×Cin，它将其分割成不重叠的补丁，并重塑为扁平化的2D补丁序列Xp∈RN×r2Cin，其中HW是输入特征图的空间分辨率，Cin是输入特征图的通道数，N=HW/r2，r是补丁大小。通过分割补丁之间的点积进行全局注意力计算，表示为：0O = softmax(φ(Xp)θ(Xp)T)0√0C0ψ(Xp)0可学习的投影函数φ、θ和ψ将Xp投影到查询、键和值中。输出O∈RN×Cout，其中Cout是输出通道的大小，是通过计算投影查询和键之间的亲和矩阵得到的投影值的加权和。由于通常设置C、Cin和Cout相同，我们将它们表示为C。尽管全局注意机制能够很好地建立长程依赖关系，但输入特征分辨率的二次复杂度O(r2N2C)使其...0对于密集预测任务，很难充分利用全局注意力。局部注意力机制[31, 33, 54, 59,67]通过在局部补丁内计算注意力来降低复杂度。输入特征图X被分割成非重叠的补丁，满足X = {xi∈Rr2×C | i = 0, ...,N - 1}。局部注意力在每个补丁内计算。0oi = softmax(φ(xi)θ(xi)T)0√0C0ψ(xi), (2)0其中oi是与xi对应的输出补丁。需要注意的是，可学习的投影函数φ、θ和ψ将r2个大小为C的元素投影，与ViT将N个大小为r2C的元素投影不同，并且对所有补丁共享。局部注意力实现了对输入特征分辨率的线性复杂度O(r4NC)。然而，由于等式（2）分别应用于每个补丁，因此补丁之间没有交换信息。因此，我们采用了一种移位方法[31, 33,59]，以在扩大的感受野中在相邻补丁之间建立补丁连接。然而，由于只有相邻补丁对查询补丁有贡献，感受野仍然受限。我们通过在局部注意力的计算中利用k-NN来克服这个限制，称为k-NN局部注意力。为了在计算局部注意力时实现非局部连接，我们利用k-NN搜索来寻找用于计算局部注意力的一组补丁候选。通过在查询补丁和k个匹配补丁之间进行成对的局部注意力计算，所提出的方法在建立对图像恢复至关重要的局部性的同时高效地捕捉到局部连接。03.2. 总体流程0图像恢复的整体框架如图2所示。为了恢复受损图像，我们首先对受损输入图像Id进行三次卷积，然后通过编码器网络和解码器的三个阶段进行处理。21420图3.所提出的k-NN局部注意力（KLA）：首先进行k-NN搜索，以在整个特征图中找到k个相似补丁。这里，我们使用局部敏感哈希（LSH）进行高效的k-NN搜索。LSH为补丁分配哈希值，并按升序对补丁进行排序。然后，使用块大小为k对排序后的补丁进行分区。在每个块中，执行补丁之间的成对局部注意力计算，然后通过局部注意力输出的加权和计算最终输出。这里，当使用单个块时，给出了对˜x1的KLA的示例。实际上，先前的块也对包含查询补丁的当前块有贡献。有关更多详细信息，请参阅第3.3节。0网络。每个阶段由补丁分区、k-NN TransformerBlocks（KTB）和插值层组成。补丁分区操作将输入特征图X分割成非重叠的补丁，补丁大小为r，满足X = {xi∈Rr2×C| i = 0, ..., N -1}。在KTB中，分割的补丁被归一化并输入到k-NN局部注意力（KLA）进行非局部聚合。KLA首先基于相似模式的补丁频繁出现在图像中，并且聚合它们有助于图像恢复。需要注意的是，现有的稀疏注意力方法[26,37]将像素聚类到单独的组中以近似全局注意力，因此自然地缺乏局部性，这是全局注意力的一个缺点，在它们的注意力模块中仍然存在。另一方面，所提出的方法通过使用成对的局部注意力计算查询补丁与k个补丁之间的稀疏相似性，并且局部性的归纳偏差反映在所提出的注意力模块中。所提出的网络具有U形层次结构，以考虑各种尺度的模式。聚合特征通过插值层传递（编码器下采样和解码器上采样）。在解码器的每个阶段，输入特征图与相应的编码器特征连接以恢复细节。在网络的末尾，进行三次卷积以从输出特征图预测恢复图像。03.3. KTB：k-NN Transformer块0对每个块应用层归一化（LN），然后使用k-NN局部注意力（KLA）将查询块和匹配块作为键和值进行局部注意力。为了增强网络的局部性，使用深度卷积（DW）[13]和多层感知机（MLP）在前馈网络（FFN）[30]中，如图所示0如图2（c）所示。形式上，KTB可以写成：0ˆXl=KLA（LN（Xl-1））+Xl-1，（3）0Xl=FFN（LN（ˆXl））+ˆXl，（4）0其中FFN（X）=MLP（DW（MLP（X）））。在每个阶段的第l个块（l =0，...，b-1）中，输出特征图Xl-10从前一个块中归一化并输入KLA。通过k个相似块特征和残差连接进行非局部聚合计算中间特征ˆXl。瓶颈阶段与KTB相同，只是不使用插值层，并且将k设置为1。k最近邻匹配。暴力k-NN匹配需要计算两个块之间的成对距离。由于这个成对距离涉及到输入长度的二次复杂度，我们利用局部敏感哈希（LSH）[2]，其具有线性计算复杂度。LSH将分割块投影到单位超球体中以建立桶。假设有m个哈希桶，通过将随机旋转矩阵R∈RN×m/2乘以球面投影块x来分配哈希值L（x）：0L（x）=argmax（[xR;−xR]），（5）0其中[∙;∙]表示两个元素的连接。通过这个哈希操作，具有高相关性的块很可能会接收相同的哈希值（在同一个哈希桶中），反之亦然。然而，由于LSH依赖于随机旋转矩阵，相似的块有时可能落入不同的哈希桶中。为了解决这个问题，采用多轮LSH，其中LSH在不同的随机旋转矩阵上应用h次。KLA：k-NN局部注意力。如图3所示，相似的块根据分配的哈希值进行分组。只有具有相同哈希值的块才能贡献op =�j∈Npwp,j · op,j,(9)wp,q =φ(−→˜xp) · θ(−→˜xq)j∈Np φ(−→˜xp) · θ(−→˜xj).(10)Lchar =�∥I − (Id + Ir)∥2 + ϵ2,Ledge =�∥△I − △(Id + Ir)∥2 + ϵ2,L = Lchar + λLedge.(11)21430为了有效地对查询块进行k-NN匹配，我们首先根据哈希值对块进行排序，然后将排序的块分成每个包含k个块（等于NN块的数量）的块，以便仅在同一块中的块上执行局部注意力。我们将π：n−→n定义为按哈希值升序排序块的排列：0π（xp）<π（xq）�L（xp）≤L（xq），（6）0为了简单起见，我们将˜x定义为排序的块，其中˜xp等于xπ（p）。然后，第i个块Pi（i = 0，...，N/k）包含k个块，0Pi={˜xki，˜xki+1，˜xki+2，...，˜xki+k-1}，（7）0然后，在块内的块对中进行局部注意力以进行非局部聚合。将排序的输入块˜X投影到查询、键和值上，使用可学习的投影函数φ，θ和ψ：Rr2×C→Rr2×C。由于块中有k个块，因此进行k2次局部注意力。对于作为查询的第p个块的成对局部注意力输出，其中第q个块用作键和值，定义如下：0op,q=softmax（φ（˜xp）θ（˜xq）0√0C）ψ（˜xq），（8）0例如，在图3的P0块中，将φ（˜x1）表示为查询块，有k个输出块{o1，j | j = 0，...，k - 1}。与ViT[18]不同，ViT采用所有块的自注意力，我们采用局部注意力，但是k次并以成对的方式计算两个块（查询和键）之间的关联矩阵，以增强局部性，如图1（c）所示。由于计算了查询块的k个输出块，因此应将成对输出聚合到查询块中。输入块˜xp的输出块op通过加权求和计算为：0其中 w p,j 是补丁之间的相对相似度，N p 是查询补丁 ˜ x p所属的块的补丁索引集合。0这里，−→ x p  ∈  R r 2C 表示补丁 ˜ x p的扁平化形式。由于哈希桶中的补丁数量在实践中通常不能被块大小整除，具有相同哈希值的补丁可能落入附近的块中。为了处理这个问题，类似于[26]，我们允许前面的块对包含查询补丁的当前块进行贡献，例如 P i − 1 对于 Pi。因此，对于每个查询补丁，局部注意力被执行 2 k 次。03.4. 训练损失0按照现有的图像恢复方法[32，0除了从退化输入图像 I d 预测残差图像 I r外，该网络还预测了一个残差图像 I r 。目标是恢复满足 I=  I d  +  I r 的清晰图像 I。我们利用Charbonnier损失 [8]L char 和边缘损失 L edge 来优化网络。0其中 ϵ 对于所有实验都被经验性地设置为 10^-3，△表示Laplacian函数。总损失 L 由 L char 和 L edge定义，其中超参数 λ 控制两个损失的比例。04. 实验04.1. 实现细节0KiT是用PyTorch实现的。我们使用AdamW优化器[34]在16张裁剪为128×128的图像批次上训练了整个网络，共进行了300个epoch。初始学习率设置为1×10^-4，采用线性预热策略和余弦退火策略来降低学习率。默认情况下，块大小k（等于NN补丁的数量）和补丁大小r设置为4。在瓶颈阶段，由于只有少量补丁（例如当HW为256×256时，补丁数量为4×4），k被设置为1。每个阶段中的KTB数量b设置为2。在KLA中，哈希数量h设置为4，用于多轮LSH。我们在各种图像恢复任务上验证了所提方法的性能，如图像去噪、去模糊和去雨。性能评估中，对于去噪和去模糊，使用RGB空间上的PSNR和SSIM进行测量。在去雨中，根据之前的工作[24，72]，在YCbCr颜色空间的Y通道上进行评估。更多结果请参见补充材料。04.2. 图像去噪0我们使用包含320个具有真实噪声的高分辨率图像的SIDD[1]数据集来训练KiT。表1显示了在SIDD [1]和DND[42]数据集上进行真实噪声去除的定量评估结果。评估结果包括经典的去噪方法[15]、基于CNN的方法[3，6，11，23，25，68，70-72，76]、基于自注意力的方法[38]和基于Transformer的方法[59]。由于DND[42]数据集没有提供真实标签，因此结果是从官方基准测试中获得的。所提方法在两个数据集上都优于现有的最先进方法。(a)(b)(c)(d)(e)(f)(g)SIMBM3D [15]25.650.68534.510.851DnCNN [76]23.660.58332.430.790MLP [6]24.710.64134.230.833CBDNet [23]30.780.80138.060.942RIDNet [3]38.710.95139.260.953AINDNet [25]38.950.95239.370.951VDN [68]39.280.95639.380.952SADNet [7]39.460.95739.590.952DANet [69]39.470.95739.580.955CycleISP [70]39.520.95739.560.956MPRNet [72]39.710.95839.800.954MIRNet [71]39.720.95839.880.956NBNet [11]39.750.97339.890.955DAGL [38]--39.830.957Uformer [59]39.770.97039.960.956KiT39.800.97239.960.956SIMDeepDeblur [39]29.230.91625.730.874SRN [52]30.260.93428.360.915PSS-NSC [21]30.920.94229.110.913DMPHN [73]31.200.94529.090.924SAPHN [50]32.020.95329.980.930MT-RNN [41]31.150.94529.150.918SPAIR [44]32.060.95330.290.931MPRNet [72]32.660.95930.960.939MIMO-UNet [12]32.680.959--KiT32.700.95930.980.94221440噪声图像 (a) (b) (c) (d) (e) (f) (g)0图4. SIDD [1]数据集上的视觉比较：(a) 裁剪图像，(b) RIDNet [3]，(c) CycleISP [70]，(d) MPRNet [72]，(e) Uformer [59]，(f)KiT，以及(g) 真实图像。0方法 SIDD DND0表1. SIDD [ 1 ]和DND [ 42]数据集上的定量结果。粗体和下划线数字分别表示最佳和次佳结果。0在SIDD数据集上取得了竞争性能，并在DND数据集上取得了竞争性能。由于DND数据集没有提供任何训练数据，因此在表1中的性能是使用在SIDD数据集上训练的网络实现的，证明了所提方法的鲁棒性。图4显示了使用各种最先进方法进行去噪的图像。现有方法输出的恢复图像丢失了细节，而所提方法通过捕捉具有非局部连接性的局部性能力成功地恢复了退化图像的细节结构。04.3. 图像去模糊0表2报告了GoPro数据集[ 39]上的图像去模糊性能。GoPro数据集提供了合成的模糊图像，每个图像是通过平均连续的清晰图像获得的。训练时使用了GoPro [ 39 ]数据集的2,103个图像和Go-0方法 GoPro HIDE0表2. GoPro [ 39 ]和HIDE [ 49]数据集上的定量结果。网络在GoPro数据集上进行了训练。0Pro [ 39 ]，评估了RealBlur [ 48 ]的3,758个图像和HIDE [49]的2,025个图像。在PSNR和SSIM指标上的出色表现验证了所提方法对恢复模糊图像也是有益的。图5显示了来自GoPro [ 39]数据集中模糊伪影的恢复图像。很容易发现，我们的结果捕捉到了锐利的细节，而其他方法无法处理高频细节。04.4. 图像去雨0按照[ 24 ]的实验设置，使用了来自多个数据集[ 20 ,  29 ,64 ,  74 ,  75]的13,712对清晰-雨图像进行网络训练。我们在五个数据集上评估了去雨结果，包括Test100 [ 75 ]，Rain100H [ 64]，Rain100L [ 64 ]，Test2800 [ 20 ]和Test1200 [ 74]。尽管利用额外的失真引导网络的SPAIR [ 44]在PSNR方面显示出与所提方法相竞争的结果，但所提的KIT在SSIM方面表现更好，证明了细节可以更好地恢复。如图4和图6的结果所示，所提方法在处理重复纹理方面具有优势，这要归功于基于KLA的聚合。在雨天的图像中，图像中存在许多具有相似模式的补丁，所提方法显示出-(a)(b)(c)21450模糊输入0真实值0图5. GoPro [ 39 ]数据集上的视觉比较：(a) MPRNet [ 72 ]，(b) MIMO-UNet [ 12 ]，(c) KiT.0多雨的输入0(a)0(b)0(c)0(d)0(e)0图6.在Test100 [75]数据集上的视觉结果：(a) DerainNet [19]，(b) PreNet [47]，(c) RESCAN [27]，(d) MPRNet [72]和(e) KiT。0定量和定性地与评估方法进行了对比，表现出了出色的性能。04.5.消融研究0我们进行了消融研究，以分析我们的方法在各个方面的性能。所有实验都在SIDD[1]上进行图像去噪任务。计算复杂度。所提出的方法由k-NN搜索的多轮哈希、特征投影和成对局部注意力组成。多轮哈希通过乘以随机旋转矩阵R对输入特征补丁进行处理，其复杂度为O(hNCm)。然后，每个补丁都通过可学习的投影函数φ、θ和ψ投影到查询、键和值：R r 2 × C → R r 2 ×C，其复杂度为O(NC^2)。计算局部注意力的复杂度为0所有补丁对的复杂度为O(khr^4N)。因此，所有操作对于输入特征分辨率的计算都是线性的。0k-NN补丁的可视化。我们的方法旨在在实现非局部连接性的同时保留细节，通过聚合具有相似特征的补丁来实现。为了验证这一点，我们进一步在图7中可视化了KLA中使用的补丁。最左边的图像被分成非重叠的补丁，其中用彩色框标记的补丁用于可视化查询补丁。由于KLA利用LSH进行k-NN搜索，相似的k个补丁在右侧的图像中被分组到一个块中。属于同一块的k个补丁也用相同的彩色框标记。红色和绿色框具有相似的模式，而蓝色框中的补丁包括非纹理区域，证明LSH有效地找到了视觉上相似的补丁。816138.7938.9238.9638.9638.96239.5839.6939.7539.7639.78439.6939.7639.8039.8039.81839.7439.7839.8039.8139.821639.7539.7839.8039.8139.8221460方法 Test100 [75] Rain100H [64] Rain100L [64] Test2800 [20] Test1200 [74] 平均0PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM0DerainNet [19] 22.77 0.810 14.92 0.592 27.03 0.884 24.31 0.861 23.38 0.835 22.48 0.796 SEMI [60] 22.35 0.788 16.560.486 25.03 0.842 24.43 0.782 26.05 0.822 22.88 0.744 DIDMDN [74] 22.56 0.818 17.35 0.524 25.23 0.741 28.130.867 29.65 0.901 24.58 0.770 UMRL [65] 24.41 0.829 26.01 0.832 29.18 0.923 29.97 0.905 30.55 0.910 28.02 0.880RESCAN [27] 25.00 0.835 26.36 0.786 29.80 0.881 31.29 0.904 30.51 0.882 28.59 0.857 PreNet [47] 24.81 0.851 26.770.858 32.44 0.950 31.75 0.916 31.36 0.911 29.42 0.897 MSPFN [24] 27.50 0.876 28.66 0.860 32.40 0.933 32.82 0.93032.39 0.916 30.75 0.903 MPRNet [72] 30.27 0.897 30.41 0.890 36.40 0.965 33.64 0.938 32.91 0.916 32.73 0.921SPAIR [44] 30.35 0.909 30.95 0.892 36.93 0.969 33.34 0.936 33.04 0.922 32.91 0.9260KiT 30.26 0.904 30.47 0.897 36.65 0.969 33.85 0.941 32.81 0.918 32.81 0.9290表3.图像去雨的定量结果。用于评估的常用五个数据集[20,64,74,75]。0PSNR h0k0表4.补丁数量k和哈希轮数h的消融研究。0k和h的数量。k的块大小决定了用于执行查询块的局部注意力的最大补丁数量。h个哈希轮用于降低相似补丁落入不同哈希桶的概率。由于这两个超参数不影响网络参数的数量，只与内存和计算复杂度有关，因此可以根据计算资源灵活调整网络容量。表4显示了根据这两个超参数提出的方法的去噪性能。当这两个超参数设置为16时，获得了最佳性能，但是我们将k和h设置为4，因为它具有相对较低的计算量并具有可比性能。共享查询和键。在现有的使用LSH进行稀疏全局注意力的方法[26，37]中，查询和键的投影函数应该共享，即φ =θ，因此共享注意力几乎总是压倒其他位置的查询与键的点积。相反，我们的方法对于共享投影函数没有这样的约束。为了进行消融研究，我们在表5中进行了共享查询和键的投影函数的额外实验，表示为KiTS。使用共享投影时，整体性能略有降低。05. 结论。0在本文中，我们提出了一种新颖的非局部图像恢复方法。具体来说，k-NN局部注意力。0(a)0(b)0图7.k-NN补丁的可视化：(a)输入图像和(b)k-NN补丁。k-NN补丁是通过k = 4的LSH发现的。属于同一块的补丁用相同颜色的框标记。0方法SIDD。0PSNR SSIM。0KiT 39.80 0.9720KiT S 39.75 0.9690表5. 共享投影φ = θ (KIT S)的消融研究。0(KLA)在相似补丁之间进行成对局部注意力与k-NN匹配。KLA在建立非局部连接性时保持了局部性的归纳偏差，并且对输入空间分辨率具有线性计算复杂度。所提出的方法在各种图像任务上的定量/定性性能方面优于现有方法。由于确定网络容量的NN补丁数量和哈希轮数与网络参数无关，因此可以灵活调整网络容量。0限制。该方法仅考虑了相同尺度的补丁之间的成对局部注意力。跨尺度注意力可以是一种有趣的方法，可以进一步改善恢复性能。我们将继续研究将跨尺度注意力整合到基于交叉位置的模型中的网络。21470参考文献。0[1] Abdelrahman Abdelhamed, Stephen Lin, and Michael SBrown.用于智能手机相机的高质量去噪数据集。在IEEE计算机视觉与模式识别会议上，页码1692-1700，2018年。5, 6, 70[2] Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razen- shteyn, andLudwig Schmidt.Angular距离的实用和最优LSH。arXiv预印本arXiv:1509.02897，2015年。40[3] Saeed Anwar and Nick Barnes.具有特征注意力的真实图像去噪。在IEEE/CVF国际计算机视觉会议上，页码3155-3164，2019年。5, 60[4] Vijay Badrinarayanan, Alex Kendall, and Roberto Cipolla.Segnet:用于图像分割的深度卷积编码器-解码器架构。IEEE模式分析与机器智能交易，39(12)：2481-2495，2017年。10[5] Antoni Buades, Bartomeu Coll, and J-M Morel.一种用于图像去噪的非局部算法。在2005年IEEE计算机学会计算机视觉与模式识别会议上，卷2，页码60-65，2005年。10[6] Harold C Burger, Christian J Schuler, and StefanHarmeling.图像去噪：普通神经网络能与bm3d竞争吗？在2012年IEEE计算机视觉与模式识别会议上，页码2392-2399，2012年。5, 60[7] Meng Chang, Qi Li, Huajun Feng, and Zhihai Xu.空间自适应网络用于单幅图像去噪。在欧洲计算机视觉会议上，页码171-187，2020年。60[8] Pierre Charbonnier, Laure Blanc-Feraud, Gilles Aubert,and Michel Barlaud.用于计算成像的两种确定性半二次正则化算法。在第1届国际图像处理会议上，卷2，页码168-172，IEEE，1994年。50[9] Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu,Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu,and Wen Gao.预训练图像处理变压器。在IEEE/CVF计算机视觉与模式识别会议论文集中，页码12299-12310，2021年。2, 30[10] Liang-Chieh Chen, George Papandreou, IasonasKokkinos, Kevin Murphy, and Alan L Yuille. Deeplab:带有深度卷积网络、空洞卷积和全连接CRF的语义图像分割。IEEE模式分析与机器智能交易，40(4)：834-848，2017年。10[11] Shen Cheng, Yuzhi Wang, Haibin Huang, Donghao Liu,Haoqiang Fan, and Shuaicheng Liu.Nbnet：噪声基础学习用于图像去噪与子空间投影。在2021年IEEE/CVF计算机视觉与模式识别会议上，第4896-4906页。5月6日0[12] Sung-Jin Cho, Seo-Won Ji, Jun-Pyo Hong, Seung-WonJung, and Sung-Jea Ko.重新思考单幅图像去模糊中的由粗到细方法。在2021年IEEE/CVF国际计算机视觉会议上，第4641-4650页。6月7日0[13] François Chollet.Xception：深度学习与深度可分离卷积。在IEEE会议上。0在2017年IEEE计算机视觉与模式识别会议上，第1251-1258页。4月0[14] Dan Ciregan, Ueli Meier, and Jürgen Schmidhuber.用于图像分类的多列深度神经网络。在2012年IEEE计算机视觉与模式识别会议上，第3642-3649页。1月0[15] Kostadin Dabov, Alessandro Foi, Vladimir Katkovnik, andKaren Egiazarian.通过稀疏3D变换域协同滤波进行图像去噪。IEEE图像处理期刊，16(8)：2080-2095，2007年。1月2月5月6月0