HAT(Hybrid Attention Transformer) 논문 리뷰

https://arxiv.org/pdf/2205.04437.pdf

HAT?
Hybrid Attention Transformer

ABSTRACT

기존의 트랜스포머 기반 방법은 super resolution에서 제한된 공간 범위만 활용할 수 있었음. Transformer의 잠재력이 아직 완전히 활용되지 않았음을 의미

채널 어텐션과 셀프 어텐션 체계를 결합하여 상호 보완적인 HAT을 제안

Cross-window information을 더 잘 집계하기 위해 오버랩되는 cross-attention module을 도입하여 인접한 window 간의 상호작용을 향상

train 단계에서는 추가 개선을 위해 동일한 task의 Pre-train 전략을 제안

논문 시점 기준 최근 방법보다 1dB이상 성능 향상

1. INTRODUCTION

SR(super resolution)?

• 컴퓨터 비전 및 이미지 처리에서 오래전부터 주제로 논의되어왔음
• 저해상도 입력에서 고해상도 이미지로 재구성하는 것이 목표

Transformer?

• 자연어 처리에서의 성공을 통해 컴퓨터 비전에 활발히 적용되고 있음

• 특히 새로 설계된 네트워크인 SwinIR은 이 분야에서 획기적인 개선

• 성공은 했지만 왜 CNN보다 transformer가 나은지는 미스터리…

• 대충 이런 종류의 네트워크가 self-attention 메커니즘의 이점을 얻고 장거리의 정보를 활용할 수 있기 때문이라고 추측됨

• 하지만 저자들은 SwinIR에서 재구성에 사용된 관련 정보의 범위를 조사하기 위해 속성 분석 방법인 LAM을 사용함

• SwinIR이 SR에서 더 높은 성능을 얻더라도, 활용하는 정보의 범위가 제한되어 있어 경우에 따라 RCAN(CNN-BASED model 중 하나)보다도 안 좋은 결과를 나타내기도 함

• 이러한 현상은 Transformer가 로컬 정보를 모델링하는 능력이 강하지만 활용되는 정보의 범위를 확장해야 더 좋다는 것을 보여줌

• 위에서 언급한 문제를 해결하고 SR Transfomer의 잠재력을 극대화하기 위해 Hybrid Attention Transformer를 제안

• HAT은 채널 어텐션과 셀프 어텐션 체계를 결합

  • 채널 어텐션 : 글로벌 정보 사용 능력이 장점
  • 셀프 어텐션 : 강력한 representation ability

• Cross window 정보를 더 잘 집계하기 위해 overlapping cross-attention module을 도입

• 위의 설계 방식은 최근의 방법보다 0.3~ 1.2dB 성능 향상

• 전반적으로 3가지 핵심을 기여

  • 더 많은 입력 정보를 활용하기 위해 Transformer에 채널 어텐션을 도입
  • Cross window information을 더 잘 집계하기 위해 Overlap 모듈 제안
  • 잠재력을 극대화하기 위해 동일한 task에 대한 pretrain 전략을 사용

2. RELATED WORK

   1. Deep Networks for Image SR

• SRCNN이 이미지 SR작업에 CNN을 처음 도입한 후로 기존의 방법보다 우수한 성능을 얻었기 때문에, 다양한 네트워크가 제안되었음

  • 모델 표현 능력을 향상하기 위해 Residual Block, Dense Block과 같은 정교한 컨볼루션 모듈 설계를 적용하는 방법

  • 재귀 신경망(recursive neural network)이나 그래프 신경망(graph neural network)과 같은 또 다른 프레임 워크를 사용하기도 함

  • 더 현실적인 결과들을 생성하기 위해 적대적 학습 방법을 도입하기도 함

  • 최근, Transformer-based 네트워크가 제안되고, SR 작업의 최신기술을 지속적으로 업데이트하여 Transformer의 강력한 표현 능력을 보여줌

  • SR 작동 메커니즘을 더 잘 이해하기 위한 SR Network 분석/해석 작업이 제안됨 [14, 35, 59, 22]

  • LAM[14]은 최종 성능에 가장 많이 기여하는 입력 픽셀을 탐색하기 위해 적분 그래디언트 방법을 채택

  • DDR[35]은 deep feature dimensionality reduction(심층 피쳐 차원 축소?) 및 시각화를 기반으로 한 SR 네트워크의 심층 의미론적 표현을 보여줌

  • FAIG[59]는 블라인드 SR에서 특정 품질 저하에 대한 판별 필터를 찾기 위해 제안됨

  • [22]는 Dropout이 실제 SR 네트워크에 대한 Co-adapting을 방지하는 데 도움을 줄 수 있음을 보여주기 위해 channel saliency(채널 돌출?)맵을 소개

  • 이번 연구에선 LAM[14]을 이용하여 SR 네트워크의 동작을 분석하고 이해함

2. Vision Transformer
- 최근, Transformer[50]는 자연어 처리 분야에서 성공으로 컴퓨터 비전 커뮤니티의 주목을 받고 있음
- 이미지 분류[36, 13, 27, 44, 49], 객체 감지[34, 48, 36, 4, 6], 세분화[55, 51, 16, 2] 등 ViT가 장거리 종속성 모델링에서 높은 성능을 보임
- [5]는 특히 ViT 스타일의 네트워크를 개발하고 이미지 처리를 위한 multi-task pre-training을 도입
- SwinIR[29]는 [36]을 기반으로 이미지 복원 Transformer를 제안함
- [3, 28]은 비디오 복원에 Transformer 기반 네트워크를 소개함

3. METHOD

   1. Motivation

• Swin Transformer[36]는 SR[29]에서 뛰어난 성능을 입증함

• 작동 메커니즘을 밝히기 위해 SR용으로 설계된 LAM[14]이라는 진단 도구를 사용

• LAM을 사용하면 선택 영역에 가장 많이 기여하는 입력 픽셀을 알 수 있음

• Transformer 기반의 방법인 SwinIR의 경우 CNN보다 더 좁은 범위를 나타냄(그림 2 (a))

  • 더 많은 정보를 활용할 수록 더 좋은 성능이 나타나는 것은 당연한 소리인데 이건 무엇을 의미하는가?
    • SwinIR이 CNN보다 훨씬 강한 매핑 능력을 가지고 있어 성능을 달성하는 데 더 적은 정보를 사용함
    • SwinIR이 더 많은 입력 픽셀을 활용할 수 있다면 성능이 더 좋아질 수 있음
  • 채널 어텐션 체계는 더 많은 픽셀을 볼 수 있도록 도와줄 수 있을 것임

• 그림 2 (b)를 보면 SwinIR의 중간 피쳐 맵에서 명확한 blocking artifacts를 관찰할 수 있음

  • 이는 window 파티션 메커니즘에 의해 발생
  • shifted window 메커니즘이 cross-window connection을 구축하는 데 비효율적임을 의미함

• 위의 두 점을 기반으로 Transformer 기반 모델에서 채널 어텐션을 조사하고, 윈도우 기반 SR Transformer에 대한 cross-window Information을 더 잘 얻기 위해 cross-window based SR Transformer를 제안함

  2. Network Architecture

• The Overall Structure

  • 그림 3과 같이 3가지(얕은 특징추출, 심층 특징 추출, 이미지 재구성)로 구성됨
  • 저해상도 입력에 대해 하나의 컨볼루션 레이어를 사용하여 얕은 피쳐를 추출
  • 전체 구조 : Shallow Feature Extraction > Deep Feature Extraction > (Element-wise sum) > Image Reconstruction
    • Deep Feature Extraction : RHAG > RHAG > RHAG > Conv

• RHAG(Residual Hybrid Attention Group)?

  • . . . > HAB > HAB > OCAB > Conv > (Element-wise sum) > . . .
  • M개의 HAB(Hybrid Attention Block)과 OCAB(overlapping Cross-Attention Block)와 3 * 3 컨볼루션 레이어를 포함
  • 일련의 HAB 매핑 후 OCAB를 삽입해 window기반 셀프 어텐션에 대한 수용 필드를 확대하고 corss-window information을 더 잘 모음
  • RHAG의 끝에서는 [29]에 근거해 컨볼루션 레이어를 넣음
  • 학습 과정을 안정화하기 위해 Residual connection도 추가

• HAB(Hybrid Attention Block)

  • . . > LayerNorm > CAB / (S) W-MSA > (Element-wise sum) > LayerNorm > MLP > (Element-wise sum) > . . .
  • 그림 2 (a)와 같이 채널 어텐션 가중치를 계산하는데 전역 정보가 포함되기 때문에 채널 어텐션이 채택될 때 더 많은 픽셀이 활성화 됨
  • 게다가 많은 작업에서 컨볼루션이 Transformer가 더 나은 시각적 표현을 얻거나 더 쉬운 최적화를 달성하는 데 도움이 될 수 있음을 보여줘왔음 [56, 58, 60, 25, 68]
  • 따라서 채널 어텐션 기반의 컨볼루션 블록을 표준 Transformer 블록에 통합하여 네트워크의 표현 능력을 더욱 향상시킴
  • 그림 3에서 볼 수 있듯이 채널 어텐션 블록(CAB)은 윈도우 기반의 멀티 헤드 어텐션(W-MSA) 모듈과 병렬로 첫 번째 LayerNorm(LN) 레이어 다음 표준 Swin Transformer 블록에 삽입됨
  • SW-MSA(shifted window-based self-attention)는 [36, 29]와 유사한 HAB의 연속으로 만들어짐(?)
  • 최적화 및 시각적 표현에 대한 CAB 및 MSA의 충돌을 피하기 위해 작은 상수인 a가 CAB의 출력에 곱해짐
  • 주어진 입력 feature X에 대해 HAB의 전체 프로세스가 계산됨
  • 특히 각 픽셀을 임베딩을 위한 토큰으로 취급함
  • [29]에 따라 패치 임베딩을 위해 패치 크기를 1로 설정함

• CAB(Channel Attention Block)

  • . . . > Conv > Activation function > Conv > Global pooling > Conv > Activation function > Conv > (Sigmoid function) > (Element-wise product) > . . .

• OCAB(Overlapping Cross-Attention Block)

  • . . . > LayerNorm > OCA > (Element-wise sum) > LayerNorm > MLP > (Element-wise sum) > . . .
  • Cross-window 연결을 직접 설정하고 window self-attention에 대한 대표 기능을 향상시킴
  • OCAB은 표준 Swin Transformer 블록[36]과 유사한 OCA(Overlapping Cross Attention) 레이어와 MLP레이어로 구성됨
  • 하지만 그림 4에서와 같이 투영된 feature들을 분할하기 위해 서로 다른 window 크기를 사용함
  • 표준 윈도우 분할 영역은 커널의 크기와 보폭이 모두 윈도우 크기 M인 슬라이딩 분할 영역으로 간주할 수 있음
  • 반대로 오버랩되는 윈도우 분할 영역은 커널 크기가 있는 슬라이딩 분할 영역으로 볼 수 있음.

3. Pre-training on ImageNet

• 최근 연구[5, 26]에선 Pre-train 이 low-level tasks에서 중요한 역할을 한다는 것을 보여줌
• IPT[5]는 denoising, deraining, super-resolution 등과 같은 다양한 저수준의 작업을 강조함 ???
• EDT[26]는 사전 훈련을 수행하기 위해 특정 작업의 다양한 degradation 을 활용함 ????
• 이 논문에서는 동일한 작업을 하되, 더 큰 규모의 데이터셋(ImageNet[8])에 대해 pre-train을 직접 수행함
• 예를 들어, 4배 SR을 하는경우 ImageNet에서 4배 모델을 교육한 후 DF2K와 같은 특정 데이터셋에서 미세 조정을 함.

4. Experiments

   1. Experimental Setup

• HAT 구조에 의해 RHAG와 HAB는 모두 6으로 설정됨

• 전체 네트워크 채널은 180으로 설정

• 어텐션 헤드와 윈도우 크기는 (S)W-MSA와 OCA에 대해 각각 6과 16으로 설정

  2. Effects of different window sizes

• 3.1에서 논의한 바와 같이 SR을 할 때 더 많은 입력 픽셀을 활성화하면 더 나은 성능을 얻음

• [26]에서 저자는 다양한 window 크기의 효과를 조사했음

• 테이블 1과 같이 window 크기가 16인 모델이 Urban100에서 더 좋은 성능을 보임

• 그림 5에서와 같이 질적으로도 크기가 16인 모델이 8인 모델보다 훨씬 더 많은 입력 픽셀을 사용함

• 이 결론에 따라 크기가 16인 모델을 기본으로 사용함

  3. Ablation Study

• Effectiveness of OCAB and CAB.

  • OCAB와 CAB의 효과를 입증하기 위한 실험을 수행함.
  • OCAB와 CAB모두 0.1dB의 성능 향상을 가져옴
  • 둘을 모두 활용하면 0.16dB의 성능 향상

• Effects of the overlapping size.

  • 0에서 0.75까지 중첩 크기의 효과를 탐색했음
  • 0은 표준 transformer 블록을 의미.
  • 0.5에서 가장 잘 수행됨
  • 0.25나 0.75로 설정하면 모델의 성능이 확연하게 향상되지 않거나 성능이 오히려 떨어짐
  • 부적절하게 겹쳐지면 오히려 도움이 안되는 것을 보여줌

• Effects of different designs of CAB.

  • 다양한 CAB 설계의 효과를 알아보기 위해 실험을 수행함
  • 컨볼루션과 채널 어텐션의 영향을 조사함
  • 테이블 4와 같이 Depth-wise Convolution을 사용하면 심각한 성능 저하가 발생하며, CAB의 대표성을 약화시키는 것을 의미함
  • 테이블 5와 같이 CAB의 가중치 계수 a가 0.01이 가장 좋은 성능을 보임.
  • CAB와 self-attention이 최적화에서 잠재적인 문제가 있을 수 있음을 나타냄
  4. Comparison with State-of-the-Art Methods

• Quantitative results.

  • 테이블 6에서 최근의 방법과 이번 접근 방식의 정량적 비교를 보여줌
  • 다른 모든 벤치마크 데이터세트에서 다른 방법들보다 훨씬 뛰어난 성능을 보임

• Visual comparison.

  • 그림 7처럼 눈으로 보기에도 더 좋았음
  5. Effectiveness of the same-task pre-training

• 테이블 6에서와 같이 사전훈련을 해주는게 훨씬 뛰어남

5. CONCLUSION

• 본 논문에서는 SR을 위한 HAT을 제안함
• 많은 픽셀을 활성화하기 위해 채널 어텐션과 셀프 어텐션을 결합함
• Cross window information을 더 잘 모으기 위해 창 크기를 설정하고, 겹치게 하는 방식을 제안함
• 또 잠재력을 활성화하기 위해 사전학습을 도입했음

[ 내용 수정 및 그림 추가 예정 ]