[논문 리뷰] Restormer: Efficient Transformer for High-Resolution Image Restoration

오늘 리뷰할 논문은 Restormer: Efficient Transformer for High-Resolution Image Restoration이라는 페이퍼입니다. 해당 논문에서는 고해상도 이미지의 품질 개선 task를 해결하기 위한 모델을 제안하였습니다. 구조만 놓고 보면 Transformer를 U-Net에 적용한 모양처럼 보입니다. 연세대학교에 계시는 Ming-Hsuan Yang 교수님이 교신저자로 참여하신 논문이라고 합니다.

일반적으로 Transformer 모델은 모델 자체가 가지고 있는 파라미터의 수와는 별개로 모델에 들어오는 token의 갯수가 늘어날 수록 연산에 드는 computational cost가 늘어나게 됩니다. 비전 분야에서는 이미지 사이즈가 커진다는 것이 곧 token 갯수의 증가를 의미하기 때문에 고해상도 이미지 관련 task에서 상당히 많은 computational cost를 필요로한다는 문제 있습니다. 심지어 늘어나는 너비, 높이의 제곱만큼 token(patch)가 급격하게 늘어납니다. 예를 들어(w, h): w x h patches 라면 (w+k, h+k): (w+k) x (h+k) patches가 됩니다.
(https://arxiv.org/abs/2111.09881)
(해당 글은 개인적인 기록을 목적으로 잘못된 내용이 있을 수 있음을 알립니다.)

0. Abstract

다양한 비전 task들을 해결하기 위해서 딥러닝 분야는 지금까지 CNN계열의 모델들을 주요 사용하였습니다. 어느정도의 성공은 있었지만, CNN을 사용하는 모델은 제한적인 receptive field 등의 구조적인 단점에 영향을 피할 수는 없었습니다. Transformer가 비전 분야에 정착하기 시작하면서 transformer block들은 receptive field의 제한이 없기 때문에 CNN기반 모델들의 문제를 어느정도 보완할 수가 있었습니다. 하지만 transformer 계열의 모델은 이미지 사이즈 혹은 해상도가 증가함에 따라 연산 부하가 제곱 배로 늘어난다는 문제가 있었습니다. Restormer는 이러한 연산 부담을 최대한으로 낮추면서 고해상도 이미지 개선 task들에 대해서 성능을 SOTA급으로 끌어올렸습니다.

1. Introduction

Image Restoration, Style Translation, Image Debluring 등과 같은 task들은 고질적으로 ill-posed한 성질을 가지고 있습니다. 보통 특정한 문제에 대한 해가 유일한 경우에 well-posed problem이라고 하며, 그렇지 않은 경우에는 ill-posed problem이라고 합니다. 위에서 말씀드린 task들을 보면 이미지를 복원하거나, 스타일을 바꿀 때 딱 정해진 해답이 존재하지 않습니다. 가령 성능이 저하된 이미지 A를 고품질 이미지 B로 바꾼다고 했을 때, 픽셀 값이 조금 다른 B'로 해를 구한다고 해서 오답이 아닌 것을 예로 생각해볼 수 있습니다.

이러한 세팅에서 task를 잘 수행하기 위해서는 강한 prior를 학습하고 대량의 이미지를 가지고 일반화를 잘하는 CNN계열 모델이 제격이었습니다. 물론 앞에서 말했던 것처럼 CNN은 limited receptive field 문제가 있고, 한 번 학습한 모델을 inference을 할 때 항상 weight 값이 고정되어있기 때문에 다양한 input context에 대해 유동적으로 대응하기가 어렵다는 문제도 있습니다. 반면 Self-attention(SA)은 들어오는 token 간의 관련성을 기반으로 계산을 하기 때문에 앞에서 언급한 문제에서 상대적으로 자유로울 수 있습니다.

하지만 글을 시작하기 앞서 말씀드렸던 것처럼 computational cost에서 자유로울 수가 없었기 때문에 이를 해결하기 위한 다양한 연구가 진행됐었습니다. SA의 계산 범위를 한정하는 등의 transformer를 사용해서 얻는 장점을 무색하게 만드는 해결책들이었습니다. Restormer는 computational cost를 낮추어 큰 이미지에 적용이 가능하면서도 SA를 이용한 global connectivity(local receptive와 반대)의 특징을 유지할 수 있도록 하였습니다.

해당 논문에서는 연산 효율은 높이기 위해서 Transformer block의 multihead-SA 부분을 multi-Dconv head 'transposed' attention(MDTA) block이라는 새로운 모듈로 교체하였고 feed-forward network는 gated-Dconv FN이라는 네트워크로 교체를 하였습니다. 여기서 Dconv는 Depth-wise convolution을 의미하고 채널 별로 convolution을 따로 해주는 연산이라고 보시면 됩니다. 아래서 자세히 말씀드리겠지만 FC layer를 최대한 Dconv layer로 교체하는게 연산 부하를 낮춘 비결 중에 하나입니다.

Dconv는 Xception: Deep Learning with Depthwise Separable Convolutions라는 논문에서 처음 제시되었으며 이를 가지고 연산 효율을 증대한 예시로는 MobileNet이 있습니다. 아래는 Xception 논문에서 나온 Depthwise separable convolution의 모양을 그림으로 표현한 것입니다. 일단 1x1 conv로 채널 수를 조절하고 채널 별로 별개의 conv연산을 해준 후에 다시 붙이는 방식으로 계산이 됩니다. 이 방식이 왜 더 효율적인지는 MobileNet 관련 글들을 찾아보시면 자세히 나올 것 같습니다.

Depthwise Separable Convolution

해당 논문의 Contribution은 아래와 같습니다.

1. Multi-scale local-global representation을 학습할 수 있는 encoder-decoder transformer 구조인 Restormer를 제안
2. MDTA를 제안하였습니다. 고해상도 이미지에서 효율적으로 local and non-local pixel interaction이 가능합니다.
3. GDFN을 제안하였습니다.

위 모듈들의 장점에 대해서는 Method 섹션에서 이야기 해보겠습니다.

2. Background

Image Restoration task와 ViT에 대한 이야기를 하고 있습니다. 궁금하신 분들은 한 번 읽어보시면 좋을 것 같습니다.

3. Method

Restormer 구조

위 그림은 Restormer의 전체 구조입니다. 맨 앞에서 말씀드렸던 것처럼 표면적으로는 U-Net구조를 보입니다. 이미지가 들어오는 낮은 층(low level)에서는 (H, W, C) 차원으로 연산을 하고 깊은 층(high level)로 내려갈 수록 H, W는 줄어들고 C(채널 수)가 encoder와 decoder에서 대칭적으로 늘어나는 모습을 보입니다. 들어오는 데이터가 고해상도 이미지라면 H, W가 너무 크기 때문에 C(채널 수)를 막 늘릴 수가 없습니다. 그래서 H와 W를 줄이면서 C채널 수를 늘리는 모습입니다. 또한 보이듯이 Transformer block 안의 multihead-SA와 FFN은 각각 MDTA와 GDFN으로 교체된 모습입니다. 들어오는 이미지는 task에 맞게 품질이 저하된 고해상도 이미지가 들어오고 모델을 고해상도 고품질로 이미지를 복원해서 결과를 반환합니다.

전체 구조를 보면 encoder를 통해 구한 low-level feature를 decoder까지 길게 통과시켜서 얻은 high-level feature와 합쳐 주고 이를 합쳐주는 구조입니다. 이는 U-Net와 비슷한 원리인 것 같습니다.

3.1. Multi-Dconv Head Transposed Attention

위의 Restormer 구조 이미지에 MDTA 구조가 나와있습니다. 기존의 SA 같은 경우에는 각 픽셀이 하나의 token(patch)라고 가정했을 때 WxH 만큼의 token이 생기고 (WxH) x (WxH) 만큼의 self-attetion 연산을 하게 됩니다. 이는 이미지 사이즈가 커짐에 따라 연산랑이 제곱 배로 늘어나는 현상을 야기합니다. 여기서는 해당 문제를 완화하고자 (WxH) 기준의 self-attetion을 하지 않고 channel 단위의 C x C self-attention 연산을 합니다. 그래서 attention map의 크기가 CxC가 됩니다.

물론 앞에서 말씀드렸던 것처럼 DConv도 연산 효율을 높이는데 기여했습니다. Query, Key, Value를 구할 때 FC layer대신 DConv를 사용하였습니다.

3.2. Gated-Dconv Feed-Forward Network

여기에서도 Feed-Forward의 Fully-connected layer 대신에 DConv를 사용하여 연산 효율을 높였습니다. 또한 DConv를 사용해서 GDFN 내부에서 계산할 때 채널 수를 한 번 늘렸다가 밖으로 결과를 내보낼 때 들어왔던 채널 수로 다시 맞춰줍니다.

3.3. Progressive Learning

작성 중

4. Experiments and Analysis

5. Conclusion