[논문 리뷰] Large-kernel Attention for Efficient and Robust Brain Lesion Segmentation

Liam Chalcroft, Ruben Lourenço Pereira외 6명|arXiv (Cornell University)|2023. 08. 14.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 대규모 커널 주의(Large-kernel Attention, LKA)를 제안하며, 트랜스포머의 전역 수용 영역와 CNN의 매개변수 효율성 및 국소 패턴에 대한 인덕티브 바이어스를 융합한 하이브리드 컨volutional-attention 모듈이다. 이는 3D 뇌 병변 분할을 위한 U-Net을 향상시키며, 도메인 이동에 대한 강건성에서 최신 기술을 능가하고 BraTS 및 ATLAS 데이터셋에서 경쟁적인 성능을 보이며, 특히 형태 기반 일반화 및 분포 외 일반화에서 뛰어난 성능을 발휘한다.

ABSTRACT

Vision transformers are effective deep learning models for vision tasks, including medical image segmentation. However, they lack efficiency and translational invariance, unlike convolutional neural networks (CNNs). To model long-range interactions in 3D brain lesion segmentation, we propose an all-convolutional transformer block variant of the U-Net architecture. We demonstrate that our model provides the greatest compromise in three factors: performance competitive with the state-of-the-art; parameter efficiency of a CNN; and the favourable inductive biases of a transformer. Our public implementation is available at https://github.com/liamchalcroft/MDUNet .

연구 동기 및 목표

• 뇌 병변 분할에 핵심적인 장거리 공간 상관관계를 포착하지 못하는 표준 CNN의 한계를 해결하기 위해.
• 의료 영상에서 표준 비전 트랜스포머의 계산 비효율성과 국소 인덕티브 바이어스 부족 문제를 극복하기 위해.
• MRI 대조 시퀀스에서 성능, 매개변수 효율성, 도메인 이동에 대한 강건성 간 균형을 이루는 모델을 개발하기 위해.
• 대규모 커널 주의를 갖춘 하이브리드 아키텍처가 트랜스포머의 인덕티브 바이어스를 달성하면서도 CNN 유사 효율성을 유지할 수 있는지 조사하기 위해.

제안 방법

• 대규모 커널 주의(LKA) 기반의 새로운 전-컨volution 트랜스포머 블록을 제안하며, 대규모 커널을 깊이 분리형, 확장된 깊이 분리형, 포인트 와이즈 컨볼루션으로 분해한다.
• 표준 자기주도 주의 또는 컨볼루션 블록을 대체하여 LKA 모듈을 U-Net 아키텍처에 통합함으로써, 매개변수 감소와 함께 전역 맥락 모델링을 가능하게 한다.
• 요소 분해 주의 메커니즘을 활용: X₁ = dDWConv(DWConv(X)), X₂ = PConv(X₁), 이는 효율적인 공간 및 채널 혼합을 가능하게 한다.
• 다중 모odal MRI 데이터를 사용하여 교차 엔트로피 및 Dice 손실을 기반으로 엔드 투 엔드로 모델을 훈련시키며, 가우시안 블러링 및 강도 저항성 조작과 같은 데이터 증강 기법을 적용한다.
• BraTS에서 추론 안정성과 속도 향상을 위해 테스트 시 시간 증강 및 하프 정밀도 추론을 사용한다.
• ISLES 2015의 새로운 MRI 시퀀스(T1, T2, DWI, FLAIR)에 대한 제로샷 전이를 통해 일반화 성능을 평가한다.

(a) LKA

실험 결과

연구 질문

• RQ1 대규모 커널 주의 메커니즘이 계산 효율성을 유지하면서도 3D 뇌 병변 분할에서 장거리 상관관계를 효과적으로 모델링할 수 있는가?

주요 결과

• BraTS 2021 공개 테스트 세트에서 LKA-E 모델은 전반 병변(WT)의 Dice 스코어가 0.939를 기록하여 CNN 기반의 nnUNet(0.902)과 Swin 기반 모델(0.938)을 모두 초월했으며, HD95에서 통계적으로 유의미한 향상(2.449 vs. 3.606, p < 0.0001)을 보였다.
• ATLAS 데이터셋에서 LKA 모델은 모든 지표에서 경쟁적인 성능를 유지하였으며, 강화된 병변(ET)의 Dice 스코어는 0.888이었고, 평균 표면 거리(AVD) 측면에서 CNN 기반 베이스라인을 뚜렷이 능가하였다.
• ISLES 2015에서의 분포 외(OOD) 평가에서 LKA 모델은 MRI 대조 시퀀스 이동에 대해 뛰어난 강건성을 보였으며, T1, T2, DWI, FLAIR 시퀀스에서 모두 높은 성능를 유지하였다. 반면 nnUNet 기반 베이스라인은 성능이 크게 저하되었다.
• 이미지 블러링(σ = 0.5에서 2.0)에 대해 LKA 모델은 안정적인 Dice 및 HD95 지표를 유지하며 더 강한 형태 바이어스를 보였으며, 이는 Swin 트랜스포머에 가까운 성향을 보였고, nnUNet의 텍스처 기반 바이어스와 대비되었다.
• 종양 코어(TC)의 경우 LKA 모델은 중앙값 HD95가 2.000을 기록하여 Swin 기반 모델(2.236)과 CNN(2.236)보다 유의미하게 뛰어난 표면 정확도를 확보하였다.
• 제거 실험을 통해 LKA에서 커널 크기와 확장률을 조정함으로써 인덕티브 바이어스를 세밀하게 조정할 수 있으며, 작업에 맞는 적응을 위해 CNN과 트랜스포머 행동 사이의 스펙트럼을 제어할 수 있음을 확인하였다.

(b) Attention