ELAN论文详解：对SwinIR进行效率优化的超分算法（Efficient Long-Range Attention Network for Image Super-resolution）

Paper: Efficient Long-Range Attention Network for
Image Super-resolution

Code: https://github.com/xindongzhang/ELAN

作者团队：香港理工大学、OPPO

1. 概述

作者认为当前基于Transformer的超分模型在效率和效果上都存在问题，首先在效果上，swinir使用8x8的窗口，限制了建立长距离依赖关系的能力；在效率上，SwinIR基本是将swin transformer照搬到low level任务中，对于超分任务来说有很多多余和碎片化的组件，比如相对位置编码（relative position bias）, attention中的掩码机制(masking mechanism), layer normalization, and several sub-branches created with residual shortcut.

本篇论文的主要贡献有两点：

1. 提出了efficient long-range attention block（ELAB），可以更好的建模图像的长距离依赖；
2. 提出了一系列的有效的优化策略，可以提高模型的效率。

2. 算法解析

2.1. Overview

算法的整体结构如下图所示，延续了SwinIR的“浅层特征提取+深层特征提取+图像重建”的结构，区别就是在深层特征提取模块，使用了新提出的ELAB替代RSTB。

2.2. Efficient Long-Range Attention Block(ELAB)

ELAB的结构如下图所示，由一个局部特征提取模块和一个group-wise multi-scale attention (GMSA)组成。

• Local feature extraction
  

通过shift conv + Relu + shift conv实现，shift conv如图所示，首先将输入特征图在通道上分成5份，分别左移、右移、上移、下移、不动，然后进行1x1卷积操作，相当于可以获取相邻像素的信息，在保持1x1卷积的计算量的同时，增加了感受野。

• Group-wise multi-scale self-attention (GMSA)

在M-MSA中越大的window尺寸可以探索更多的信息，但是计算量也会平方增加，因此作者提出了GMSA，

如图所示，GMSA是一个分组的多尺度的自注意力机制，通过将输入特征图在通道维度上分成K个group，每个group内部进行W-MSA，而且每个组中使用的window的大小是不同的，这就是多尺度，完成注意计算后，将不同group的特征图进行concat，最后通过一个1x1卷积进行特征融合。

原始的W-MSA的计算复杂度是$2M^{2}HWC$，M是window的大小，HWC是特征图的尺寸；GMSA的计算复杂度是$2({\sum }_k(M_k^2))HWC/K$，其中K是group的数量，M_k是第k个group的window大小。

• Accelerated self-attention (ASA)

对slef-attention加速改进

（1）使用BN替代LN，BN可以融合到卷积中；
（2）之前的self-attention，q,k,v分别由1x1卷积得到，现在改为q,k相互复用，节省了计算开销；

作者验证在不影响SR效果的前提下，可以提升速度；

• Shared attention

Shared attention是在连续的n层self-attention之间共享attention scores，这样可以节省n次1x1卷积和2n次reshape以及n次矩阵乘，作者发现在n比较小（1或2）时，可在轻微降低超分效果的同时，提升速度。

• Shifted window
  

作者观察到swin transformer中的circular shift + mask attention对超分任务来说是冗余的，因为在circular shift之后，尽管会有像素被移动到很远的地方，但是超分任务中的特征图尺寸很大，这部分像素占所有像素的比例很小，因此它们对超分效果的影响可以忽略。因此，本文改进了swin transformer中的SW-MSA，去除mask机制和相对位置编码，只保留了circular shift + GMSA + inverse circular shift，提升了效率。

3. 结果

• 模型配置
  ELAN：36个ELAB，通道数为180
  ELAN-light: 24个ELAB，通道数为60
  GMSA的window size为：4 × 4, 8 × 8 and 16 × 16
  shared attention中n=1，也就是相邻的两个msa共享attention scores
• 速度对比
  ELAN-light比SwinIR-light快4.5倍
• ablation
  表3中的第2行是在swinIR中去除了残差连接，271ms -> 247ms
  shifted windows: 247ms -> 177ms
  ASA: 177ms -> 66ms