[논문 리뷰] Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation

Reza Azad, Leon Niggemeier외 6명|arXiv (Cornell University)|2023. 08. 31.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 의료 영상 분할을 위한 비전 트랜스포머에서 자기주의를 대체하기 위해 경량이며 효율적인 주의 메커니즘인 변형 가능 큰 커널 주의(D-LKA)를 제안한다. 변형 가능한 샘플링 격자를 사용하는 큰 컨volutional 커널을 통해 D-LKA는 선형 복잡도로 전역적 맥락을 포착하며, 이는 우수한 3D 부피 특징 학습을 가능하게 한다. D-LKA Net은 더 적은 파라미터와 FLOPs로 Synapse, NIH 췌장, 피부 병변 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Medical image segmentation has seen significant improvements with transformer models, which excel in grasping far-reaching contexts and global contextual information. However, the increasing computational demands of these models, proportional to the squared token count, limit their depth and resolution capabilities. Most current methods process D volumetric image data slice-by-slice (called pseudo 3D), missing crucial inter-slice information and thus reducing the model's overall performance. To address these challenges, we introduce the concept of extbf{Deformable Large Kernel Attention (D-LKA Attention)}, a streamlined attention mechanism employing large convolution kernels to fully appreciate volumetric context. This mechanism operates within a receptive field akin to self-attention while sidestepping the computational overhead. Additionally, our proposed attention mechanism benefits from deformable convolutions to flexibly warp the sampling grid, enabling the model to adapt appropriately to diverse data patterns. We designed both 2D and 3D adaptations of the D-LKA Attention, with the latter excelling in cross-depth data understanding. Together, these components shape our novel hierarchical Vision Transformer architecture, the extit{D-LKA Net}. Evaluations of our model against leading methods on popular medical segmentation datasets (Synapse, NIH Pancreas, and Skin lesion) demonstrate its superior performance. Our code implementation is publicly available at the: https://github.com/mindflow-institue/deformableLKA

연구 동기 및 목표

• 의료 영상 분할을 위한 비전 트랜스포머에서 자기주의의 높은 계산 비용을 해결하기 위해.
• 3D 의료 영상 분석에서 2D 슬라이스 기반 처리의 한계를 극복하기 위해 상하좌우 슬라이스 간 특징 통합을 가능하게 하기 위해.
• 로컬 세부 정보를 유지하면서 전역 맥락 모델링을 향상시켜 CNN의 국소 수용 영역 제약을 극복하기 위해.
• 의료 분할 벤치마크에서 높은 성능을 유지하면서도 경량이며 파라미터 효율적인 주의 메커니즘을 개발하기 위해.

제안 방법

• 자기주의의 장거리 모델링을 선형 복잡도로 시뮬레이션하는 큰 컨volutional 커널을 사용하는 변형 가능 큰 커널 주의(D-LKA)를 제안한다.
• 학습 가능한, 미분 가능한 변형 격자를 도입하여 샘플링 위치를 복잡한 해부학적 구조에 맞게 동적으로 왜곡한다.
• 부피적 의료 영상에서의 상하좌우 슬라이스 간 의존성을 포착하기 위해 D-LKA의 3D 버전을 설계한다.
• D-LKA 블록과 스킵 커넥션을 조합하여 특징 정제를 위한 계층적 비전 트랜스포머 아키텍처(D-LKA Net)를 구축한다.
• 로컬 세부 정보와 전역 맥락을 균형 있게 조화시키기 위해 CNN과 트랜스포머 구성 요소를 융합한 하이브리드 인코더-디코더 구조를 사용한다.
• 스킵 커넥션을 통한 다중 척도 특징 집약 전략을 활용하여 공간 해상도를 유지하고 분할 정확도를 향상시킨다.

Figure 1 : Proposed network architecture of the 3D D-LKA model on the left and the 2D D-LKA model on the right.

실험 결과

연구 질문

• RQ1 큰 커널 컨volution 메커니즘이 의료 영상 분할을 위한 비전 트랜스포머에서 자기주의를 효과적으로 대체하면서 계산 비용을 줄일 수 있는가?
• RQ2 큰 커널 주의에서의 변형 가능한 샘플링이 3D 의료 영상에서 비정상적인 형태의 기관에 대한 특징 표현을 어떻게 향상시키는가?
• RQ3 2D 가짜 3D 접근 방식에 비해 3D D-LKA는 상하좌우 슬라이스 간 특징 학습을 어느 정도 향상시키는가?
• RQ4 스킵 커넥션과 변형 가능한 컨볼루션은 D-LKA Net 아키텍처의 최종 분할 성능에 어떤 기여를 하는가?
• RQ5 경량 주의 메커니즘이 의료 분할 벤치마크에서 정확도와 모델 효율성 측면에서 기존의 SOTA 모델을 초월할 수 있는가?

주요 결과

• Synapse 데이터셋에서 D-LKA Net은 DSC 87.49%를 기록하여 파라미터 수가 37.73M 적은 UNETR++(80.59%)과 UNETR(77.42%)를 능가했다.
• NIH 췌장 데이터셋에서 D-LKA Net은 DSC 81.22%를 달성하여 UNETR++보다 0.63% 향상되었고, HD95는 1.04점 감소했다.
• 변형 가능한 컨볼루션을 적용한 3D LKA 버전은 비변형 버전보다 DSC가 0.63% 높아, 적응형 샘플링의 이점이 입증되었다.
• 제거 실험 결과, 스킵 커넥션은 성능 향상에 크게 기여하며, 최상위 수준의 스킵 커넥션만으로도 DSC가 0.42% 향상되었다.
• 췌장 데이터셋에서 D-LKA Net은 SOTA 방법들 중에서 가장 낮은 파라미터 수(62.07M)를 기록하여 높은 효율성을 보였다.
• 정성적 결과에서는 D-LKA Net이 UNETR와 UNETR++보다 췌장의 비정상적인 형태를 더 잘 포착하여 분할 아티팩트를 감소시켰다.

Figure 2 : Architecture of the deformable LKA module.