Rethinking Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising

基于Transformer的盲点网络自监督图像去噪算法的再思考

在这里插入图片描述

论文链接:https://arxiv.org/abs/2404.07846

项目链接:https://github.com/nagejacob/TBSN

Abstract

盲点网络(Blind-spot networks,BSN)是自监督图像去噪(self-supervised image notifying,SSID)中常用的神经网络结构,但大多数现有的BSN都是在卷积层上进行的。虽然transformers在许多图像恢复任务中表现出克服卷积限制的潜力,但其注意机制可能违反盲点要求,从而限制了其在BSN中的适用性。为此,我们提出分析和重新设计通道和空间注意以满足盲点要求。具体来说,通道自注意可能会泄漏多尺度结构中的盲点信息,由于下采样将空间特征混洗到通道维度中。为了解决这个问题,我们将通道分成几组,分别进行通道注意。对于空间自注意,我们在注意矩阵中应用了一个精心设计的掩模来限制和模仿膨胀卷积的感受野。基于重新设计的通道和窗口注意,我们构建了一个基于Transformer的盲点网络(TBSN),该算法具有较强的局部拟合能力与全局视角能力,并引入了一种知识蒸馏策略,将TBSN蒸馏成较小的去噪器,在保持性能的同时提高了计算效率。在真实图像去噪数据集上的大量实验表明,TBSN极大地扩展了感受野,与现有的SSID方法相比表现出良好的性能。

Introduction

图像去噪是一项基本的低层视觉任务,旨在从噪声观测中恢复隐含的干净图像。随着卷积神经网络的发展,基于学习的方法(Mao,Shen,and Yang 2016; Zhang et al 2017; Tai et al 2017; Lin et al 2024)与传统方法相比,(Buades,Coll, and Morel 2005; Dabov et al 2007)。为了便于网络训练,通常将加性白色高斯噪声(AWGN)与噪声-干净图像对合成用于监督学习。由于AWGN和相机噪声之间的分布间隙,它们在现实场景中表现出降低的去噪性能。一个可行的解决方案是捕获具有严格对齐的噪声-干净对的数据集(Plotz and Roth 2017; Abdelhamed,Lin and Brown 2018)用于网络训练(Guo et al 2019; Kim et al 2020)。然而,数据收集过程需要严格控制的环境和大量的人力,这是不太实际的。

近年来,自监督图像去噪(SSID)(Krull, Buchholz, and Jug 2019; Lee, Son, and Lee 2022; Li et al 2023; Wang et al 2023b; Jang et al 2023a; Zhang et al 2023)方法被提出,旨在规避对配对数据集的依赖。Noise2Void的开创性工作随机屏蔽了噪声输入的一些位置,并训练网络从周围的位置重建它们。在像素独立噪声的情况下,网络学习预测没有随机噪声的掩蔽像素,即干净像素,从而实现自监督去噪。盲点网络(BSN)(Laine et al 2019; Wu et al 2020; Byun, Cha, and Moon 2021)进一步利用专用设计的网络架构来实现掩模模式场,这在性能和训练效率上都显示出优越性。对于真实场景中的空间相关噪声去除,一些工作(Wu et al 2020; Zhou et al 2020)建议首先使用像素重排下采样(PD)打破噪声相关性,然后使用BSN去噪。(Lee,Son,and Lee 2022; Wang et al 2023 b; Jang et al 2023 a)已经显示出在噪声去除和细节保留之间的更好的权衡。

现有的BSN结构多为卷积神经网络(CNN),但卷积运算在捕获长程相关性方面能力有限,卷积滤波器的静态权值不能灵活地适应输入内容(Liang et al,2021; Zamir et al,2022; Chen et al,2023 b),这些限制可以通过Transformer模型来缓解(Vaswani et al,2017)。然而,Transformer算子可能违反盲点要求,并导致对噪声输入的过拟合。尽管存在困难,但很少有人尝试将变换器应用到BSN中。例如,LG-BPN(Wang et al 2023 b)将通道自注意(Zamir et al 2022)用于全局特征增强,但仍然使用卷积层进行局部信息集成。SwinIA(Papkov and Chizhov 2023)实现了一种基于swin-transformer(Liu et al 2021 b)的改进窗口注意的BSN,但受限于对盲点的要求,只能利用注意层中噪声输入的浅层特征,表现出较低的性能,可见在BSN中发挥transformer的有效能力是非常具有挑战性的。

在本文中,我们提出分析空间和通道自注意机制,并重新设计它们以满足盲点的要求。对于通道自注意,我们观察到,简单地应用它可能会泄漏盲点信息,尤其是在多这种架构的深层特征已经被多次下采样,并且空间信息被混洗到通道维度。通道可能在盲点处泄漏空间信息,导致对噪声输入的过拟合。经验发现,当通道维数大于空间分辨率时,会出现这种效应。为了消除这种不良影响,我们将通道分成组,并对每组单独进行通道注意,其中组通道数控制在小于空间分辨率。对于空间自注意,我们通过限制窗口的感受野来重新设计窗口的注意力,以保持对盲点的要求。具体地说,我们对注意力矩阵应用了固定的掩模,使得每个像素只能注意偶数坐标上的像素。结合所设计的空间和通道自注意机制,我们提出了一种扩展的Transformer注意力块(DTAB)。我们将DTAB嵌入到基于编码器-解码器的U-Net架构中,从而提出了基于Transformer的盲点网络(TBSN)。

此外,BSN结构由于需要额外的设计来满足盲点要求,计算效率低下,随着模型规模的增加和复杂的后细化过程,计算效率会变得更低(Lee,Son and Lee 2022)。然而,一些简单有效的监督去噪器有可能达到最先进的SSID方法的性能。在这项工作中,我们利用这一特性,探索了一种知识蒸馏策略,以减少推理过程中的计算成本。具体来说,我们将预训练的TBSN的结果视为伪ground truth,并将其作为监督来训练一个普通的U-Net,即TBSN2UNet。

在真实世界的去噪数据集上进行了大量的实验(Abdelhamed,Lin and Brown 2018; Plotz and Roth 2017)来评估TBSN和TBSN 2UNet的有效性。如图1所示,受益于提出的空间和通道自注意机制,TBSN增强了局部自适应能力,大大扩展了感受野,对静态噪声有较好的抑制作用。此外,TBSN 2UNet保持了TBSN的性能,同时显著降低了推理成本。

在这里插入图片描述

我们的主要贡献可以总结如下:

  • 我们提出了一种基于Transformer的盲点网络(TBSN),它包含空间和通道自关注,用于自监督图像去噪。
  • 对于通道自注意,我们发现当通道数变大时可能会泄漏盲点信息,因此我们对每个划分的通道组分别执行它以消除这种不利影响。对于空间自注意,我们引入掩膜窗口注意,其中对注意矩阵应用精心制作的掩膜以保持盲点要求。
  • 大量的实验表明,TBSN在真实世界的图像去噪数据集上实现了最先进的性能,而我们从TBSN中提取的U-Net有效地降低了推理过程中的计算成本。

Related Work

深度图像去噪

基于学习的方法的发展(Zhang et al 2017)已显示出优于传统补丁程序的上级性能(Buades, Coll, and Morel 2005; Dabov et al 2007)关于合成高斯去噪。(Mao,Shen,and Yang 2016; Tai et al 2017; Liu et al 2018)进一步提出了改进去噪能力的方法。NBNet(Cheng et al 2021)通过在特征空间中学习一组重建基,提出了一种噪声基网络。(Liu et al 2021 a)提出了一种基于规范化流架构的轻量级去噪网络。(Vaswani et al 2017)已成功应用于视觉任务(Dosovitskiy et al. 2020; Liu et al. 2021b)。对于图像去噪,使用大规模图像预训练(Chen et al 2021)和Swin Transformer架构(Liang et al 2021)研究了Transformer。Restormer (Zamir et al 2022)和Uformer(Wang et al 2022 a)提出了多尺度分层网络设计,实现了性能和效率之间的更好权衡。然而,由于盲点要求,将Transformer应用于自监督图像去噪的努力有限(Wang et al 2023 b; Papkov and Chizhov 2023)。

自监督图像去噪

自监督图像去噪(SSID)试图利用噪声图像本身的信息作为监督(Krull, Buchholz, and Jug 2019; Batson and Royer 2019)。为避免产生过拟合到恒等映射等平凡解,盲点网络(BSN)(Wu et al 2020)通过在每一空间位置将对应的噪声像素从感受野中排除。进一步引入概率推理(Laine et al,2019)和规则损失函数(Huang et al,2021; Zhang et al,2022)来恢复盲点处缺失的信息。在实际RGB图像去噪中,由于图像信号处理(ISP)流水线中的去噪操作,噪声具有空间相关性(Guo et al 2019)。当部署专为空间独立噪声去除而设计的BSN时,它可以轻松拟合输入噪声。一个可行的解决方案是通过像素重排下采样来打破噪声相关性(Zhou et al 2020),然后将BSN应用于下采样图像(Lee,Son和Lee 2022; Wang et al 2023 b; Pan et al 2023; Jang et al 2023 b,a)。此外,CVFSID(Neshatavar et al,2022)学习循环函数以将噪声图像分解成干净和噪声分量。(Li et al 2023)检测平坦区域和纹理区域,然后分别为它们构造监督。(Cheng,Liu,and Tan 2023; Lin et al 2023; Chen et al 2023 a; Zou,Yan,and Fu 2023; Wang et al 2023 a),但对BSN体系结构的研究还不够深入。在这项工作中,我们将Transformer机制应用于BSN,以进一步释放盲点方式的潜力。

Method

网络体系结构概述

在这里插入图片描述

如图2所示,TBSN跟随扩张的BSN(Wu et al 2020),在第一层采用3×3中心掩蔽卷积,在其余层采用膨胀的Transformer注意块(DTAB),网络结构为UNet,采用patch-unshuffle/shuffle(Jang et al,2023a)基于下采样/上采样操作来维持盲点要求。构建块,即DTAB,分别由分组通道自注意(G-CSA)、基于掩膜窗口的自注意(M-WSA)和前馈网络(FFN)的扩张对应物形成。TBSN既受益于通道注意力的全局交互作用,又受益于窗口注意力的局部拟合能力。

分组通道自注意力(G-CSA)

通道注意力(Hu,Shen and Sun 2018)通过显式地对通道之间的相互依赖性进行建模来重新校准通道特征响应。给定输入特征 X ∈ R H × W × C \mathbf{X}\in\mathbb{R}^{H\times W\times C} XRH×W×C,通道注意力可以形式化为,
C A ( X ) = X ∗ ϕ ( X ) . (1) \mathrm{CA}(\mathbf{X})=\mathbf{X}*\phi(\mathbf{X}). \tag{1} CA(X)=Xϕ(X).(1)
其中,函数 ϕ ( ⋅ ) \phi(·) ϕ()聚合每个通道中的空间信息,并且 ϕ ( ⋅ ) \phi(·) ϕ()是通道级乘法运算。例如,NAFNet(Chen et al 2022)通过全局平均池化来实现 ϕ ( ⋅ ) \phi(·) ϕ(),而Restormer(Zamir et al 2022)在通道维度中应用转置矩阵乘法。然而,在SSID任务中,通道注意力可能会泄漏盲点信息,因为在先前的方法中忽略了所有空间位置的内容。

在这项工作中,我们系统地分析了通道注意(CA)在BSN的影响,并实证发现它依赖于通道数量与空间分辨率。对于单层结构(Wang et al 2023 b),空间信息在很大程度上被全局平均池化压缩,因此CA有利于性能。对于多尺度架构(Jang et al 2023 a),空间信息通过下采样操作被混洗到各种通道。因此,CA可以部分地等效于空间交互,从而泄漏盲点值。为此,我们提出控制通道数小于空间分辨率。具体来说,我们引入分组通道自注意(G-CSA)将深度特征分成多个通道组并单独执行CA。我们的G-CSA可以公式化为,
G − C S A ( X ) = C o n c a t ( X 1 ∗ ϕ ( X 1 ) , ⋯   , X G ∗ ϕ ( X G ) ) . (2) \mathrm{G-CSA}(\mathbf{X})=\mathrm{Concat}(\mathbf{X}_1*\phi(\mathbf{X}_1),\cdots,\mathbf{X}_G*\phi(\mathbf{X}_G)). \tag{2} GCSA(X)=Concat(X1ϕ(X1),,XGϕ(XG)).(2)
其中 X = Concat ( X 1 , ⋯   , X G ) , G \mathbf{X}=\text{Concat}(\mathbf{X}_1,\cdots,\mathbf{X}_G),G X=Concat(X1,,XG),G,G是组号。我们将每个组的通道数(即 C G \frac CG GC)设置得足够小,以避免空间信息的泄漏。在实现中,我们将MDTA(Zamir et al 2022)调整到我们的G-CSA中,并使用公式(2)进行全局交互。我们还将3 × 3深度卷积替换为它们的膨胀卷积,以达到盲点要求,如图2(b)所示。

在这里插入图片描述

基于掩模窗口的自注意(M-WSA)

基于窗口的自注意(Liu et al 2021 b)已经广泛应用于图像恢复(Liang et al 2021; Chen et al 2023 b)。在这项工作中,我们模拟了扩张卷积的行为(Wu et al. 2020; Wang et al. 2023b; Jang et al. 2023a)提出了一种用于SSID的基于掩码窗口的自注意(M-WSA),其可以即插即用到任何层中,并利用当前的深度特征作为查询/键/值。如图2(d)所示,我们巧妙地设计了一个固定的注意力掩码,添加到注意力矩阵中,以限制查询和键/值标记之间的交互。从图3(a)可以看出,在原始窗口注意力中,每个查询标记都与窗口内所有空间位置的键/值标记交互。在我们的M-WSA中,查询令牌关注偶数坐标处的空间位置(参见图3(b))。因此,M-WSA表现出与用于构建BSN的扩张卷积相同的功能,但具有更大的感受野和更强的局部拟合能力。

在这里插入图片描述

在这里,我们正式地说明了我们的注意力掩码。在窗口注意力中,在大小为 M × M M\times M M×M的局部窗口内,当前特征首先被分别投影到查询、键和值令牌 Q , K , V ∈ R M 2 × d \mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{M^2\times d} Q,K,VRM2×d。然后,原始窗口注意力可以被公式化为,
Attention ( Q , K , V ) = SoftMax ( Q K T / d ) V . (3) \text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{SoftMax}(\mathbf{Q}\mathbf{K}^T/\sqrt{d})\mathbf{V}. \tag{3} Attention(Q,K,V)=SoftMax(QKT/d )V.(3)
其中 d d d是特征维度。在我们的M-WSA中,我们的注意力掩码 M ∈ R M 2 × M 2 \mathbf{M}\in\mathbb{R}^{M^2\times M^2} MRM2×M2被应用于注意力矩阵,该矩阵限制每个查询仅关注偶数坐标处的键/值,如图2(d)所示。因此,公式(3)可以修改为,
A t t e n t i o n ( Q , K , V ) = S o f t M a x ( Q K T / d + M ) V , M ( i , j ) = { 0 , if  x i − x j ≡ y i − y j ≡ 0 ( m o d 2 ) − ∞ , otherwise . \begin{align} \mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{SoftMax}(\mathbf{QK}^T/\sqrt{d}+\mathbf{M})\mathbf{V}, \tag{4}\\ \left.\mathbf{M}(i,j)=\left\{\begin{array}{ll}0,&\text{if }x_i-x_j\equiv y_i-y_j\equiv0\pmod{2}\\-\infty,&\text{otherwise}\end{array}\right.\right.. \tag{5} \end{align} Attention(Q,K,V)=SoftMax(QKT/d +M)V,M(i,j)={0,,if xixjyiyj0(mod2)otherwise.(4)(5)
具体来说, M M M是一个二值矩阵,它根据查询(在 i i i处)和键/值(在 j j j处)标记的相对位置屏蔽了某些位置。 ( x i , y i ) (x_i,y_i) (xi,yi) ( x j , y j ) (x_j,y_j) (xj,yj) i i i j j j的空间位置。当 i i i j j j在两个轴上的距离相等时, M ( i , j ) = 0 \mathbf{M}(i,j)=0 M(i,j)=0,注意力值不变。受相对位置嵌入(Liu et al 2021 b)的启发,可以根据 i i i j j j的相对位置从较小的二进制矩阵 M ^ ∈ R ( 2 M − 1 ) × ( 2 M − 1 ) \hat{\mathbf{M}}\in\mathbb{R}^{(2M-1)\times(2M-1)} M^R(2M1)×(2M1)计算 M ∈ R M 2 × M 2 \mathbf{M}\in\mathbb{R}^{M^2\times M^2} MRM2×M2,以提高效率,即,
M ( i , j ) = { 0 , if  M ^ ( x i − x j , y i − y j ) = 0 − ∞ , if  M ^ ( x i − x j , y i − y j ) = 1 , M ^ ( x , y ) = { 0 , if   x ≡ y ≡ 0   ( m o d   2 ) 1 , otherwise . \begin{align} \mathbf{M}(i,j)=\left\{\begin{array}{ll}0,&\text{if }\hat{\mathbf{M}}(x_i-x_j,y_i-y_j)=0\\-\infty,&\text{if }\hat{\mathbf{M}}(x_i-x_j,y_i-y_j)=1\end{array}\right., \tag{6}\\ \hat{\mathbf{M}}(x,y)=\left\{\begin{array}{ll}0,&\text{if}\:x\equiv y\equiv0\:(mod\:2)\\1,&\text{otherwise}\end{array}\right.. \tag{7} \end{align} M(i,j)={0,,if M^(xixj,yiyj)=0if M^(xixj,yiyj)=1,M^(x,y)={0,1,ifxy0(mod2)otherwise.(6)(7)
在实现方面,我们采用重叠交叉注意(Chen et al 2023b),从更大的领域计算键/值标记,以进一步扩大感受野。

讨论。建议的G-CSA和M-WSA与以前BSN中的Transformer算子不同。如图4(a)(b)所示,LGBPN中的信道注意力(Wang et al 2023 b)在应用于多尺度架构时有泄漏盲点信息的风险,而我们的GCSA在单独的组中执行通道注意以缓解此问题。在SwinIA窗口注意(Papkov和Chizhov 2023)掩盖了注意力矩阵的主对角线以维持盲点要求。其键/值标记仅限于来自噪声输入的像素级浅层特征,因此显示出较差的结果。相比之下,我们的M-WSA应用专门设计的掩模来模拟膨胀卷积的行为,这可以在深特征上灵活地执行。

在这里插入图片描述

知识蒸馏实现高效推理

自监督图像去噪方法通常计算成本较高,这源于复杂的网络设计(Wu et al 2020)、增大的网络规模(Jang et al. 2021)以及后处理操作(Lee, Son, and Lee 2022)。计算负担在很大程度上限制了它们在某些情况下的适用性,例如在移动的设备上。尽管如此,SSID方法的性能仍然不如于相应的监督方法。即使是轻量级监督方法也可能比复杂的自监督方法获得更好的性能。换句话说,轻量级网络可能完全足以适应某些复杂的自监督方法的结果。利用这一点,我们提出了一种知识蒸馏策略,以减少推理成本,同时保持性能。

具体来说,我们采用高效的U-Net(Ronneberger,Fischer and Brox 2015)架构作为我们的学生网络,它是从自监督学习的TBSN(即TBSN2UNet)中提取出来的,
L d i s t i l l = ∥ s g ( T B S N ( y ) ) − U N e t ( y ) ∥ 1 (8) \mathcal{L}_{distill}=\left\|sg(\mathbf{TBSN}(y))-\mathbf{UNet}(y)\right\|_{1} \tag{8} Ldistill=sg(TBSN(y))UNet(y)1(8)
其中 y y y是噪声图像, s g ( ⋅ ) sg(·) sg()是停止梯度操作。请注意,我们的目标是减少推理过程中的计算成本。这与应用知识蒸馏以获得更好性能的方法不同(Wu et al 2020; Jang et al 2021; Li et al 2023)

Experiments

实现细节

数据集。我们对两个广泛使用的真实世界图像去噪数据集进行了实验,即SIDD(Abdelhamed,Lin and Brown 2018)和DND(Plotz and Roth 2017)。SIDD数据集的噪声-干净对是从五个智能手机摄像头中收集的,其中每个噪声图像被多次捕获,平均图像作为ground truth。DND是从数码单反相机中采集的一个基准数据集,噪声图像是用短曝光时间拍摄的,而相应的干净图像是用长曝光时间拍摄的。它包含50对仅供测试。我们以完全自监督的方式在测试图像上训练和测试我们的网络。

训练细节。对于TBSN的自监督训练,我们遵循AP-BSN(Lee,Son and Lee 2022)应用像素洗牌下采样(PD)来打破噪声相关性,并在训练和推理过程中采用非对称PD因子来权衡去噪效果和细节保留。我们还采用随机替换细化(R3)策略对去噪效果进行了改进,批量大小和块大小分别设置为4和128×128(Loshchilov and Hutter 2018)优化器来训练网络。学习率最初设置为3×10−4,每40k次迭代减少10次,总共100k次训练迭代。对于知识提取,训练设置与自监督学习相同。所有实验都在PyTorch框架和Nvidia RTX 2080TiGPU上进行。

与最新方法的比较

定量比较。表1显示了拟议的TBSN和最先进的自监督方法的定量结果:Noise2Void(Krull,Buchholz, and Jug 2019),Noise 2Self(Buchholz, and Jug 2019)、NAC(Xu et al 2020)、R2R(Pang et al 2021)、CVF-SID(Neshatavar et al 2022)、AP-BSN(Lee、Son and Lee 2022)、SASL(Li et al 2023)、LG-BPN(Wang et al 2023b)、PUCA(Jang et al 2023a)和AT-BSN(Chen et al 2024)。其中,针对空间独立噪声设计的盲点技术(Noise2Void、Noise2Self和R2R)对真实世界的噪声图像表现出很小的去噪效果。虽然像素重排下采样(PD)打破了噪声相关性并成功地去除了噪声(Lee,Son and Lee 2022),但性能仍然受到其普通卷积BSN架构的限制。最近的一些工作通过寻找先进的BSN结构来解决这个问题。例如,LG-BPN将Transformer块(Zamir et al 2022)并入BSN以获得全局信息,并且比AP-BSN基线提高了0.37dB。PUCA设计了具有信道关注的多尺度BSN,并且实现了0.63dB的提高。尽管如此,得益于通道和窗口注意机制,我们的TBSN在SIDD基准数据集上将改进提高到0.87dB。AT-BSN引入了多教师蒸馏策略,相比之下,我们的TBSN2UNet在SIDD和DND基准数据集上分别比AT-BSN提高了0.01dB和0.33dB,显示了我们方法的优越性。

在这里插入图片描述

定性比较。自监督图像去噪方法的定性结果如图5和图6所示。图5中颜色图表的去噪取决于全局信息,其中以前的方法不能完全消除噪声。利用通道自注意,我们的TBSN平滑地去除了空间相关噪声。

在这里插入图片描述

图5和图6中的杯和壁表明,由于其窗口注意的局部拟合能力,TBSN可以保持细节。

在这里插入图片描述

模型复杂度比较

基于通道(Zamir et al 2022)和基于窗口的(Liu et al 2021 b)TBSN中的自注意力是为图像恢复任务设计的高效Transformer模块。此外,TBSN采用分层多尺度架构以进一步提高其效率。如表2所示,TBSN保持与卷积对应物PUCA相似的计算复杂度,并且比LG-BPN更高效。此外,由于SASL的U-Net结构,SASL显示了有吸引力的#Param和#FLOPs结果。为了公平的比较,从我们的TBSN提取的U-Net表现出与SASL相同的计算复杂度,但具有更高的性能,这表明我们的知识提取策略的优越性。

在这里插入图片描述

Ablation Study

感受野的可视化

感受野的扩大是Transformer·成功的一个主要因素(Zamir et al 2022)。在本小节中,我们绘制了用于恢复输出中心像素的输入像素,以获得TBSN的有效感受野。具体来说,我们通过网络传递一个输入图像,选择输出图像的中心像素,并计算其相对于输入图像的梯度。PyTorch格式的伪代码如下:

import torch as t 
image = t.rand((1,3,H,W),requires_grad=True) 
output = model(image) 
center_pixel = t.mean(output[...,H//2,W//2]) 
center_pixel.backward() 
gradient = t.sum(t.abs(x.grad),dim=1,keepdim =True)

这种梯度表示输出像素随每个输入像素上的干扰而变化的程度。我们将梯度的绝对值沿着通道轴求和以进行可视化。从图1中可以看出,TBSN比以前的BSN显示出更宽的感受野(Lee,Son and Lee 2022; Wang et al 2023 b; Jang et al 2023 a),这可能是TBSN吸引人的性能的一个解释。

在这里插入图片描述

DTAB分析

表3分析了我们的膨胀Transformer注意力块(DTAB)中组件的有效性。我们开始于一个基本模型(1),它将窗口和通道注意力退化为膨胀卷积(Wu et al 2020)。相比之下,我们的扩展M-WSA(2)增强了具有局部拟合能力的基础模型,并提供了0.27dB的改进。我们的扩展G-CSA(3)表现出全局交互作用,表现出0.65dB的改善。结合通道和窗口注意力,TBSN实现了0.81的总改善,这证明了局部和全局操作的互补性。此外,我们在图4中评估了其他窗口和通道注意实现的效果。用SwinIA(4)代替M-WSA导致0.59dB性能下降,而用LG代替G-CSA,BPN(6)导致0.05 dB性能下降。其他注意力机制,Swin Transformer(Liu et al 2021b)(5),SE(Hu,Shen,and Sun 2018)(7)和SCA(Chen et al 2022)(8)也显示出较差的性能。总之,消融研究结果表明,DTAB是最佳网络选择。

在这里插入图片描述

通道注意力分析

如该方法所示,通道自注意(CSA)(Wang et al 2023 b)可能会在多尺度架构中泄漏盲点信息。表4分析了CSA的效果,下采样尺度从1到5。相应地,最深层的通道数量从48增加到768,空间分辨率从1282降低到8 2,从表4的中间行可以看出,普通的扩展CSA在小于等于3的尺度上提供了积极的效果,但是在4-和5-尺度上导致了明显的性能下降。这是由于在4-和5-尺度上信道维度大于空间分辨率,所以它泄漏了盲点信息。相反,我们的分组扩展G-CSA将通道分成几组,并分别执行通道注意。由于每组内的通道维度被控制为小于空间分辨率,因此扩展G-CSA在所有尺度上都提供恒定的改进。

在这里插入图片描述

知识蒸馏分析

如表5所示,我们使用U-Net架构进行实验,以评估知识蒸馏的有效性。尽管参数和FLOP较少,但在SIDD验证数据集上,以监督方式训练的U-Net的性能比TBSN高1.21dB。这表明轻量级U-Net具有足够的学习能力来获得TBSN的去噪性能。因此,实验结果表明,知识蒸馏是一种有效的方法,可以有效地减少学生U-Net的模型规模和推理过程中的计算量。

在这里插入图片描述

Conclusion

为此,本文提出了一种基于变换的自监督图像去噪盲点网络(TBSN)。提出了一种新的自适应空间和通道自注意算子的BSN构造方法.对于空间自注意,我们对窗口自注意矩阵应用了一个精心设计的掩码,从而限制了它的感受域,使之模仿膨胀卷积。对于通道自注意的空间信息泄漏问题,我们提出了在不同的组中执行通道自注意,以消除其有害影响.实验结果表明,TBSN算法在去噪数据集上扩展了有效的感受野,取得了较好的性能.

Appendix

感受野可视化

在这里插入图片描述

如图A所示,与卷积算子相比,我们的扩展M-WSA执行更长范围的交互,而我们的扩展G-CSA聚合了全局感知的所有空间位置。图B显示了TBSN中各个注意层的接收场。从图中可以看出,我们的GCSA和M-WSA可以模仿以前BSN中的膨胀卷积(Wu et al 2020)以保持盲点要求。

在这里插入图片描述

合成去噪结果

在本节中,我们验证了TBSN在合成去噪方面的有效性。具体来说,我们用负对数似然损失训练TBSN,并用后验推理测试它们。(Laine et al 2019)。训练图像被裁剪成具有翻转/旋转增强的patch。批量大小设置为4,patch大小设置为128×128。网络使用Adam进行训练(Loshchilov and Hutter 2018)优化器,总共40万次迭代。学习率最初设置为1×10−4,并通过余弦退火调度器降至零(Loshchilov and Hutter 2016)我们观察到通道注意力更容易泄漏合成噪声上的结合点信息,因此我们将每个注意力组中的通道数量设置为4。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

合成图像去噪的定量结果如表A所示。从表中可以看出,由于盲点处的信息丢失,Noise2Void(Krull, Buchholz, and Jug 2019)无法充分恢复干净的信号。Laine19(Laine et al 2019)通过贝叶斯推理缓解了这个问题,贝叶斯推理有可能完全恢复干净的信息。此外,R2R(Quan et al 2020)从单个噪声图像中合成训练对,用于网络训练。(Huang et al 2021)和Blind2Unblind(Wang et al 2022 b)应用规则损失函数来重用盲点信息。然而,所有上述方法都基于普通卷积架构,受益于Transformer算子,TBSN在高斯噪声σ=25的Kodak数据集上比Laine 19(Laine et al 2019)提高了0.16 dB,这证明了TBSN的优越性。

由于自监督训练的TBSN的性能已经与最先进的监督基线相匹配,因此轻量级U-Net(Ronneberger,Fischer和Brox 2015)只能提高推理效率,而不能实现可比的性能。(Ronneberger,Fischer,and Brox 2015),即TBSN2UNet显示出比TBSN稍差的结果。然而,从主论文的表5来看,U-Net比TBSN效率更高。这表明我们的知识蒸馏策略可以有效地降低推理过程中的计算成本。

模型大小影响

在本节中,我们研究了TBSN模型大小对性能的影响。从表B中可以看出,性能首先随着模型大小的增加而增加,然后趋于稳定。性能最初受到模型容量的限制(表B(1)(2)(3)),较大的模型有利于自监督学习。在某个点上,性能反而受到学习框架、训练样本数量等的限制,进一步扩大模型大小并不能提供改进(表B(5)(6))。因此,我们选择表B(4)作为TBSN的型号配置。

在这里插入图片描述

其他定性结果

图C和图D分别提供了SIDD(Abdelhamed,Lin and Brown 2018)和DND(Plotz and Roth 2017)数据集的额外视觉比较。

在这里插入图片描述
在这里插入图片描述

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值