[논문 리뷰] P-Mamba: Marrying Perona Malik Diffusion with Mamba for Efficient Pediatric Echocardiographic Left Ventricular Segmentation

Zi Ye, Tianxiang Chen외 3명|arXiv (Cornell University)|2024. 02. 13.

Cardiovascular Function and Risk FactorsMedicine인용 수 3

한 줄 요약

P-Mamba는 시각 Mamba를 사용해 전역적 종속성 모델링을 하고, DWT 기반 Perona-Malik 확산(PMD) 블록을 통해 노이즈 억제와 에지 유지 보존을 실현함으로써 소아 에코카디오그래픽 좌심실 분할을 위한 새로운 효율적인 딥러닝 프레임워크를 제안한다. 이 모델은 PSAX에서 0.9221, A4C에서 0.9046의 Dice 점수로 최신 기술(SOTA) 수준의 정확도를 달성하였으며, 이 Fourfold 복잡도와 선형 복잡도를 갖는 시각 트랜스포머보다도 뛰어난 추론 효율성을 보였다.

ABSTRACT

In pediatric cardiology, the accurate and immediate assessment of cardiac function through echocardiography is crucial since it can determine whether urgent intervention is required in many emergencies. However, echocardiography is characterized by ambiguity and heavy background noise interference, causing more difficulty in accurate segmentation. Present methods lack efficiency and are prone to mistakenly segmenting some background noise areas, such as the left ventricular area, due to noise disturbance. To address these issues, we introduce P-Mamba, which integrates the Mixture of Experts (MoE) concept for efficient pediatric echocardiographic left ventricular segmentation. Specifically, we utilize the recently proposed ViM layers from the vision mamba to enhance our model's computational and memory efficiency while modeling global dependencies.In the DWT-based Perona-Malik Diffusion (PMD) Block, we devise a PMD Block for noise suppression while preserving the left ventricle's local shape cues. Consequently, our proposed P-Mamba innovatively combines the PMD's noise suppression and local feature extraction capabilities with Mamba's efficient design for global dependency modeling. We conducted segmentation experiments on two pediatric ultrasound datasets and a general ultrasound dataset, namely Echonet-dynamic, and achieved state-of-the-art (SOTA) results. Leveraging the strengths of the P-Mamba block, our model demonstrates superior accuracy and efficiency compared to established models, including vision transformers with quadratic and linear computational complexity.

연구 동기 및 목표

• 소아 에코카디오그래머에서의 노이즈 간섭과 낮은 효율성 문제를 해결한다.
• 배경 노이즈 억제를 통해 정확도를 향상시키면서도 미세한 좌심실 경계 세부 정보를 유지한다.
• 기존의 시각 트랜스포머 기반 모델들과 비교해 계산 및 메모리 효율성을 향상시킨다.
• 선택적 상태공간 모델링과 확산 기반 노이즈 제거 기법을 통합된 아키텍처로 통합하여 초음파 영상 분할에 적용한다.
• 최소한의 계산 비용으로 소아 에코카디오그래머 데이터셋에서 뛰어난 성능을 달성한다.

제안 방법

• 선형 복잡도와 하드웨어 효율성을 갖춘 시각 Mamba 인코더를 사용한 이중 인코더 아키텍처를 도입한다.
• 이산 웨이블릿 변환(DWT) 분해 후 특징 맵에 Perona-Malik 확산을 적용하는 DWT 기반 PMD 블록을 제안하여 노이즈 억제와 에지 강화를 실현한다.
• 입력 특징을 DWT를 통해 하위대역으로 분해함으로써, 좌심실 영역의 구조적 세부 정보를 유지하면서 국소적 노이즈 필터링을 가능하게 한다.
• 다양한 척도에서 Mamba 및 PMD 브랜치의 특징 맵을 융합한 후, 마스크 예측을 위해 분류 헤드와 FCN 디코더를 통과시킨다.
• 경계 정확도 향상을 위해 Dice 손실 및 기타 표준 분할 손실을 사용한 다중 척도 감독 학습을 수행한다.
• 자기주의 주의(self-attention) 없이도 장거리 종속성을 효율적으로 모델링할 수 있도록 Mamba 블록의 선택적 구조적 상태공간 메커니즘을 활용한다.

실험 결과

연구 질문

• RQ1 Mamba와 Perona-Malik 확산을 융합한 하이브리드 아키텍처가 노이즈가 많은 소아 에코카디오그래머에서 분할 정확도 향상에 기여하는가?
• RQ2 DWT 기반 PMD 블록의 통합이 표준 CNN 또는 주의 메커니즘과 비교해 에지 유지 보존과 노이즈 억제 성능을 향상시키는가?
• RQ3 Mamba 기반 인코더가 시각 트랜스포머보다 계산 비용을 얼마나 줄일 수 있으며, 성능을 유지하거나 향상시키는가?
• RQ4 소아 에코카디오그래머 벤치마크에서 제안된 P-Mamba 모델은 최신 기술 모델들과 비교해 효율성과 정확도 면에서 어떤가?
• RQ5 아블레이션 스터디를 통해 각 구성요소(Mamba 대비 PMD)가 최종 분할 성능에 기여하는 정도는 어떠한가?

주요 결과

• P-Mamba는 PSAX 데이터셋에서 0.9221, A4C 데이터셋에서 0.9046의 Dice 점수를 기록하여 비교한 모든 최신 기술(SOTA) 방법들을 능가했다.
• 모델의 추론 시간은 23.49ms로 단축되었으며, 파rameter 수는 183.37M, GPU 메모리 사용량은 12.22GB로, PVT 및 MaxViT 등의 모델보다 효율성이 뛰어나게 향상되었다.
• 아블레이션 스터디 결과, DWT 기반 PMD 블록을 제거할 경우 PSAX에서 Dice 점수가 0.0028 감소함을 확인하여, 노이즈 억제 및 에지 충실도 확보에 핵심적인 역할을 함을 입증했다.
• Mamba 블록을 ViT 기반 모델(PVT, Flatten Transformer 등)로 대체할 경우 정확도가 낮아짐을 확인하여, Mamba의 뛰어난 모델링 능력을 입증했다.
• 간단한 소벨 연산자보다 DWT 기반 PMD 블록이 더 높은 성능을 보였으며, 이는 노이즈 억제 기능이 없는 소벨 연산자와의 비교를 통해 확증되었다.
• PSAX에서 정밀도(0.9316)와 재현율(0.9128)이 모두 높아, 임의의 양성 결과나 누락 감지에 대해 강건함을 입증했다.