Background
이미지 복원은 노이즈, 블러, 비, 왜곡 등 다양한 열화를 제거해 "깨끗한 이미지"를 복원하는 고전적인 저수준 비전 과제입니다. 최근 SOTA는 주로 CNN 기반이지만 이는 지역적인 패턴 복원에는 강하지만 장거리 의존성 포착에는 한계를 보인다는 점이 핵심적인 문제로 지적되고 있습니다. 이 한계를 보완하기 위해 일부 연구들은 self-attention을 도입했지만 전역 self-attention의 계산량이 토큰 수에 대해 제곱으로 증가해 고해상도 특징맵에 직접 적용하기 어렵기 때문에 주로 저해상도에서 "일부 레이어만" 제한적으로 쓰는 경우가 많았습니다.
본 논문에서는 이미지의 디테일있는 복원을 위해 여러 해상도 (멀티스케일) 특징맵에서 self-attention을 적극적으로 활용하되 계산량 문제를 피할 수 있도록 활용하고자 하였습니다. 그래서 본 논문에서는 UFormer를 제안하였으며 이름에서도 알 수 있다싶이 UNet처럼 계층적인 인코더-디코더 그리고 skip connection을 유지하면서 내부 블록을 convolution 대신 Transformer 블록으로 구성한 U-Shaped Transformer 아키텍쳐를 설계하였습니다.
UFormer의 핵심적인 구조로 Locally-enhanced Window (LeWin) Transformer Block에서는 전역 self-attention 대신 non-overlapping window 기반 self-attention을 사용해 고해상도에서 계산량을 크게 줄이고 계층 구조에서 해상도가 낮아질수록 상대적으로 더 넓은 문맥도 포착할 수 있게 만들고자 하였습니다. 또한, 이미지 복원에는 지역 문맥도 중요하다는 점을 반영해 Transformer Block의 FFN에 Depth-wise convolution을 추가하여 로컬 정보를 더 잘 담도록 하였습니다. 또한 Learnable Restoration Modulator에서는 다양한 열화등에 유연하게 대응하기 위해 디코더 여러 계층에서 특징을 조정하는 멀티 스케일 spatial bias 형태를 가벼운 모듈을 추가합니다. 구현 방식은 매우 단순하게 윈도우 크기의 learnable tensor를 각 블록 특징에 더하는 방식으로 파라미터 및 연산 증가를 최소화하면서 디테일 복원을 돕는다고 설명합니다.
그림 1은 기존 모델들과 UFormer 사이의 efficiency graph를 보여주고 있습니다. 결과에서 볼 수 있는 것과 같이 UFormer는 기존의 이미지 복원 모델들보다 훨씬 적은 연산량에 높은 PSNR를 보여주고 있습니다.
Method
1) Overall Pipeline
그림 2는 UFormer의 전체적인 구조를 보여주고 있습니다. 그림에서 볼 수 있다싶이 UFormer는 U-Shaped로 계층적으로 구성된 인코더-디코더 구조 사이에 skip connection이 추가되어 멀티스케일 특징을 효율적으로 전달하고 있습니다. 기본적으로 구조는 UNet에 Transformer block를 추가하였기 때문에 내부적으로는 Transformer의 전처리 과정과 매우 유사합니다. 초기 특징을 추출하기 위해 열화된 이미지 $\mathbf{I} \in \mathbb{R}^{3 \times H \times W}$가 주어지면 먼저 $3 \times 3$ Convolution에 LeakyReLU를 통해 저수준 특정을 추출하여 $\mathbf{X}_{0} \in \mathbb{R}^{C \times H \times W}$를 얻습니다.
다음 단계에서 $\mathbf{X}_{0}$는 LeWin Transformer Block 여러 개와 다운 샘플링 계층 1개로 구성된 $K$개의 인코더 스테이지를 통과하게 됩니다. 여기서 LeWin 블록은 window 기반 self-attention을 사용해 고해상도에서의 계산량을 줄이면서 장거리 의존성을 학습하도록 설계되어 있습니다. 다운 샘플링은 flatten된 feature map을 2D로 변환한 뒤 stride 2에 $4 \times 4$ convolution으로 공간 해상도를 절반으로 줄이고 채널 수를 두 배로 증가시킵니다. 예를 들어 $l$번째 인코더 스테이지의 출력은 $\mathbf{X}_{l} \in \mathbb{R}^{2^{l}C \times \frac{H}{2^{l}} \times \frac{W}{2^{l}}}$이 됩니다.
그리고 bottleneck에서는 LeWin Transformer Block을 통해 처리가 더 진행됩니다. 이때, 계층구조 상 이 지점의 특징맵은 해상도가 낮아 더 긴 의존성을 잘 포착할 수 있게 됩니다.
그리고 디코딩 단계에서는 $K$단계의 업샘플링과 skip connection feature와의 결합 그리고 LeWin Transformer Block이 순차적으로 구성됩니다. 이때, 업샘플링은 stride 2의 $2 \times 2$의 Transposed Convolution을 사용해 공간 해상도를 두 배로 늘리고 채널수는 절반으로 줄이게 됩니다. 그리고 업샘플된 특징맵과 인코더에서 넘어온 특징을 concatentation하여 LeWin Transformer Block들이 이를 입력으로 받아 최종 복원에 필요한 표현을 학습합니다.
최종적으로 $K$개의 디코더 스테이지 이후 특징을 다시 2D 특징으로 변환하여 $3 \times 3$ convolution을 적용해 residual 영상 $\mathbf{R} \in \mathbb{R}^{3 \times H \times W}$를 예측합니다. 최종 복원 결과는 residual을 입력에 더하는 방식으로 $\mathbf{I}^{'} = \mathbf{I} + \mathbf{R}$로 얻습니다.
이때, UFormer는 $\mathbf{I}^{'}$와 $\hat{\mathbf{I}}$ 사이의 Charbonnier Loss로 학습합니다.
$$\mathcal{L}(\mathbf{I}^{'}, \hat{\mathbf{I}}) = \sqrt{||\mathbf{I}^{'} - \hat{\mathbf{I}}||^{2} + \epsilon^{2}}$$
여기서 본 논문에서는 $\epsilon = 10^{-3}$을 사용하여 모든 실험에서 동일하게 사용했다고 명시하였습니다.
2) LeWin Transformer Block
이미지 복원에서 Transformer Block을 그대로 사용할 때 문제점은 크게 두 가지로 볼 수 있습니다.
- 전통적인 Transformer는 모든 토큰 간 global self-attention을 계산해 토큰 수에 대해 연산량이 quadratic 하게 증가하므로 고해상도 feature map에서 쓰기 부적합
- 이미지 복원은 주변 픽셀 정보 (지역적 정보)가 중요한데 Transformer를 사용하는 기존 연구들에서는 local dependency 포착이 약할 수 있음을 지적
본 논문에서는 이러한 문제를 해결하기 위해 Locally-enhanced Window (LeWin) Transformer Block을 제안합니다. LeWin Transformer Block은 non-overlapping window-based self-attention을 활용하여 계산량을 줄이면서도 global dependency를 효율적으로 학습합니다. 또한, Locally-enhanced FFN (LeFFN)에서는 depth-wise convolution을 추가하여 로컬 문맥을 강화합니다. 이러한 새로운 Transformer Block의 전체 연산과정은 다음과 같이 Pre-LN과 Residual의 형태로 구성되어 있습니다.
$$\begin{cases} \mathbf{X}^{'}_{l} &= \text{W-MSA}(\text{LN}(\mathbf{X}_{l - 1})) + \mathbf{X}_{l - 1} \\ \mathbf{X}_{l} &= \text{LeFF}(\text{LN}(\mathbf{X}^{'}_{l})) + \mathbf{X}^{'}_{l} \end{cases}$$
2-1) Window-based Multi-head Self-Attention (W-MSA)
본 논문에서의 핵심은 global self-attention 대신에 feature map $\mathbf{X} \in \mathbb{R}^{C \times H \times W}$를 $M \times M$ 크기의 non-overlapping window로 나눈 뒤 각 window 내부에서만 self-attention을 수행합니다. 각 window $i$는 flatten+transpose되어 $\mathbf{X}_{i} \in \mathbb{R}^{M^{2} \times C}$의 토큰 시퀀스가 만들어집니다. 헤드 수가 $k$ 그리고 헤드 차원이 $d_{k} = \frac{C}{k}$일 때 각 윈도우에서의 multi-head attention은 다음과 같이 정의됩니다.
$$\mathbf{Y}^{(h)}_{i} = \text{Attention} (\mathbf{X}_{i}\mathbf{W}_{Q}^{(h)}, \mathbf{X}_{i}\mathbf{W}_{K}^{(h)}, \mathbf{X}_{i}\mathbf{W}_{V}^{(h)})$$
그리고 모든 head에서의 출력을 concatenation하여 선형 투영을 수행합니다. 또한, relative position을 반영하기 위해 기존의 Swin Transformer에서 활용하던 relative position bias $\mathbf{B}$를 어텐션 로짓에 다음과 같이 더해줍니다.
$$\text{Attention} (Q, K, V) = \text{Softmax} \left( \frac{QK^{\top}}{\sqrt{d_{k}}} + \mathbf{B} \right)V$$
복잡도 관점에서 global attention의 $\mathcal{O}(H^{2}W^{2}C)$가 window attention에서는 $\mathcal{O}(M^{2}HWC)$로 감소하게 됩니다. 또한, UFormer는 계층적인 구조이기 때문에 해상도가 감소할 수록 동일한 윈도우가 더 넓은 receptive field를 커버해 장거리 의존성 학습에도 좋은 영향을 끼친다고 합니다.
2-2) Locally-enhanced FFN (LeFFN)
그림 3은 본 논문에서 제안하는 LeFFN입니다. 이 방식 역시 매우 단순합니다. 토큰 별로 $1 \times 1$ convolution을 통해 채널을 확장한 뒤 토큰들을 2D feature map으로 reshape하여 $3 \times 3$ depth-wise convolution으로 로컬 정보를 얻습니다. 그리고 다시 토큰으로 flatten을 수행한 뒤 $1 \times 1$ convolution을 통해 채널을 다시 축소시키게 되죠. 이때, 각 활성화 함수는 GELU를 사용하게 됩니다.
3) Multi-Scale Restoration Modulator
서로 다른 열화는 교란 패턴이 제각각이기 때문에 UFormer가 다양한 복원 과제에서 디테일을 더 잘 살리는 가벼운 Multi-Scale Restoration Modulator를 추가로 제안합니다. 이 모듈을 특징을 보정 (calibration)하여 더 많은 디테일 복원을 유도하는 것이 목적입니다.
모듈레이터는 UFormer 디코더의 서로 다른 해상도를 가지는 여러 stage에 걸쳐 여러 개가 삽입됩니다. 즉, 디코더가 복원 과정에서 점진적으로 해상도를 올리며 세부적인 디테일을 구성할 때 가 스케일에서 특징을 조정할 수 있게 도와주는 것이죠.
각 LeWin Transformer Block에서 모듈레이터는 모두 학습 가능한 텐서로 정의되며 윈도우 크기와 동일하게 적용됩니다. 이 텐션의 self-attention을 계산하기 전 모든 non-overlapping window에 대해 공유 bias처럼 동일하게 더해집니다.
정성적으로 모듈레이터를 넣으면 더 많은 모션블러 및 노이즈 패턴이 제거되고 더 깨끗한 복원 결과가 나온다고 설명합니다. 저자들은 이에 대해서 디코더 각 스케일에 모듈레이터를 넣는 것은 특징 맵을 더 유연하게 조정해 디테일 복원 성능을 올린다고 언급하며 이는 StyleGAN이 멀티스케일 노이즈를 특징에 더해 변화를 유도하는 방식과 유사하다고 언급하고 있습니다.
Experiment Results
1) Real Noise Removal
2) Motion Blur Removal
3) Defocus Blur Removal
4) Real Rain Removal
5) Ablation Study
5-1) Transformer vs. convolution & Hierarchical structure vs. single scale
5-2) Where to enhance locality?
5-3) Effect of the multi-scale restoration modulator
