[논문 리뷰] HAT(Hybrid Attention Transformer) 논문리뷰

Activating More Pixels in Image Super-Resolution Transformer

Xiangyu Chen, Xintao Wang, Jiantao Zhou, Yu Qiao, Chao Dong

https://arxiv.org/abs/2205.04437

Activating More Pixels in Image Super-Resolution Transformer

 

 

Abstract

transformer 기반의 image super resolution이 low level visison task에서는 좋은 성능을 보였으나, 이는 제한된 범위의 입력 정보만을 사용한다. 이것은 기존의 네트워크에 비해 활용도가 낮아지는 것을 의미한다. 그렇기 때문에 HAT는 더 많은 픽셀을 활성화하기 위해 Channel Attention과 Self Attention을 결합하였다. 더불어 cross-window information을 잘 통합하기 위해 overlapping cross-attention 모듈을 도입하여 window feature간의 상호작용을 향상시키도록 했다. 성능을 개선시키기 위해 훈련 단계에서는 same-task pre-training 전략을 사용하였다.

 

1. Introduction

Single image super resolution은 low resolution이 주어졌을 때 high resolution을 만드는 것을 목표로 한다. 최근에는 transformer를 사용하고 있는데, high level task에 대해 빠른 진전을 이룬 이후 super resolution뿐만 아니라 low-level vision에서도 개발되고 있다.

그러나 Transformer의 성능이 CNN보다 좋은 이유는 여전히 미지수로 남아 있다. Self attention을 통해 long range정보를 활용할 수는 있으나, 그림 2(a)를 확인하면 SwinIR이 CNN보다 더 적은 입력 픽셀을 사용했다는 점을 알 수 있다. 이로부터 transformer의 local information을 모델링하는 능력은 강해도 활용되는 정보의 범위는 확장할 필요가 있음을 알 수 있다.

Contribution

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

 

2. Methodology

2.1 Motivation

Swin Transformer는 CNN보다 Single image super resolution에서 좋은 성능을 낸다는 것을 입증하였다. LAM진단 도구를 사용해서 SwinIR의 작동 매커니즘을 살펴본 결과, 다른 CNN에 비해 더 좋은 성능을 보이지만 더 적은 범위를 표시하고 있는 것을 확인할 수 있었다. 여기서 붉은 점은 이미지를 재구성하는 데에 유용한 픽셀값 정보들을 표시한 것이다. SwinIR은 더 적은 정보량으로도 더 좋은 성능을 발휘하는 mapping capability를 갖고 있다는 통찰력을 주었고, 더 많은 픽셀 정보를 사용하면 더 좋은 성능을 낼 수 있을 것이라는 성능 개선의 여지를 보였다.

그림 2(b)에서는 SwinIR의 중간 feature map에서 나타나는 blocking artifacts를 관찰할 수 있다. 이러한 artifacts는 Shifted window mechanism이 cross-window connection을 구축하는 데에 도움을 주지 못하기 때문에 본 논문에서는 transformer-based 모델의 channel attention을 연구하고 cross-window 정보를 더 효과적으로 집계할 overlapping cross-attention 모듈을 제시했다.

2.2 Network Architecture

2.2.1 the overall structure

 

전체 네트워크는 그림 3과 같이 Shallow Feature Extraction, Deep Feature Extraction, Image Reconstruction으로 구성된다. Shallow Feature Extraction에서는 3x3 conv로 shallow feature를 추출한다. Shallow feature extraction은 각 픽셀 토큰에 대한 고차원 임베딩을 달성하는 동시에 저차원에서 고차원 공간으로 input을 mapping할 수 있다.

Deep Feature Extraction은 N개의 RHAG(Residual Hybrid Attention Groups)와 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를 사용하였다.

 

2.2.2 Residual Hybrid Attention Group (RHAG)

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

 

2.2.3 Hybrid Attention Block (HAB)

channel attention weights계산 과정에 global information이 포함되기 때문에, channel attention을 고를 때 더 많은 픽셀이 활성화 된다. 네트워크의 representation ability를 강화시키기 위해 channel attention-based convolution block을 Transformer block에 통합한다. (원래는 어떻게 되어 있었다는 건지 조금 의문??잘이해안됨) 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의 간격으로 채택된다. 

 

 

2.2.4. 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간의 정보를 더 잘 집계할 수 있도록 한다.