学习稀疏 Transformer 网络以实现高效图像去雨

基于 Transformer 的方法因能建模对高质量图像重建至关重要的非局部信息，在图像去雨上取得了显著性能。本文中，我们发现大多数现有 Transformer 通常利用查询-键对中令牌的所有相似度进行特征聚合。然而，若查询令牌与键令牌不同，由这些令牌估计的自注意力值也会参与特征聚合，从而干扰清晰图像复原。为克服该问题，我们提出一种高效的去雨网络——稀疏 Transformer（DRSformer），其能自适应保留最有用的自注意力值用于特征聚合，使聚合特征更好地促进高质量图像重建。具体而言，我们设计可学习的 top-k 选择算子，为每个查询自适应保留来自键的最关键注意力分数以优化特征聚合。同时，由于 Transformer 中朴素前馈网络未建模对潜在清晰图像复原重要的多尺度信息，我们开发了有效的混合尺度前馈网络以生成更优去雨特征。为学习结合 CNN 算子局部上下文的丰富混合特征集，我们为模型配备专家混合特征补偿器，提出一种协作精化去雨方案。在常用基准上的大量实验结果表明，所提方法相较最先进方法取得了优越性能。源代码与训练模型见 https://github.com/cschenxiang/DRSformer。