图像恢复论文简记——Uformer: A General U-Shaped Transformer for Image Restoration

2021年6月6日提交到arxiv上的文章。后来的医学图像去噪架构Eformer就是Uformer基础上改进而来，看来还是值得一读，简单记录下。

知乎同名账号同步发布。

一、架构设计

架构如图：

比起普通的UNet，不同之处在于采用了LeWin Transformer，这种Transformer也是本工作的创新点。

所谓LeWin Transformer，就是local-enhanced window Transformer，其中包含W-MSA和LeFF：

• W-MSA：non-overlapping window-based self-attention，作用是减小计算开销（传统transformer是在全局计算self-attention，而它不是）；
• LeFF：传统transformer中采用前馈神经网络，不能很好利用local context，LeFF的采用可以capture local information。

⚠️两个创新点：

• 提出LeWin Transformer，引入UNet
• 三种跳跃连接

二、主要模块细节

2.1，W-MSA

这是本工作最大的创新点。 （经提醒，swin Transformer里就有）

首先将将C×H×W的X分为N个C×M×M个patch，每个patch视为有M×M个C维vector（N = H × W / M²），这C个vector就输入W-MSA中。根据上述公式，简单理解就是将X分为不重叠的N片，然后对每一片进行self-attention的计算即可。

作者表示，虽说是在一片上进行self-attention的计算，但是在UNet的encode阶段，由于下采样的存在，所以在这一片上计算自注意力，对应在下采样前更大感受野上计算自注意力。

采用了relative position encoding，所以计算公式可以表示为：

这种位置编码的引用[48,41]分别是：

[48] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repre-
sentations. arXiv preprint arXiv:1803.02155, 2018.
[41] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining
Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint
arXiv:2103.14030, 2021.

2.2，LeFF

LeFF是Incorporating Convolution Designs into Visual Transformers发明的，其中的Convolution-enhanced image Transformer (CeiT)包含了这个设计。

精髓在于对self-attention计算模块输出的N个token（vector），重新排列为$\sqrt{N}\times \sqrt{N}$的“image”，然后进行depth-wise的卷积操作。看完CeiT作者给出的图解，再看Uformer作者给出的图解，就不难理解含义了：

每个线性层/卷积层之后，用的都是GELU激活函数。
（depth-wise的卷积网上一搜就有，作用是减少参数，提升计算速度）

2.3，三种跳跃连接

UNet架构是有跳跃连接的，在本工作中就是将encoder部分的Transformer的输出传递到decoder部分，但是利用这些跳跃连接传递的信息的方式有很多种，作者探索了三种：

• 第一种是直接concat过来；
• 第二种是：每一个decode stage有一个上采样和两个Transformer block，意思是第一个用self-attention，第二个用cross attention；
• 第三种是concat的信息作为key和value的cross attention。

作者认为三种差不多，但第一种好一点点，所以采用第一种作为Uformer的默认设置。

Uformer的架构设计细节就在这，其他内容不细看了。

三、计算开销

既然LeWin Transformer中的W-MSA主打减小计算开销，那么自然要看下算法复杂度：

给定feature map X，维度为C×H×W，如果是传统的self-attention，那么复杂度为$O(H^2{W}^{2}C)$，分割为M×M的patch再做self-attention，则是$O(\frac{HW}{M^2}{M}^{4}C)=O(M^{2}HWC)$，复杂度有所减少。

四、实验结果

作者做了去噪、去雨、去模糊的实验。