摘要：近日，一篇发表于计算机视觉顶级期刊 TPAMI 的文章提出了一项为高效图像复原任务量身打造的深度学习模型，自适应稀疏 Transformer(AST-v2)模型。针对现有方法在性能与效率难以兼顾的问题，AST-v2 通过降低特征冗余、抑制无关区域的噪声交互，为
作者丨周世豪、潘金山、杨巨峰
近日，一篇发表于计算机视觉顶级期刊 TPAMI 的文章提出了一项为高效图像复原任务量身打造的深度学习模型，自适应稀疏 Transformer(AST-v2)模型。针对现有方法在性能与效率难以兼顾的问题，AST-v2 通过降低特征冗余、抑制无关区域的噪声交互，为图像去雨、去雾、去模糊等多类复原任务提供了新的解决思路。模型的核心设计包括 自适应稀疏自注意力 (ASSA) 与 特征细化前馈网络 (FRFN) ，二者协同作用，使得 AST-v2 在六个经典图像复原基准任务上均取得了领先性能。
论文标题： Learning An Adaptive Sparse Transformer for Efficient Image Restoration
作 者： Shihao Zhou, Jinshan Pan, and Jufeng Yang
作者机构：南开大学、鹏城实验室，南京理工大学
论文地址：https://ieeexplore.ieee.org/abstract/document/11106710
项目主页：https://github.com/joshyZhou/ASTv2
Hugging face：https://huggingface.co/spaces/yssszzzzzzzzy/AStv2
在图像复原领域，现有研究面临一个棘手的“两难困境”：一方面，我们希望模型能拥有像 Transformer 的强大全局建模能力，能看到图像的每一个角落，捕捉长距离依赖关系；但另一方面，这种“一览无余”的密集注意力（Dense Attention）机制，却不可避免地将无关区域的噪声交互和冗余的特征信息一并带入计算，干扰了模型对清晰图像的重建能力。
尽管社区已提出多种高效注意力机制尝试解，但仍面临两大核心挑战：
噪声交互： 标准的密集注意力机制会不可避免地引入来自无关区域的噪声交互，干扰有效特征的学习。信息冗余： 密集聚合后的特征图中往往包含大量冗余信息，这会阻碍模型关注真正关键的特征。现有解决方案，如 Top-K 选择算子或超像素空间建模，或存在对特定任务参数敏感的问题，或未能从根本上解决冗余噪声交互问题。而简单地应用稀疏注意力又容易引发信息丢失。因此，设计一种能够智能识别最有价值的特征，同时又能平衡稀疏性与信息完整性的高效机制，成为了本研究的核心动机。
为了解决上述问题，论文提出了一个新的解决范式：AST-v2。AST-v2 的整体框架基于一个流行的编码器-解码器结构。
图1:本文提出的自适应稀疏框架。它主要由自适应稀疏自注意力（ASSA）模块和特征细化前馈网络 （FRFN）构成
图1:本文提出的自适应稀疏 Transformer (AST-v2) 框架。它主要由自适应稀疏自注意力（ASSA）模块和特征细化前馈网络 （FRFN）构成
如图1所示，AST-v2的总体架构主要有两个核心模块组成：自适应稀疏自注意力(ASSA)和特征细化前馈网络(FRFN)。
如图2所示，该模块采用了一种精巧的双分支范式，旨在动态地平衡信息密度与稀疏性。
密集分支(DSA): 采用标准 Softmax激活函数的注意力机制，以保留完整的上下文信息。稀疏分支(SSA): 采用平方 ReLU 激活函数滤除低相关性的查询-键（Query-Key）对，生成稀疏的注意力分数。自适应调制:利用稀疏分支生成的注意力图，通过一个门控机制和线性投影，产生尺度(γ)和位移(β)因子用于对密集分支的注意力权重进行像素级的特征变换(γ⊙DSA +β)。这种设计使得模型能够根据输入特征自适应地调整注意力的稀疏程度，既能滤除噪声，又避免了因过度稀疏导致的关键信息丢失问题。图2: 自适应稀疏自注意力（ASSA）机制示意图
为了解决通道维度上的特征冗余问题，论文设计了如图3所示的一种遵循“增强-消除”(enhance-and-ease) 的处理策略的FRFN 模块。
增强: 通过分组的部分深度可分离卷积，模型可以专注于增强特征图中信息量最丰富的通道子集，而非对所有通道进行无差别处理。消除: 在特征增强后，利用一个门控机制来抑制或“简化”那些信息量较少的冗余特征。通过这一机制，FRFN 能够学习到更具判别性的特征表示，从而提升最终的图像重建质量。
图3: 特征细化前馈网络（FRFN）示意图
AST-v2在六种不同的图像复原任务多个公开基准数据集（如SPAD,Snow100K,GoPro,SOTS等）上进行了广泛的实验评估。结果显示，该模型在各项任务中的关键指标（PSNR,SSIM）上均达到或超越了此前的 SOTA 方法。
AST-v2 模型在图像去雨任务上展现了卓越的性能，通过在SPAD数据集上的定量指标（表1）和定性视觉效果（图4）得到了双重验证。定量分析： AST-v2 取得49.79 dB 的PSNR最高分，超越了包括 Restormer(46.25)、DRSformer(48.53)和其会议版本 AST(49.51)在内的所有模型，这表明其复原结果的优越性。模型的SSIM分数为 0.9939，证明其生成的图像在亮度、对比度和结构上都与参考图像高度一致。表1: 在 SPAD数据集上针对雨痕去除任务的定量比较
定性分析：在第一个栅栏的例子和第二个车轮标志的例子中，其他方法或多或少地错误的将图像内容理解为雨痕残影，而 AST-v2(i)成功地分辨了这种干扰，复原出的标志清晰、自然，最接近参考图像。图4: 数据集上雨痕去除的定性结果。其他方法都残留了雨痕，而 AST-v2 生成的结果不仅干净，也更接近参考图像
图4: SPAD 数据集上雨痕去除任务的定性结果。相比之下，其他方法都残留了雨痕，而 AST-v2 生成的结果不仅干净，也更接近参考图像
在Snow100K数据集上的测试结果表明，AST-v2在图像去雪任务中展现了性能优势，无论是在客观的评价指标（表2）还是在主观的视觉效果（图5）上，都超越了其他对比方法。定量分析： AST-v2模型的PSNR分数达到了34.18 dB，是所有参与比较的十五种方法中最高的，这证明其生成的图像在像素级别上最接近原始无雪图像。SSIM分数为 0.94，与其它先进方法（如 NAFNet、SFNet 等）的0.95分数非常接近，表明其复原结果在结构和纹理上保持了高度的保真度。定性分析：在图中展示的灯塔、马拉松运动员和路牌三个场景中，AST-v2(f)都能非常有效地去除图像中的雪花干扰，其效果明显优于JSTASR(c)和HDCW(d)等方法。更重要的是， AST-v2在去除雪花的同时，没有牺牲图像原有的精细纹理。例如，在灯塔的特写中，墙体的纹理和颜色得到了很好的保留；在运动员的特写中，面部细节也比其他方法更加清晰。图5: 在数据集上进行去雪的定性结果。相比之下，AST-v2 在有效去除雪花的同时，没有牺牲细节纹理
图5: 在 Snow100K数据集上进行去雪任务的定性结果。与所对比的其他技术相比，AST-v2 在有效去除雪花的同时，没有牺牲细节纹理
在 ISTD 数据集上进行的图像去阴影实验表明，AST-v2 模型在处理该项复杂任务时，无论是在综合量化指标（表3）还是在视觉效果的自然度上（图6），都表现出卓越的性能。定量分析：在重要的全图(All)评估中，AST-v2 在 PSNR(30.26)和 SSIM(0.959)两项指标上均取得了最高分。各区域表现优异: 在非阴影区域，AST-v2 的 SSIM分数(0.973)与DMTN并列第一；在阴影区域的PSNR(36.41)和非阴影区域的PSNR(32.19)也均位列第二，展现了其均衡且强大的恢复能力。值得注意的是，表中部分方法（标注了 † 符号）利用了额外的阴影掩码（shadow mask）信息来辅助处理。而 AST-v2 在不使用这些额外信息的情况下，依然取得了最佳的综合性能，更凸显了其算法的先进性。表3: 在 ISTD数据集上针对图像去阴影任务的定量比较。† 表示该方法利用了额外的阴影掩码信息
定性分析：对比输入图像(b)，AST-v2(f)能够非常彻底地移除地砖和路面上的阴影。相较之下，其他一些方法存在明显复原失败的区域。例如，DC(c)和 AST(e)未能完全去除阴影，留下了模糊的暗色块；而LG(d)则在去除阴影后留下了清晰的“残影”或“鬼影”伪影。AST-v2 生成的图像不仅没有阴影，而且很好地恢复了地面原有的纹理，效果干净自然，最接近参考图像(a)。图6: 在 ISTD数据集上进行阴影去除任务的定性结果。相比之下，AST-v2 生成的无阴影图像不会引入明显的伪影
AST-v2(L)模型在图像去模糊任务上展现了顶尖的性能和卓越的泛化能力。通过在合成数据集(GoPro)和真实世界数据集(RealBlur)上的综合测试，该模型在定量指标（表4）、视觉效果（图7与图8）和模型效率三方面均表现出色。定量分析：双数据集 SOTA 性能:在GoPro(合成模糊)数据集上，AST-v2(L)的 PSNR 达到 33.12，SSIM 达到 0.962，均为所有对比方法中的最高分。在RealBlur(真实模糊)数据集上，AST-v2(L)同样取得了PSNR(32.62)和SSIM(0.921)的双料冠军。所有模型都仅在GoPro数据集上进行训练。AST-v2(L)在 RealBlur真实场景测试中依然保持高性能，这证明了其强大的泛化能力，能将从合成数据中学到的知识有效应用于真实世界的复杂模糊模式。表4: 在 GoPro和 RealBlur基准数据集上针对图像去模糊任务的定量比较。所有方法仅在 GoPro 数据集上训练
定性分析：在GoPro数据集的车牌修复案例中，多种方法处理后的车牌号码依然存在模糊或扭曲的伪影。相比之下，AST-v2(L)生成的结果最为清晰，还原了参考图像(b)的细节。在RealBlur数据集的真实场景文字修复案例中，输入图像(b)和其他方法的结果(c,d,e)中的韩语文字都非常模糊，难以辨认。AST-v2(L)的结果(f)则显著提升了文字的清晰度和可读性，效果最接近清晰的参考图像(a)，再次验证了其在真实世界场景中的有效性。图7: 在 GoPro数据集上针对合成运动模糊去除任务的定性结果。AST-v2 (L) 生成的结果更清晰，模糊伪影也更少
图7: 在 GoPro数据集上针对合成运动模糊去除任务的定性结果。与其他所考虑的方法相比，AST-v2 (L) 生成的结果更清晰，模糊伪影也更少
图8: 在 RealBlur数据集上进行真实模糊伪影去除的定性结果。AST-v2 (L) 恢复的结果中，字符更为清晰
AST-v2模型在暗光图像增强任务上展现了优越的性能，这在两个广泛使用的基准数据集LOL-v2和SMID上的定量（表5与表6）与定性评估（图9与图10）中均得到了充分证明。定量分析：在LOL-v2数据集上，该数据集包含真实(real)和合成(syn)两个子集，AST-v2 在综合评估中表现最佳。它取得了24.40 的最高平均PSNR。在真实世界的暗光基准SMID上，AST-v2的优势更为明显，在PSNR(29.51)和SSIM(0.824)两项指标上均取得了第一名。其性能超越了包括 MambaLLIE(PSNR 29.26)和Retinexformer(PSNR 29.15)在内的所有先进方法。表5: 在 LOL-v2数据集上针对暗光增强任务的定量比较
表6: 在 SMID数据集上针对暗光增强任务的定量比较
定性分析：在LOL-v2数据集上:无论是合成场景（上排）还是真实场景（下排），AST-v2 增强后的结果在色彩和纹理上都最接近参考图像。相较于其他方法，AST-v2 的效果更为自然和均衡。在SMID数据集上，其他对比方法在增强后都出现了不同程度的色彩偏移，而AST-v2生成的图像颜色最接近真实值。图9: 在颜色和纹理方面，经 AST-v2 增强后的结果在视觉上更接近参考图像
图9: 在 LOL-v2 数据集上进行暗光增强的定性结果。两个例子分别来自合成子集（上）和真实子集（下）。在颜色和纹理方面，经 AST-v2 增强后的结果在视觉上更接近参考图像
图10: 在 SMID 数据集上进行暗光增强的定性结果。与所对比的方法相比，我们方法生成的图像颜色更接近参考图像
在SOTS基准数据集上，AST-v2(L)模型展现了其在图像去雾任务上的高性能，其优势在定量指标（表7）和定性视觉效果（图11）上都得到了验证。定量分析：在SOTS数据集上，AST-v2(L)的PSNR分数达到了32.28，位列第一。同时，其SSIM分数也达到了0.980，与NDR-Restore方法并列第一。表7: 在 SOTS数据集上针对去雾任务的定量比较
定性分析：在湖中大黄鸭的场景中，输入图像(b)覆盖着浓厚的雾气。与其他方法如 PromptIR(c)和AST(d)相比，AST-v2(L)(e)恢复出的结果更为清晰。通过观察图中放大的桥梁细节可以看出，AST-v2(L)不仅有效地去除了雾气，还避免了其他方法可能引入的色彩失真或模糊残留问题，其视觉效果最接近无雾的参考图像(a)。图11: 在 SOTS 基准上进行去雾任务的定性结果。与其他被考虑的方法相比，AST-v2 (L) 恢复的结果更为清晰
图表8所示，AST-v2在实现卓越性能的同时，展现了极高的模型效率。在与MPRNet、SwinIR、Restormer 等七种主流SOTA模型的效率分析中，AST-v2展现出强大的竞争力。在所有参与对比的模型中，AST-v2的PSNR分数最高，达到了49.79。AST-v2 的参数量仅为 8.7M，是所有模型中最轻量的。同时，AST-v2的运行时间为0.05秒，在所有模型中位列第二，展现了极快的推理速度。此外，如表9所示，与该工作的前期工作(AST)相比，AST-v2在去雨任务上PSNR提升了0.28dB，同时参数量减少了87%，FLOPs 减少了18%，推理时间缩短了49%。
表8: 在 SPAD数据集上的模型效率分析
表9: 与 AST 相比，AST-v2 更准确，同时模型也显著更轻量、更快速。
表9: 在 SPAD 基准上的图像去雨任务，于相同实验设置下对 AST-v2 和 AST [21] 进行的比较。FLOPs 和推理时间是在256×256 尺寸的图像上计算的。与 AST 相比，AST-v2 更准确，同时模型也显著更轻量、更快速。值得注意的是，模型复杂度上的改进在各种图像修复任务中是一致的
经过 AST-v2 修复后的图像不仅在视觉质量上得到提升，更能显著改善如目标检测和语义分割等高级视觉任务的性能（图12），让之前因图像降质而失败的检测/分割样本被成功识别。
图12: 由 AST-v2 修复的图像能够提升下游任务的性能，其中，原先的失败样本可以被成功地检测或分割
图12: 来自退化图像和修复后图像的高级计算机视觉任务结果。由 AST-v2 修复的图像能够提升下游任务（即目标检测和语义分割）的性能，其中，原先的失败样本（例如，修复后图像中的手提箱和公交车门）可以被成功地检测或分割
深入剖析了现有 Transformer 图像复原方法中，密集注意力的“噪声与冗余”和稀疏注意力的“信息丢失”之间的根本矛盾。AST-v2 并非在两者间进行取舍，而是提出了一种全新的解决范式。
在去雨、去雪、去模糊等六大极具挑战性的图像复原任务中均取得了当前最佳（SOTA）性能。更重要的是，它在实现性能超越的同时，大幅降低了模型复杂度（参数量锐减87%），实现了性能、效率和速度的完美平衡。
作者团队将代码开源(https://github.com/joshyZhou/ASTv2)，极大地便利了社区进行后续的研究和应用。
总的来说，AST-v2通过重新审视底层视觉模型中注意力的设计范式，为如何构建兼具高性能与高效率的图像复原模型提供了一个全新的、极具潜力的方向。这项工作不仅对学术界在 Transformer 架构上的探索有重要的启发意义，也为工业界在资源受限的真实场景下部署和应用先进的图像复原技术，提供了实践指导。
来源：极市平台
