[논문리뷰] Simple Baselines for Image Restoration

Sep 13, 2023

Simple Baselines for Image Restoration

Paper: https://arxiv.org/pdf/2204.04676.pdf (ECCV 2022)

Introduction

최근 Computer vision task에서 다양한 방법론과 구조가 등장함에 따라 성능이 매우 향상되었습니다.

하지만 이러한 성능 향상을 위해 computational cost가 커지고, 모델 구조 (혹은 레이어) 는 더욱 복잡해졌습니다.

논문의 저자는 이러한 system complexity 를 1)inter-block complexity, 2)intra-block complexity 두가지로 나누어 분류합니다.

Inter-block complexity

Inter-block complexity에는 대표적으로 multi-stage, multi-scale fusion 등이 있습니다.

각각의 간단한 예로는

FastDVDNet처럼 앞의 Unet 에서는 feature를 추출하고 다음 stage의 Unet에서 해당 feature를 refine하는 구조는 multi-stage architecture에 해당하고

MPRNet 처럼 여러 size의 feature map 끼리 fusion 해주는 경우가 multi-scale fusion에 해당한다고 할 수 있습니다.

Intra-block complexity

Intra-block complexity는 Restormer, SwinIR 과 같이, 별도의 커널이나 복잡한 로직을 구현해야하는 경우라고 할 수 있습니다.

python이 지원되는 GPU server 에서는 해당 system-complexity가 크게 문제가 되지는 않지만, edge-device나 NPU 등 system, resource가 제약적인 환경에서는 문제가 될 수 있습니다.

Simple Baseline Network

논문의 저자는 이러한 system-complexity를 최소화 하면서 SOTA 급의 성능을 낼 수 있는 방법을 연구합니다.

이를 위해 우선 Baseline이 될 simple network를 구성하는데, inter-block complexity를 최소화 하기 위해서 일반적으로 vision task에서 사용되는 single-stage UNet (Plain UNet) 의 구조를 차용하였습니다.

UNet을 구성하는 block은 마찬가지로 conv-relu + skip connection의 plain 조합을 사용하였습니다.

이와 같은 Plain Unet을 기반으로, 이후에 서술될 내용과 같이 component들을 교체/추가 하여 SOTA급의 성능을 낼 수 있는 16GFLOPs (256x256) 크기의 Baseline Network를 구성하게됩니다.

Normalization

Normalization은 모델의 학습에 있어서 도움이 되는것은 맞지만, batch norm, Instance norm 과 같은 경우 batch size 등과 같은 요소들에 의한 불확실성이 존재합니다.

그렇기 때문에 실제로 최근 vision model의 경우 normalization을 제외시키는 경우가 대부분입니다.

하지만 저자는 최근 transformer 구조를 사용하는 SOTA 모델에서 사용하는 Layer Normalization에 주목하고, 이를 network에 도입하여 실험하였습니다.

그 결과, 학습에 안정성이 크게 증가하여 더 높은 learning rate를 사용할 수 있으며, 성능 또한 상향되었다고 하였습니다.

(denoising과 deblurring에서 각각 +0.44dB, +3.39dB의 PSNR 상승이 있었다고 하는데 크게 유의미한 수치인지는 모르겠습니다.)

Activation

일반적으로 Vision task에서는 activation function으로 ReLU나 Leaky ReLU를 주로 사용합니다.

저자는 최근 transformer 기반 SOTA 모델들이 GeLU function을 사용하는점에 주목하여 activation function을 GeLU로 교체하였습니다.

이를 통해 denoising task에서는 성능을 유지하고 (-0.02dB), deblurring task에서는 +0.2dB 의 성능 향상이 있었다고 합니다.

(이또한 수치적으로 의미는 별로 없고, 교체 이유또한 타당성이 있어보이진 않습니다.)

Attention

Transformer 구조가 vision task에 적용되면서 SOTA 성능을 낼 수 있게된 가장 핵심적인 요인은 self-attention 이라고 할 수 있습니다.

하지만 self-attention은 input size에 비례하여 연산량이 quadratic 하게 증가한다는 한계점이 있습니다.

이를 해소하기 위해서 SwinIR 등에서 window attention을 도입한 image restoration 모델을 제안하였지만, 연산량을 해결하는 대신 global context를 캐치하는 능력과 구현이 복잡하여 intra-block complexity가 상승한다는 단점이 있습니다.

따라서, Baseline network에는 self-attention 대신 local information을 위해 depthwise convolution을, global information을 위해 channel attention을 추가하였고, 이는 각 feature를 추출하는데에 충분하다고 주장합니다.

Baseline Network Summary

최종적으로 Baseline network는 위와 같은 Block을 가지는 UNet architecture가 됩니다.

이렇게 구성된 baseline network가 system-complexity가 높은 다른 모델들과 비교하여 성능면에서 뒤떨어지지 않는다는것을 실험을 통해 입증합니다.

Nonlinear Activation Free Network

결과에서 확인할 수 있듯, Baseline network는 기존 SOTA 모델들과 비교하여도 comparable한 성능을 보여줍니다.

하지만 논문에서는 이에 그치지 않고 simplicity를 유지하며 성능을 더 향상시킬수는 없는가? 성능을 유지하며 모델을 더 simple하게 만들수 없는가에 대해서 추가적인 고민을 하게됩니다.

Gated Linear Unit

저자는 고민에 대한 해답을 GLU(Gated Linear Unit) 에서 찾습니다.

Gate(X, f, g, \sigma)=f(x)\odot \sigma({g(X))},

GLU는 이미 수많은 transformer 기반 SOTA 모델에서 차용하고 있지만, 이것을 단순히 모델에 끼워넣는것은 intra-block complexity를 증가시키므로 목표에 알맞지 않습니다. 그렇기 때문에 아래의 과정을 통해 GLU를 단순화 합니다.

GELU(x)=x\Phi(x),

0.5x(1+tanh[\sqrt{2/\pi}(x+0.0044715x^3)]).

GELU의 근사식에 따라서, GELU는 GLU의 special case라고 할 수 있습니다.

따라서 GLU는 1)activation function으로 사용될 수 있고, $\sigma$ 를 제거하더라도 2)그 자체로 nonlinearity를 갖습니다.

이에 근거하여 저자는 GLU의 simple 버전인 SimpleGate를 제안하고 GELU를 대체합니다.

SimpleGate(X,Y) = X\odot Y

GELU의 복잡함과 달리, SimpleGate는 단순히 feature map의 channel을 반으로 나누어 곱해주는것으로 구현됩니다.

Baseline network의 GELU를 SimpleGate로 교체하였을때, 기존과 비교하여 성능 차이가 없거나 성능이 소폭 상승하는것을 확인했습니다.

Simplified Channel Attention

GELU를 SimpleGate로 대체하게 되면, Baseline network에서 nonlinear activation function은 오직 channel attention 내의 ReLU, Sigmoid만 존재하게 됩니다.

#Channel attention block
channel_attention = nn.Sequential(
	nn.AdaptiveAvgPool2d(1)
  nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
  nn.ReLU(inplace=True),
  nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
  nn.Sigmoid()
)

x = x * channel_attention(x)

일반적으로 channel attention은 위와 같은 과정으로 이루어지게 되는데, channel attention 을 하나의 function $\Psi$ 로 본다면 다음과 같은 형태로 다시 쓸 수 있습니다.

CA(X)=X*\Psi(X)

이 형태는 앞선 GLU 와 유사하게 볼 수 있고, 그렇다면 Channel Attention 또한 simplify 할 수 있다고 추측할 수 있습니다.

Channel attention의 핵심적인 역할만을 남겨두고 단순화를 하면 다음과 같이 Simple Channel Attention을 사용할 수 있습니다.

SCA(X)=X*Wpool(X)

Channel attention을 Simple Channel Attention으로 교체하였을때, 성능이 유지되거나 소폭 상승하는 결과를 얻을 수 있었다고 합니다.

Summary

GELU와 Channel Attention을 교체하는것으로 network에는 non linear activation function이 존재하지 않게 되었고, 성능 또한 유지되거나 소폭 상승하였습니다.

Conclusions

결론적으로 해당 논문에서는 single-stage UNet 구조를 사용하여 기존 목적인 inter, intra complexity를 줄였고, 거기에 더해 network 내의 nonlinear activation function을 모두 제거하여 simplicity를 더욱 높임과 동시에 SOTA급의 성능을 유지하였습니다.

저자는 이러한 network를 NAFNet(Nonlinear Activation Free Network) 라고 제안하였습니다.

NAFNet은 단순한 구조와 적은 계산량으로도 기존의 denoising, deblurring SOTA 모델과 동등하거나 더 좋은 성능을 나타냈습니다.

또한 Block 수를 조정하더라도 작은 모델부터 큰 모델까지 안정적으로 우수한 성능을 보여주어 scalable 하다는 장점을 갖고 있습니다.

Simple Baselines for Image Restoration