[논문 리뷰] PMFSNet: Polarized Multi-scale Feature Self-attention Network For Lightweight Medical Image Segmentation

Jiahui Zhong, Wenhong Tian외 5명|arXiv (Cornell University)|2024. 01. 15.

Advanced Neural Network Applications인용 수 9

한 줄 요약

PMFSNet는 전역 컨텍스트와 로컬 특징의 균형을 맞추는 극성 다중 스케일 자기 주의 블록을 사용하는 가벼운 UNet 스타일 분할 모델(매개변수 ≤1M)로, 2D/3D 의학 영상에서 계산량을 줄이며 경쟁력 있는 정확도를 달성합니다.

ABSTRACT

Current state-of-the-art medical image segmentation methods prioritize accuracy but often at the expense of increased computational demands and larger model sizes. Applying these large-scale models to the relatively limited scale of medical image datasets tends to induce redundant computation, complicating the process without the necessary benefits. This approach not only adds complexity but also presents challenges for the integration and deployment of lightweight models on edge devices. For instance, recent transformer-based models have excelled in 2D and 3D medical image segmentation due to their extensive receptive fields and high parameter count. However, their effectiveness comes with a risk of overfitting when applied to small datasets and often neglects the vital inductive biases of Convolutional Neural Networks (CNNs), essential for local feature representation. In this work, we propose PMFSNet, a novel medical imaging segmentation model that effectively balances global and local feature processing while avoiding the computational redundancy typical in larger models. PMFSNet streamlines the UNet-based hierarchical structure and simplifies the self-attention mechanism's computational complexity, making it suitable for lightweight applications. It incorporates a plug-and-play PMFS block, a multi-scale feature enhancement module based on attention mechanisms, to capture long-term dependencies. Extensive comprehensive results demonstrate that even with a model (less than 1 million parameters), our method achieves superior performance in various segmentation tasks across different data scales. It achieves (IoU) metrics of 84.68%, 82.02%, and 78.82% on public datasets of teeth CT (CBCT), ovarian tumors ultrasound(MMOTU), and skin lesions dermoscopy images (ISIC 2018), respectively. The source code is available at https://github.com/yykzjh/PMFSNet.

연구 동기 및 목표

• 정확한 의료 영상 분할이 필요성과 낮은 계산 부담 및 작은 모델 크기 사이의 필요성에 대응합니다.
• 전역 컨텍스트 모델링과 로컬 특징 표현을 경량 아키텍처에서 균형 있게 다룹니다.
• 중복 제거 및 엣지 디바이스 배치를 가능하게 하는 플러그 앤 플레이 주의 블록을 제안합니다.
• 다중 스케일 특징 융합을 활용해 제곱 비용 없이 장거리 의존성을 포착합니다.]
• method_0_translation_placeholder?
• method_1_translation_placeholder?

제안 방법

• PMFSNet를 엔코더-디코더 유사 구조와 병목 부근의 플러그 앤 플레이 PMFS 블록으로 도입합니다.
• Adaptive Multi-branch Feature Fusion (AMFF)을 사용해 서로 다른 단계의 다중 분기 특징을 통합합니다.
• Polarized Multi-scale Channel Self-attention (PMCS)과 Polarized Multi-scale Spatial Self-attention (PMSS)을 구현해 채널과 공간 전반의 주의 포인트를 확장합니다.
• 선형화된 다중 스케일 형식으로 쿼리, 키, 값으로 주의를 계산하고 특징을 요소별로 강화합니다.
• 클래스 불균형을 다루기 위해 Weighted Extended Dice Loss (WEDL)로 학습합니다.]
• research_questions_0_placeholder?
• research_questions_1_placeholder?
• research_questions_2_placeholder?
• research_questions_3_placeholder?

실험 결과

연구 질문

• RQ1 PMFSNet가 2D 및 3D 의료 영상에서 ≤1M 매개변수로 경쟁력 있는 분할 성능을 달성할 수 있는가?
• RQ2 PMFS 구성요소가 SOTA 모델에 비해 계산 복잡도 감소를 유지하거나 향상시키면서 정확도를 보이는가?
• RQ3 적응형 다중 분기 융합과 극성 다중 스케일 주의가 분할의 전역 및 로컬 특징 표현에 어떤 영향을 미치는가?
• RQ4 의료 분할 과제에서 WEDL이 클래스 불균형 문제를 다루는 데 효과적인가?

주요 결과

Method	FLOPs(G) ↓	Params(M) ↓	HD(mm) ↓	ASSD(mm) ↓	IoU(%) ↑	SO(%) ↑	DSC(%) ↑
UNet3D	2223.03	16.32	113.79	22.40	70.62	70.72	36.67
DenseVNet	23.73	0.87	8.21	1.14	84.57	94.88	91.15
AttentionUNet3D	2720.79	94.48	147.10	61.10	52.52	42.49	64.08
DenseVoxelNet	402.32	1.78	41.18	3.88	81.51	92.50	89.58
MultiResUNet3D	1505.38	17.93	74.06	8.17	76.19	81.70	65.45
UNETR	229.19	93.08	107.89	17.95	74.30	73.14	81.84
SwinUNETR	912.35	62.19	82.71	7.50	83.10	86.80	89.74
TransBTS	306.80	33.15	29.03	4.10	82.94	90.68	39.32
nnFormer	583.49	149.25	51.28	5.08	83.54	90.89	90.66
3D UX-Net	1754.79	53.01	108.52	19.69	75.40	73.48	84.89
PMFSNet3D (Ours)	15.14	0.63	5.57	0.79	84.68	95.10	91.30

• PMFSNet3D는 3D CBCT 치아 데이터셋에서 IoU 84.68%, SO 95.10%, DSC 91.30%를 0.63M 매개변수와 15.14 GFLOPs로 달성합니다.
• PMFSNet3D는 여러 SOTA 모델과 비교하여 FLOPs와 매개변수를 크게 줄이면서도 높은 IoU 및 DSC를 유지합니다.
• 3개의 데이터셋(CBCT, MMOTU 초음파, ISIC 2018)에서 PMFSNet 변형들이 경쟁 방법들보다 훨씬 적은 매개변수로도 강한 정확도를 달성합니다.
• PMFS 블록은 선형 주의의 복잡도에 근접한 긴 거리 의존성 모델링을 가능하게 하며 엣지 디바이스 배치를 돕습니다.]
• table_headers_translation_placeholder?
• table_rows translation_placeholder?