[논문 리뷰] SwinIR: Image Restoration Using Swin Transformer

2020년에 Vision Transformer가 등장하면서 많은 vision task에서 transformer를 사용한 모델이 매우 좋은 성과를 보이고 있다. 본 논문에서는 Swin Transformer와 여러 structural design을 통해서 attention-based architecture가 image reconstruction tasks (ex. denoising, super-resolution, compression)에서 SOTA성능을 낼 수 있다는 것을 보여준다. 

https://arxiv.org/abs/2108.10257

SwinIR: Image Restoration Using Swin Transformer

 

최근 image reconstruction과 ViT에 관심이 많기 때문에 당연하게 눈이 간 논문이었고 성능도 뛰어났기 때문에 한번 리뷰 해보고 싶었다. 이제 뛰어들어보도록 하자!!

---

1. Introduction

(왼쪽) denoising, (오른쪽) super-resolution

 

image restoration이란, low-quality image로부터 high-quality image를 reconstruct하는 문제를 말한다. 그 예로는 noise가 있는 image로부터 noise가 제거된 clean image를 얻는 denoising task와 low-resolution image로부터 high-resolution image를 만들어내는 super-resolution 등이 있다.

 

이 image restoration 분야에서는 오랜 기간 동안 Convolutional Neural Network (CNN)이 dominant한 method로써 지금까지도 많이 사용되는데, 몇 가지 단점이 존재한다.

1. images와 convolutional kernels 간의 interactions이 content-independent하다.
   • convolution kernel은 shared parameters로써 모든 image region에 같은 kernel이 사용되는데 이러한 방법이 별로 좋지 않을 수 있다. 
2. convolution은 long-range dependency를 효과적으로 모델링 할 수 없다. 

 

위와 같은 문제를 해결하기 위해서 최근 CNN과 비슷하거나 더 높은 성능을 보여주고 있는 Vision Transformer를 사용하여 image restoration을 하는 몇 가지 논문이 있었는데, 이들은 이미지를 patch 단위로 나누면서 각 patch를 독립적으로 처리하기 때문에 다른 문제가 발생한다.

1. patch 단위로 attention을 하기 때문에 border pixels가 neighbouring pixels에 대한 정보를 이용하지 못한다.
2. 마찬가지로 patch 단위의 restoration이기 때문에 restored image에 border artifacts가 생길 수 있다. 이는 patch를 overlapping 시킴으로써 어느 정도 해결이 가능하지만 추가적인 computation이 생기게 된다.

 

이러한 문제들은 patch 내에서만 attention을 한다는 것이 큰 원인이므로 (ex. stand alone self attention) ViT처럼 patch 간 attention을 할 수 있다면 어느 정도 해결이 될 것이다. 하지만 ViT의 경우 연산량과 model size가 너무 크기 때문에, 논문에서는 local attention mechanism을 사용하여 large size image도 처리할 수 있는 Swin Transformer를 기반으로 architecture를 구성한다.

 

2. Method

 

Swin Transformer를 기반으로 하는 architecture이기 때문에 SwinIR이라고 이름을 정해주었고 크게 3 개의 modules로 구성되어 있다.

 

첫 번째로 Shallow feature extraction module은 convolution layer를 사용하여 shallow feature를 extract하며, reconstruction module에 direct하게 전달되어 low-frequency information이 보존되도록 해준다.

 

두 번째로 Deep feature extraction module은 residual Swin Transformer blocks (RSTB)으로 이루어지며, 여러 개의 swin transformer layers를 통해 local attention과 corss-window interation으로 deep features를 extract한다. 추가로 feature enhancement를 위해 block의 끝에 convolution layer를 추가하였고 feature aggregation을 주기 위해 residual connection을 사용하였다.

 

마지막으로 Reconstruction module에서는 위 두 모듈에서 뽑아낸 shallow와 deep features를 합쳐서 최종 이미지를 출력한다.

 

2.1 Shallow  feature extraction

low quality input $I_{LQ} \in \mathbb{R}^{H\times W \times C_{in}}$이 있다고 할 때 3x3 convolutional layer $H_{SF}(\cdot)$을 사용하여 shallow feature $F_0 \in \mathbb{R}^{H\times W \times C}$를 뽑아낸다.

$$F_0=H_{SF}(I_{LQ})$$

 

초반 shallow layer에 CNN을 사용하는 이유는 convolution이 low dimensional information을 처리하는 데 좋은 성능을 보이고 간단하게 input image를 고차원 feature space로 mapping할 수 있기 때문이다.

2.2 Deep feature extraction

deep feature extraction module $H_{DF}(\cdot)$는 K개의 residual Swin Transformer blocks (RSTB)와 3x3 convolutional layer로 이루어져있고, input으로 shallow feature extraction의 output인 $F_0$가 들어간다.

$$F_i = H_{RSTB_i}(F_{i-1}), \ \ \ i=1,2,\dots ,K$$

$$F_{DF}=H_{CONV}(F_K)$$

$H_{RSTB_i}(\cdot )$은 i번째 RSTB를 의미하고 가장 마지막에 3x3 convolution인 $H_{CONV}$를 거치게 되는데, 이렇게 마지막에 convolution을 사용함으로써 transformer-based network에 convolution operation의 inductive bias를 더해주는 효과가 있다고 한다.

Residual Swin Transformer Block (RSTB)

RSTB는 그냥 이전 Swin transformer block에다 residual connection과 convolutional layers를 더해준 것이 끝이다. RSTB는 L개의 Swin transformer layers로 구성되어있고 i 번째 RSTB의 input feature $F_{i,0}$를 고려할 때 식은 다음과 같다.

$$F_{i,j} = H_{STL_{i,j}}(F_{i,j-1}), \ \ \ j=1, 2,\dots ,L$$

 

여기서 $H_{STL_{i,j}}$는 i번째 RSTB의 j번째 Swin transformer layer를 뜻하고 여기서도 모든 layers를 거친 뒤에 convolution layer와 residual connection이 추가 된다. 

$$F_{i, out} = H_{CONV_i}(F_i, L)+F_{i,0}$$

**$H_{CONV_i}(\cdot )$은 i번째 RSTB의 convolutional layer

 

이와 같은 구조는 두 가지 장점을 가지는데 첫 번째는 convolutional layers의 spatially invariant filters (image region마다 같은 kernel을 가지는 것)가 SwinIR의 translation equivariance를 enhancement 시켜준다는 점이고, 두 번째는 residual connection을 통해서 다른 levels의 features를 aggregation할 수 있도록 만들어준다는 점이다.

 

**Residual connection을 layer, block, network 단위로 사용하기 때문에, Patch Merging을 사용한 down-sampling을 하지 않는다. 

Swin Transformer layer (STL)

여기서 사용되는 STL은 기존 Swin transformer의 layer와 같으므로 깊이 다루지는 않을 것이다.

간단히만 설명하면 self-attention이 token length에 quadratic하게 연산량이 커지는 문제가 있어서 large resolution input에 사용하기 어렵다는 문제가 있었는데, 이를 고정된 크기 M의 non-overlapping window를 지정하여 $H \times W \times C$ features를 $\frac{HW}{M^2} \times M^2 \times C$ 형태로 바꾼 다음에, 먼저 각 $M \times M$ window 내에서 self-attention을 하고 그 후에 shifted window partitioning을 통해서 window 간의 connection까지 더해주어서, 연산량을 크게 줄이면서 성능의 저하는 최소한으로 만든 것이 Swin인 것이다. 

 

자세한 내용은 추후에 리뷰를 할 수도 있고 빨리 알고 싶다면 논문을 읽거나 다른 블로그 리뷰를 보는 것도 추천한다.

2.3 Image Reconstruction

image reconstruction module은 shallow features와 deep features를 aggreagateg하는 역할을 한다.

$$I_{RHQ} = H_{REC}(F_0+F_{DF})$$

**$H_{REC}$: reconstruction module, $F_0$: shallow features, $F_{DF}$: deep features

 

reconstruction module $H_{REC}$는 task에 따라 달라질 수 있는데, SR의 경우 sub-pixel convolution layer를 통해 upsample을 하고, 그 외에 denoising과 compression은 single convolution layer를 사용한다.

Loss function

loss도 SR의 경우 L1 pixel loss, denoising과 compression은 Charbonnier loss를 사용한다.

$$L_1 =||I_{RHQ}-I_{HQ}||_1$$

$$L_{Charbonnier}=\sqrt{||I_{RHQ}-I_{HQ}||^2+\epsilon^2} $$

**논문에서는 $\epsilon=10^{-3}$으로 정해주었다.

3. Experiments

결론을 먼저 말하자면 SwinIR이 Super-resolution, Denoising, Compression 3가지 tasks에서 이전 methods보다 적은 parameters를 가지면서 SOTA의 성능을 달성하였다. 이번 리뷰에서도 experiments를 자세히 다루지는 않을 것이고 결과만 살펴보도록 하겠다. ㅎㅎ

3.1 Super-resolution

먼저 super-resolution task에 대해서 swinIR을 사용했을 때 모든 경우에서 SOTA 성능을 보여준다. 위 이미지에서도 다른 method보다 blur가 덜하고 detail이 잘 살아있다.

**빨간색이 best score이고 파란색이 second best score이다.

3.2 Image Compression

image compression에서도 SOTA성능을 보여주고 있다..ㅎㅎ..별로 설명할 게 없네..

 

3.3 Image Denoising

denoising 또한 SOTA성능을 보여준다.

 

3.4 Ablation Study

 

마지막으로 SwinIR의 여러 다른 setting에 따라서 성능이 어떻게 변화하는 지를 실험한 ablation study이다. 차례대로 살펴보면,

1. channel number가 커질 수록 성능이 높아진다.
2. RSTB의 수를 늘릴 수록 성능이 높아진다.
3. RSTB의 layer 수를 늘릴 수록 성능이 높아진다.
4. patch size를 크게 할수록 성능이 높아지고 CNN-based method보다 그 폭이 더 높다.
5. training images의 수를 늘릴 수록 성능이 높아진다. 
   • 이전에 알려진 transformer-based models은 CNN보다 더 큰 datasets을 필요로 하였는데, SwinIR은 같거나 더 적은 데이터셋으로도 CNN보다 높은 성능을 보여준다.
6. CNN-based method (RCAN)보다 더 빠르게 수렴한다.

 

RSTB design에 대한 실험도 있는데, residual connection을 사용하지 않았을 때 0.16 dB정도의 성능 차이가 있었고 ( 오....생각보다 별로 도움이 안된다고 생각이 들 수 있지만 SOTA를 찍기 위해선 소수점 단위도 중요하기 때문에...) 1x1 convolution을 썼을 때 local neighbouring information을 뽑아낼 수 없기 때문에 3x3 보다 성능이 낮게 나온다 ( 근데 파라메터 수 측면에서 보면 3x3보다 더 적게 사용하기 때문에 마냥 안 좋다고 볼 순 없을 것 같다). 마지막으로 3x3 conv 3개를 사용한 것인데 채널은 4분의 1정도 낮게 준 것으로 보이고 파라메터 수는 줄어들었으나 성능이 약간 감소하였다고 말한다.

---

이렇게 SwinIR 논문을 살펴보았는데, Vision transformer 모델로 image reconstruction tasks에서 SOTA 성능을 낼 수 있다는 점이 흥미로웠다. 이 논문의 reference와 searching을 통해서 다른 transformer-based IR 논문을 더 찾아볼 예정이고 좋은 논문이 있다면 리뷰를 하도록 하겠다 ㅎㅎ.