[논문 리뷰] More than Encoder: Introducing Transformer Decoder to Upsample

Yijiang Li, Wentian Cai외 3명|arXiv (Cornell University)|2021. 06. 20.

Radiomics and Machine Learning in Medical Imaging참고 문헌 52인용 수 6

한 줄 요약

이 논문은 의료 영상 분할에서 특징 재구성 향상을 위해 디코더를 활용하는 새로운 트랜스포머 기반 업샘플링 방법인 Window Attention Upsample (WAU)을 제안한다. 다운샘플링 경로의 특징에 국소 윈도우형 자기주의를 적용함으로써 공간적 세부 정보와 의미 일관성을 향상시켜, Synapse, MSD Brain, ACDC 데이터셋에서 다양한 아키텍처에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Medical image segmentation methods downsample images for feature extraction and then upsample them to restore resolution for pixel-level predictions. In such a schema, upsample technique is vital in restoring information for better performance. However, existing upsample techniques leverage little information from downsampling paths. The local and detailed feature from the shallower layer such as boundary and tissue texture is particularly more important in medical segmentation compared with natural image segmentation. To this end, we propose a novel upsample approach for medical image segmentation, Window Attention Upsample (WAU), which upsamples features conditioned on local and detailed features from downsampling path in local windows by introducing attention decoders of Transformer. WAU could serve as a general upsample method and be incorporated into any segmentation model that possesses lateral connections. We first propose the Attention Upsample which consists of Attention Decoder (AD) and bilinear upsample. AD leverages pixel-level attention to model long-range dependency and global information for a better upsample. Bilinear upsample is introduced as the residual connection to complement the upsampled features. Moreover, considering the extensive memory and computation cost of pixel-level attention, we further design a window attention scheme to restrict attention computation in local windows instead of the global range. We evaluate our method (WAU) on classic U-Net structure with lateral connections and achieve state-of-the-art performance on Synapse multi-organ segmentation, Medical Segmentation Decathlon (MSD) Brain, and Automatic Cardiac Diagnosis Challenge (ACDC) datasets. We also validate the effectiveness of our method on multiple classic architectures and achieve consistent improvement.

연구 동기 및 목표

• 기존 업샘플링 기법이 다운샘플링 경로의 국소적이고 세밀한 특징을 충분히 활용하지 못하는 한계를 해결하기 위해.
• 트랜스포머 디코더의 특징 확대 잠재력을 탐색하여, 일반적으로 시퀀스 생성에 사용되는 것 외의 응용 가능성을 탐색하기 위해.
• 의료 영상에서 중요한 미세한 세부 정보를 유지하는 메모리 및 계산 비용이 효율적인 업샘플링 기법 개발을 위해.
• 스킵 커넥션을 사용하는 모든 분할 모델에 호환 가능한 일반적인 업샘플링 모듈을 만들기 위해.
• 2D 및 3D U-Net, FCN, DeepLabV3를 포함한 다양한 아키텍처에서 일관된 성능 향상을 입증하기 위해.

제안 방법

• 트랜스포머 기반의 어텐션 디코더(AD)와 이중선형 업샘플링을 잔차 연결로 조합한 이중 스트림 모듈인 어텐션 업샘플링(AU)을 제안한다.
• 자기주의 계산을 국소적 공간 윈도우로 제한함으로써 계산 비용을 줄이고 장거리 모델링 능력을 유지하는 윈도우형 어텐션 디코더(WAD)를 설계한다.
• 각 윈도우 내에서 효율적인 특징 투영을 위해 딥라이즈드 분리형 컨볼루션을 사용하여, 작은 의료 데이터셋에서의 성능 향상을 도모한다.
• 국소적 특징 윈도우에서 쿼리, 키, 밸류 행렬을 계산하기 위해 윈도우형 자기주의를 적용하여 맥락 인식 업샘플링을 가능하게 한다.
• U-Net 유사 아키텍처의 스킵 커넥션 경로에 WAU를 통합하여 표준 분할 헤드와 함께 엔드 투 엔드 학습이 가능하게 한다.
• 이중선형 업샘플링이 어텐션 기반 정밀 조정을 보완하는 잔차 학습 체계를 도입하여 학습 안정성 향상과 특징 품질 향상을 달성한다.

실험 결과

연구 질문

• RQ1 트랜스포머의 디코더가 의료 영상 분할에서 특징 확대에 효과적으로 재사용될 수 있는가?
• RQ2 다운샘플링 경로의 국소적 세밀한 특징을 통합할 경우, 표준의 역컨볼루션 또는 이중선형 업샘플링에 비해 업샘플링 품질은 어떻게 향상되는가?
• RQ3 윈도우형 어텐션 메커니즘에서 국소적 세부 정보와 장거리 맥락 모델링 간의 균형을 고려할 때 최적의 윈도우 크기는 무엇인가?
• RQ4 제안된 WAU 방법은 2D 및 3D 모델을 포함한 다양한 분할 아키텍처로 어떻게 일반화되는가?
• RQ5 업샘플링에서 전역 어텐션 대비 국소 윈도우형 어텐션 사용 시 계산 비용과 성능 향상 간의 상충 관계는 어떠한가?

주요 결과

• WAU는 Synapse 다기관 분할 데이터셋에서 최신 기술 수준의 성능을 달성하여 기존 방법을 능가한다.
• 의료 분할 디카테론(MSD) 뇌 데이터셋에서 기준 ResUNet 대비 평균 DSC가 1.47 포인트 향상되었다.
• 3D U-Net에서 DSC가 0.36 향상되어 3D 의료 볼륨 분할에 대한 강력한 일반화 능력을 입증했다.
• WAU의 최적 윈도우 크기는 4이며, 이보다 크거나 작아지면 집중력 상실 또는 충분한 맥락 부족으로 성능 저하가 발생한다.
• 딥라이즈드 분리형 컨볼루션은 특히 파rameter가 제한된 작은 의료 데이터셋에서 다른 컨볼루션 유형 대비 최고의 성능을 보였다.
• 시각화 결과 WAU의 어텐션 가중치가 관련 해부학적 구조에 공간적으로 집중되어 있음을 확인하여 특징의 구분력 향상과 병변의 가시성 향상을 확인했다.