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

Abstract

 

Image restoration은 오랫동안 지속되어온 low-level vision problem으로서, low-quality images로부터 high-quality images를 복원해내는 작업이다. SOTA 모델이 CNN-based인 상황에서, high-level vision task에서 impressive한 performance를 보여주던 Transformer를 Image restoration에 적용해보려는 시도가 있었다. 이 논문에서는, Swin Transformer를 기반으로 한 image restoration 방식인 SwinIR을 제안한다.

 

 

SwinIR은 크게 3가지 파트로 나뉜다. 

1. Shallow feature extraction
   
   
2. Deep feature extraction
   • several residual Swin Transformer blocks(RSTB) 존재
   • 각 RSTB에는 several Swin Transformer layers 존재(with a residual connection)
     
     
3. High-quality image reconstruction

저자들은 SwinIR을 이용하여 3가지 대표적인 task에 대한 실험을 진행하였다.

1. Image super-resolution
   • classical, lightweight(benchmark dataset), real-world image(현실에서 찍은 이미지)
     
     
2. Image denoising
   • grayscale, color image
     
     
3. JPEG compression artifact reduction 

SwinIR은 각각의 task에서 SOTA 모델보다 최대 0.14~0.45 dB(PSNR 기준) 만큼 outperform하였다.

Total number of parameters도 최대 67%까지 감소시켰다.

 

 

여기서는 Image super-resolution의 관점에서만 SwinIR에 대해 알아볼것이다.

Introduction

 

Image restoration에서 가장 중요한 것은 low-quality image에서 high-quliaty의 clean한 image를 만들어내는 것이다.

 

Limits of CNN

지금까지는 CNN이 image restoration에서 가장 주력 모델로 사용되고 있었다. 대부분의 CNN-based methods는 residual learning, dense connection과 같이 정교한 architecture design에 집중하고 있었다. 전통적인 모델에 비해서 CNN 기반 모델로 인해 performance가 많이 향상된 것은 맞지만, 여전히 convolution layer에서 나오는 2가지 basic problem이 존재하였다.

1. Images와 convolution kernel(filter)간 interaction은 content-independent하다.
   • 다른 이미지 영역을 복구하는데 동일한 convolutional kernel을 사용하는 것은 best choice가 아닐 수 있다.
     
     
     
2. local processing의 원리에 의하면, convolution은 long-range depedency modelling에 효과적이지 않다.
   • locality : 입력에서 각 entity(pixel)간의 관계가 서로 가까운 요소들에 존재
   • CNN은 계속 filter를 통해 local하게만 feature를 보기 때문에, layer를 쌓아도 전체 범위를 한번에 볼 수 없음

Limits of vision Transformer(viT)

이러한 문제를 가진 CNN을 대체하기 위해, self-attention mechanism을 탑재하여  contexts간 global interactions을 포착할 수 있는 Transformer가 등장해서 많은 vision problem에 대해 좋은 모습을 보여주었다. 그러나 image restoration을 위한 vision Transformers는 input image를 48*48의 고정된 size의 patch로 분할하여 각 pach를 independently하게 processing 한다는 특징이 있었다. (이미지에서 patch = NLP에서의 token) 

 

이러한 특징은 두가지 결점을 야기하였다.

 

1.  Border pixels들이 patch 바깥으로 나가는 neighbouring pixel를 utilize할 수 없다.
   • patch 단위로 attention을 하기 때문에, border pixel이 바깥으로 나가는 pixel 정보 이용 X
     
     
2. Restored image들이 각 patch 주위로 border atrifacts를 만들어낼 수도 있다.
   • Patch overlapping으로 이 문제를 완화할 수 는 있지만, 추가적인 computing 연산이 필요하다.
     
     

즉, patch 내에서만 attention이 진행된다는 것이 가장 큰 원인

>> SwinIR에서는 local attention + shifted window scheme으로 이 문제 해결

 

Apperance of Swin Transformer and  Advantages of Swin Transformer

최근에 등장한 Swin Transformers는 CNN과 Transformer의 장점을 결합하여 엄청난 가능성을 보여주고 있다.

 

1. Swin Transformer는 local attention mechanism을 활용하여 large size의 이미지를 처리할 수 있는 CNN의 장점을 가지고 있다. 
   
   
2. Swin Transformer는  shifted window scheme을 활용하여 long-range dependency를 modelling 할 수 있는 Transformer의 장점을 가지고 있다.
   
   

Outline of SwinIR 

SwinIR은 크게 3가지 모듈로 이루어져있다. 

 

1. Shallow feature extraction
   
   • shallow feature를 추출하기 위해 convolution layer를 사용
   • 추출된 feature는 reconstruction module에 직접 전달되어 low-frequency information을 보존할 수 있다.
2. Deep feature extraction
   • several residual Swin Transformer blocks(RSTB) 존재
   • 각 RSTB에는 several Swin Transformer layers 존재 : local attention + cross-window interaction 진행
   • RSTB의 끝에는 convolution layer를 추가하여 feature enhancement 진행
   • RSTB 내부에서는 residual connection을 사용하여 feature aggregation을 위한 shortcut 제공
3. High-quality image reconstruction

 

SwinIR은 CNN-based image restoration 방식에 비해 여러가지 이점이 존재한다.

 

1. Image content와 attention weight간 content-based interaction이 존재한다. 이는 spatially varying convolution으로 해석될 수 있다. 
   
   
2. shifted window mechanism을 이용하여 long-range dependency modelling이 가능하다.
   
   
3. Better performance with less parameters.

PSNR results of several models

 

 

Method

 

Network Architecture

 

The architecture of the proposed SwinIR for image restoration

 

Shallow feature extraction 

Low quality input ILQ∈ H×W×Cin 이 있다고 할 때 3x3 convolutional layer HSF(⋅)을 사용하여 shallow feature F0∈ H×W×C를 뽑아낸다.  

*H, W, Cin, C : image height, width, input channel number, feature channel number( = filter channel number)

 

 

Convolution layer는 초기 visual processing에 좋고, more stable optimization이 가능하게 하며 better result를 만들어준다.

또한, 그것은 input image space를 high dimension space로 만들어주는 가장 간단한 방법이기도 하다.(filter channel만 늘리면 되기 때문에)

 

Deep feature extraction 

Deep feature extraction module HDF(⋅)는 K개의 residual Swin Transformer blocks (RSTB)와 3x3 convolutional layer로 이루어져있고, input으로 shallow feature extraction의 output인 F0가 들어간다. Input이 RSTB block들을 차례차례 지나고 마지막에 convolution layer를 만나게 되는 형태이다.  



Feature extraction 끝부분에 convolutional layer를 사용하는 데에는 2가지 이유가 있다.

 

1. Convolution operation의 inductive bias를 Transformer based network에 넣어줄 수 있다.
   *inductive bias: 보지못한 data에 대해서도 귀납적 추론이 가능하도록 하는, 알고리즘이 가지고 있는 가정의 집합
   • Transformers는 기본적으로 global context, 모든 정보를 활용하기 때문에 inductive bias가 부족하다.
   • CNN은 filter를 적용하는 위치에 따라 결과값이 달라진다, inductvie bias 존재.
2.  Shallow and deep features의 later aggregation을 위한 초석을 놓는 역할을 한다.

 

Residual Swin Transformer Block(RSTB)

 

 

 

RSTB는 Swin Transformer layers(STL, with residual block)과 convolutional layer로 구성되어 있다.

F(i,0)을 i번째 RSTB의 input feature라고 했을 때, L개의 Swin Transformer layers을 통해 F(i,1)부터 F(i,L)까지 추출한다(= intermediate features). 여기서 HSTL(i,j)는 i번째 RSTB의 j번째 Swin transformer layer를 뜻한다.

 

모든 layer를 거쳤으면 convolution layer와 residual connection이 추가된다.

 

이러한 design에는 2가지 장점이 있다. 

 

1.  Transformer가 spatially varying convolution의 특정 부분으로 실현된 것처럼 보여도, spatially invariant filters를 가진 convolutional layers는 SwinIR의 translational equivariance를 강화해줄 수 있다.
   
   • spatially invariant filter : image 위치에 따라 불변하는 필터. 모두 동일한 필터를 사용한다
   • Translational equivariance(=variance) : CNN에서는 필터가 동일해도 그 필터를 어떤 이미지 영역에 쓰는지에 따라 연산결과가 달라진다.
2. Residual connection은 different block에서부터 reconstruction module까지 identity-based connection을 제공한다.
   즉, different level of features에 대한 aggregation이 가능하게 한다.

 

Swin Transformer layer(STL)

 

 

RSTB안에 있는 STL은  Transformer의 standard multi-head self-attention을 base로 하고 있다. 

Transformer와의 차이점은 local attention과 the shifted window mechanism이다. 

 

MSA(Multi-head Self-Attention)

기존 self-attention은  token length에 quadratic하게 연산량이 커지는 문제가 있어서 large resolution input에 사용하기 어렵다는 문제가 있었는데,  Swin Transformer에서는 이를 고정된 크기 M의 non-overlapping window를 지정하여 local attention으로 만들어 해결하였다. 

 

예를 들면 input의 size가 H*W*C 일때, Swin Transformer는 H*W 이미지를 non-overlappng M*M local windows로 분할하여 input을 HW/M2 * M2 * C size로 변환한다. 여기서 HW/M2는 window의 total number를 의미한다. 

 

출처 :  DSBA 유튜브

 

이미지가 window로 분할되면, 각 window에서 standard self-attention이 따로따로 진행된다(= local attention)

 

Query, Key, Value matrices

 

Local window feature X∈ M2 × C에 대해 각각 PQ, PK,PV라는 projection matrices(다른 window도 공유하는 가중치)를 곱하여 query, key, value를 생성한다. Q, K, V의 size는 M2 × d 이다. 

 

Self-attention mechanism in a local window

 

앞에서 구한 Q, K, V를 바탕으로 local self-attention을 진행하면 위와 같은 식으로 연산이 수행된다. 여기서 나머지는 Transformer의 scaled dot-product attention과 동일하며 B가 추가된 것이 차이점이다. 

 

원래 Transformer에서는 attention mechanism의 특성상 시점이 반영이 안되서 postional encoding이라는 행렬을 input, output embedding에 더하여 위치정보를 표현하였다. 그 위치정보는 sin, cos함수의 주기로 구한 절대좌표였다.  

 

하지만 Swin Transformer에서는 절대좌표가 아닌 relative positional encoding을 사용하는데, 축 기준에 따라(x축, y축) 현 위치에서 해당하는 인덱스의 위치까지 얼마만큼 이동해야하는지를 표시한 값이다. 그 예시는 아래와 같다. 

 

출처 : DSBA 유튜브

 

위 예시를 참고하면 M*M local window의 relative positional encoding matrix의 size는 M2*M2 이기 때문에 위의 Attention 수식에서 QK(T)와 더해질 수 있다.

 

이런 local attention을 h번 병렬로 진행하고(multi-head attention), 최종적으로 진행한 결과를 concatenate한다.

 

하지만 이렇게만 local attention을 진행한다면 different layer에 대하여 partition(window 분할방법)이 고정되므로, local window간 connection은 없게된다. 이러한 문제를 해결하기 위해 실제로는 방금 진행한 regular window partitioning과 shifted window partitioning가 교차되면서 attention이 수행되어 cross-window connection이 가능해진다. ( MSA = W-MSA(window) + SW-MSA(shifted window)) 

 

여기서 shifted window partitioning이란 feature map을 ([M/2],[M/2]) 만큼 우측하단으로 shift하고 다시 window partitioning을 진행하는 것이다.

 

출처: DSBA 유튜브

 

MLP(Multi-layer Perceptron)

MSA 단계가 끝나면  further feature transformation을 위해,  GELU 활성화 함수를 가진 two fully-connected layers 구조로 이루어진 MLP 과정이 수행된다.

 

 

Whole process

STL의 전체 과정을 요약하면 다음과 같다. 

 

X가 STL layer의 input이라고 할 때, LayerNorm layer가 MSA와 MLP가 일어나기 전 X에 추가된다.

그리고 residual connection이 있기 때문에  identity X가 MSA와 MLP가 일어난 후 다시 더해진다.

 

 

Image Reconstruction

 

F0 : shallow features, F(DF): deep features

 

SwinIR의 마지막 모듈인 Image reconstruction에서는 shallow and deep features를 aggregate하여 high-quality image인 I(RHQ)를 만들어낸다. H(REC)는 reconstruction module의 함수이다. 

 

Shallow feature는 주로 low-frequencies를 포함하고 있고, deep features는 lost high-frequencies를 복구하기 위해 집중한다.

*low-frequency: pixel values that are changing slowly over space(body part) -  original image에 convolution을 진행하여 get

 high-frequency: pixel values that are changing rapidly(edge part)

 

그림에서 보이는 Long skip connection이 shallow feature, 즉 low-frequency information이 directly하게 reconstruction modeule에 전달되도록 해준다. 이는 deep feature extraction module이 high-frequency information에만 집중하게 해주고, training을 안정화시켜준다. 

 

reconstruction module function에는 sub-pixel convolution layer가 들어가 있어 원하는 scale대로 feature를 upsampling할 수 있다. (ex) 이미지의 H, W를 각각 r배하고 싶으면, channel을 r*r 개만큼 만든뒤, sub-pixel을 거쳐서 channel을 1개로 변환시킨다)

 

Sub-pixel convolution layer

 

 

최종적으로, 저자들은 residual learning을 이용해 HQ image를 reconstruct하는 것이 아닌, LQ와 HQ image의 residual을 reconstruct하는 방식을 제안하였다. 이 식은 다음과 같다. H(SwinIR)은 SwinIR 모델 전체의 function이다.

 

Loss function

일반적인 Image super resolution에서 SwinIR의 loss로는 L1 pixel loss(이미지 자체)를 사용한다. (classical, lightweight image)

I(RHQ)는 SwinIR의 output, I(HQ)는 ground-truth HQ image이다. 

 

그러나, real-world image에서는 visual quality를 높이기 위해 pixel loss, GAN loss, perceptual loss (ESRGAN에서 차용)를 combination한 loss를 사용한다.