[논문리뷰] Activating More Pixels in Image Super-Resolution Transformer(HAT)

"Activating More Pixel in Image Super-Resolution Transformer"

paper: https://arxiv.org/abs/2205.04437

GitHub: https://github.com/XPixelGroup/HAT

Activating More Pixels in Image Super-Resolution Transformer

---

[Abstract]

 Transformer 기반의 image super-resolution 같은 low-level vision task는 좋은 성능을 보였다. 하지만, Transformer 기반의 네트워크는 attribution 분석을 통해 제한된 공간 범위의 입력 정보만 사용한다. 이러한 것은 transformer의 잠재력이 여전히 기존 네트워크에 비해 완전히 활용되지 않는다는것을 의미한다. 그래서 더 많은 입력 픽셀을 활성화 하기 위해 HAT(Hybrid Attention Transformer)를 제안한다. HAT는 CA(Channel Attention)과 self-attention을 결합해 서로 상호보완적인 장점을 이용해서 만들었다. 게다가 Cross-window information을 더 잘 통합하기 위해 이웃한 window feature 간의상호작용을 향상 시키기 위해 overlapping cross-attention 모듈을 도입한다. 그리고 훈련단계에서는 성능을 개선시키기 위해 same-task pre-training 전략을 사용한다. 실험 결과에서는 제안한 모듈이 효율적인것을 보여주며 다른 SotA 방법보다 1dB 더 높았다.

---

[Intorduction]

 SR(single image Super-Resolution)은 computer vision과 image processing에서의 고전적인 문제이며 LR(Low Resolution) 이 주어졌을때 HR(High Resoltuion)을 만드는것이 목표이다. SR은 최근 몇년 동안 CNN을 접목하면서 좋은 성능을 보였고, 최근에는 transformer를 사용하고 있다. High-level vision task에 대해 빠른 진전을 이룬 이후, transformer 기반의 방법은 SR 뿐만 아니라 low-level vision에서도 개발되고 있다. 

  하지만 이러한 성공에도 불구하고, "Transformer가 CNN보다 나은 이유"는 여전히 미스테리로 남아있다. 직관적인 설명은 이러한 종류의 네트워크가 self-attention mechanism으로 부터 이익을 얻고 long-range 정보를 활용할수 있다는 것이다. 하지만 LAM에서 transformer 기반의 SwinIR이 CNN 기반의 방법보다 입력 픽셀을 덜 사용했다(그림 2(a)). 이러한 점에 의해 SwinIR이 RCAN 보다 정량적 결과는 좋을수 있지만, 정성적 결과에서는 더 안좋을수도 있다. 그래서 이러한 현상은 transformer가 local information을 모델링하는 능력이 더 강하지만 활용되는 정보의 범위를 확장할 필요가 있음을 보인다.

  이러한 문제점을 해결하기 위해 HAT(Hybrid Attention Transformer)를 제안한다. HAT는 global information을 이용하는 channel attention과 powerful representative ability를 사용하기 위해 self-attention을 결합한다. 또한, Cross-window information을 더 잘 통합하기 위해 overlapping cross-attention module을 도입하고, same-task pre-training 전략을 사용한다. Same-task pre-training은 multi-related-task pre-training보다 더 나은 결과를 보였다. 그림 1은 제안하는 모델이 SotA보다 0.3dB ~ 1.2dB더 높은 성능을 보인다.

 

Contribution

• 더 많은 input information을 사용하기 위해 channel attention을 도입함
• Cross-window information을 더 잘 통합하기 위해 overlapping cross-attention 모듈을 제안함
• Same-task pre-training 전략을 사용함

[Method]

2.1 Motivation

  Swin Transformer는 SR에서 좋은 성능을 낸다는것을 입증했고, CNN보다 더 나은 성능을 보였다. Transformer의 작동 mechanism을 밝히기 위해 SR을 위해 설계된 attribution 방법인 LAM 진단도구를 사용한다. 그림 2(a)의 붉은 점은 재구성하는데 기여한 유용한 픽셀인다. 직관적으로 더 많은 정보를 사용했을때 더 좋은 성능을 보이지만, transformer 기반의 SwinIR은 RCAN보다 더 적은 범위를 표시한다. 이러한 부분은 상식과 모순되지만, 또한 추가적인 통찰력을 제공해준다. 그림 2(a)의  파란 박스 안의 재구성한 결과는 SwinIR이 RCAN보다 poor 하다. 이러한 점은 CA가 RCAN보다 더 많은 픽셀을 보도록 할 수 있다.

1. SwinIR은 CNN보다 더 강한 mapping capability을 가져 적은 정보량으로도 더 좋은 성능을 발휘함
2. SwinIR은 더 많은 입력 픽셀의 정보를 사용함으로 성능의 개선 여지가 보임

  그림 2(b)에서는 SwinIR의 중간 feature map에서 blocking artifacts가 관찰된다. 이러한 artifacts는 window partition mechanism으로 발생하며, shifted window mechanism이 cross-window connection을 구축하는데 비효율적이라고 한다. 몇몇 high-level vision tasks들은 windows를 따라 connection을 강화시키면 window-based self-attention 방법을 개선할수 있을 것이라고 한다. 그래서 이 두 관점을 기반해 Transforemr-baesd 모델의 channel attention을 연구하고 window-based SR Transformer에 대한 cross-window 정보를 더 효과적으로 집계할 overlapping cross-attention 모듈을 제시한다.

 

 

1. Blocking Artifacts (블로킹 아티팩트):
- 디지털 영상 또는 음성 압축에서 발생하는 현상으로, 데이터를 압축하거나 전송할 때 발생하는 낮은 품질의 영상 또는 음성 부분입니다.
- 주로 JPEG, MPEG 등과 같은 압축 표준에서 발생하는데, 이러한 현상은 압축 과정에서 데이터를 나누어 처리할 때 발생하는 정보 손실로 인해 나타납니다.
- 블로킹 아티팩트는 일반적으로 이미지나 비디오의 경계선 주변에 나타나며, 낮은 해상도로 인해 발생하는 깨진 블록과 같은 형태로 나타날 수 있습니다.
- 이러한 아티팩트는 영상 또는 음성의 시각적 또는 청각적 품질을 저하시키며, 고화질 영상이나 음성을 재생성하는데 있어서 문제가 될 수 있습니다.

2. Window Partition Mechanism (윈도우 분할 메커니즘):
- 데이터를 처리하거나 표시하는 과정에서 사용되는 기술 중 하나입니다.
- 윈도우는 데이터의 특정 부분에 대한 조각이며, 윈도우 분할 메커니즘은 데이터를 이러한 윈도우로 분할하여 처리하거나 표시하는 방법을 의미합니다.
- 이러한 분할은 시간적, 공간적으로 이루어질 수 있습니다. 예를 들어, 시간적인 분할은 시간을 기준으로 데이터를 분할하여 처리하는 것을 의미하고, 공간적인 분할은 공간을 기준으로 데이터를 분할하여 처리하는 것을 의미합니다.
- 윈도우 분할 메커니즘은 데이터를 효율적으로 처리하거나 표시하기 위해 사용되며, 다양한 응용 분야에서 활용됩니다. 예를 들어, 음성 인식에서는 오디오 신호를 시간적인 윈도우로 분할하여 각 윈도우에 대한 특징을 추출하고 처리합니다.

 

2.2 Network Architecture

The Overall Structure.

 

 

  전체 네트워크는 그림 3과 같이 Shallow Feature Extraction, Deep Feature Extractino, Image Reconstruction으로 구성된다. Shallow Feature Extraction은 3x3 conv로 shallow feature를 추출한다. 이때 channel 수는 3에서 96이 된다. Shallow feature extraction은 각 픽셀 토큰에 대한 고차원 임베딩을 달성하는 동시에 저 차원에서 고차원 공간으로 input을 mapping할 수 있다. 또한, 초기 convolutional layer는 더 나은 visual representation을 학습하도록 돕고 안정적인 최적화로 이어지게 할 수 있다.

  Deep Feature Extraction은 N개의 RHAG(Residual Hybrid Attention Groups)와 중간 feature를 점진적으로 처리할 수 있는 1개의 3x3 Convolutional layer로 구성된다. RHAG에서 또한, deep feature 정보를 더 잘 통합하기 위해 RHAG 끝 부분에 3x3 Convolutional layer를 도입한다. 이후 shallow feature와 deep feature를 혼합하기 위해 global residual connection(Element-wise sum)을 수행한다.

  Image Reconstruction은 HR 영상을 재구성하기 위한 모듈로 3x3 Conv, Pixel-shuffle method, 3x3 Conv를 한다. 이때 Pixel-shuffle 방법으로 혼합된 feature를 up-sample 한다. 파라미터를 최적화 하기 위해 L1 loss를 사용한다.

 

Residual Hybrid Attention Group (RHAG)

  그림 3에 나타나듯이 RHAG 모듈은 M개의 HAB(Hybrid Attention Blocks)와 1개의 OCAB(Overlapping Cross-Attention Block) 그리고 3x3 Convolutional layer로 구성된다. HAB들로 mapping 후, OCAB를 삽입해 window-based self-attention에 대한 receptive field를 확장하고 cross-window informatoin을 더 잘 집계하도록 한다. 마지막에 Convolutional layer를 사용하게 되면 Transformer-based network에 있는 conv 연상의 귀납적 편향을 가져올 수 있고, shallow와 deep 한 feature를 통합하기 위한 더 나은 기반을 만들수 있다. 또한 training process를 안정화 하기 위해 residual connection을 추가한다.

 

Hybrid Attention Block (HAB)

  그림 2(a)에 나온것 처럼 channel attention weights를 계산하기 위해 global information이 포함되기 때문에, channel attention을 채택할 때 더 많은 픽셀이 활성화 된다. 게다가 많은 연구는 Convolution이 Transformer가 더 나은 visual representation을 얻거나 더 쉽게 최적화를 달성하는데 도움이 될 수 있다고 한다. 그래서 네트워크의 representation ability를 강화시키기 위해 channel attention-based convolution block을 Transforemr block에 통합한다. 그림 3에 나타나듯이 LN(Layer Norm) layer 이후,  standard Swin Transformer block에 삽입된 CAB(Channel Attention Block)와 W-MSA(Window-based Multi-head Self-Attention)이 있다. SW-MSA(Shifted Window-based Self-Attention)은 Swin Transformer와 비슷하게 연속적인 HABs의 간격으로 채택된다. 최적화 및 시각적 representation에 대한 CAB와 MSA의 connflict 가능성을 피하기 위해 작은 상수 α를 CAB 출력에 곱한다. HAB의 연산은 아래와 같다.

 

 

  Self-attention 모듈은 입력 feature의 크기가 HxWxC일때 처음에 HW/M^2로 분할되어 (local window size가 MxM) 각 window마다 self-attention이 계산된다. Local window feature X_W ∈ R^(M^2xC)의 경우, query, key and value 행렬은 Q, K, V로 linear mapping을 통해 계산된다. Window-based self-attention 수식은 아래와 같다.

 

 

d 는 qurey / key의 차원을 나타낸다. B는 relative position encoding으로 각 위치의 상대적인 위치 정보를 인코딩해 attention 계산에 활용하는 방법이다. B는 Attention is All you Need에서 제시된 방법으로 계산된다. 이를 통해 모델은 입력 시퀀스 내의 각 위치가 다른 위치와 어떤 관계에 있는지를 학습할 수 있다. 또한, 인접한 non-overlapping windows 사이의 연결을 구축하기 위해, shifted window partition 접근법을 사용하고 sifted size를 window size의 절반으로 설정한다.

  CAB는 CA 모듈 앞에 Conv layer와 GELU activation function을 사용한다. Transformer 기반의 구조는 종종 토큰 임베딩을 위해 많은 채널을 필요로 하기 때문에, 일정한 폭의 conv를 직접 사용하는것은 큰 계산 비용을 초래한다. 따라서 channel number를 두 conv layer 사이의 상수 β로 압축한다. 첫번째 conv layer에서는 channel을 β만큼 나누어 압축하고 두번째 conv layer에서는 다시 원본 채널 수로 팽창한다.

 

Overlapping Cross-Attention Block (OCAB)

  OCAB는 Swin Transformer 기반의 모델에서 window간의 연결성을 향상시키기 위해 제안된 모듈이다. 즉, Cross-window connection을 직접적으로 구축하고 window self-attention을 위해 representative 능력을 강화한다. OCAB는 Swin Transformer block 처럼 OCA와 MLP로 구성된다. 하지만 OCA의 경우 그림 4처럼 다른 window size를 사용해 투영된 feature를 분할한다. OCA layer는 input feature의 다른 window size 간의 attention을 계산해 window간의 정보를 더 잘 집계할 수 있도록 한다. 구체적으로, Q, K, V가 input feature X에 포함된다고 할때, Q는 window size가 MxM인 HW/M^2 non-overlapping으로 분할되고, K, V는 크기가 MoxMo인 HW/M^2 overlapping으로 펼쳐진다. 식은 아래와 같다.

 

 

γ는 overlapping size를 조절하는 상수이다. 이 연산을 더 쉽게 이해하기 위해, standard window partition은 kernel size와 stride가 window size M가 동일한 sliding partition으로 고려되어질 수 있다. 대조적으로 overlapping window partition은 kernel size가 Mo와 같은 sliding partition 볼수 있으며, stride는 M과 같다. Overlapping window의 size 일관성을 보장하기 위해 γM 크기의 Zero padding을 사용한다. Attention matrix는 위의 attention식처럼 계산된다. OCA는 q,k,v가 같은 window feature에서 계산되는 WSA와는 다르게, q에 더 유용한 정보를 활용할 수 있는 더 큰 field에서 k/v를 계산한다. MOA(Multi-resolution Overlapped Attention)이 overlapping window partition과 비슷하게 수행하지만 MOA는 window feature를 토큰으로 사용해 global attention을 계산하는 반면 OCA는 픽셀 토큰을 사용해 각 window feature 내부의 cross-attention을 계산하기 때문에 MOA는 OCA과 다르다.