[논문 리뷰] UKAN-EP: Enhancing U-KAN with Efficient Attention and Pyramid Aggregation for 3D Multi-Modal MRI Brain Tumor Segmentation

Yanbing Chen, Tianze Tang외 2명|arXiv (Cornell University)|2024. 08. 01.

Medical Image Segmentation Techniques인용 수 6

한 줄 요약

본 논문은 U-KAN을 3D 다중 모달 MRI 뇌종양 분할에 적용하고 SE 주의(attention)를 갖춘 UKAN-SE(U-KAN with SE attention)를 도입하여, BraTS 2024에서 U-Net, Attention U-Net, Swin UNETR와 비교하고, 상당히 짧아진 학습 시간과 약 10.6M 파라미터로 강한 정확도를 보임을 보인다.

ABSTRACT

Background: Gliomas are among the most common malignant brain tumors and exhibit substantial heterogeneity, complicating accurate detection and segmentation. Although multi-modal MRI is the clinical standard for glioma imaging, variability across modalities and high computational demands hamper effective automated segmentation. Methods: We propose UKAN-EP, a novel 3D extension of the original 2D U-KAN model for multi-modal MRI brain tumor segmentation. While U-KAN integrates Kolmogorov-Arnold Network (KAN) layers into a U-Net backbone, UKAN-EP further incorporates Efficient Channel Attention (ECA) and Pyramid Feature Aggregation (PFA) modules to enhance inter-modality feature fusion and multi-scale feature representation. We also introduce a dynamic loss weighting strategy that adaptively balances cross-entropy and Dice losses during training. Results: On the 2024 BraTS-GLI dataset, UKAN-EP achieves superior segmentation performance (e.g., Dice = 0.9001 $\pm$ 0.0127 and IoU = 0.8257 $\pm$ 0.0186 for the whole tumor) while requiring substantially fewer computational resources (223.57 GFLOPs and 11.30M parameters) compared to strong baselines including U-Net, Attention U-Net, Swin UNETR, VT-Unet, TransBTS, and 3D U-KAN. An extensive ablation study further confirms the effectiveness of ECA and PFA and shows the limited utility of self-attention and spatial attention alternatives. Conclusion: UKAN-EP demonstrates that combining the expressive power of KAN layers with lightweight channel-wise attention and multi-scale feature aggregation improves the accuracy and efficiency of brain tumor segmentation.

연구 동기 및 목표

• 다중 모달 MRI 뇌종양 분할을 위한 3D U-Net 프레임워크 내에서 Kolmogorov-Arnold Networks (KAN)의 활용에 대한 동기를 제시한다.
• 2D U-KAN 모델을 3D에 적응시키고 Global attention을 위한 Squeeze-and-Excitation(SE) 모듈을 도입하여 UKAN-SE를 제안한다.
• BraTS 2024 Task 1 데이터세트에서 U-KAN과 UKAN-SE를 U-Net, Attention U-Net, Swin UNETR과 비교하여 평가한다.
• 분할 성능을 유지하거나 향상시키면서 파라미터 수와 학습 시간의 효율성을 강조한다.

제안 방법

• Convolution 단계 뒤에 Tokenized KAN(Tok-KAN) 단계를 포함하는 구성을 갖는 3D U-KAN의 적응 모델을 사용한다.
• 패턴 모델링을 향상시키기 위해 학습 가능한 활성화 함수를 갖는 KAN 계층을 도입한다.
• 각 컨볼루션 블록 뒤에 SE 모듈을 추가한 글로벌 어텐션을 갖는 U-KAN 변형인 UKAN-SE를 도입한다.
• BraTS 2024 Task 1에서 네 가지 MRI 모달리티(T1, T1Gd, T2, FLAIR)를 사용하여 다섯 개 모델(U-Net, Attention U-Net, Swin UNETR, U-KAN, UKAN-SE)을 학습 및 평가한다.
• 손실은 교차 엔트로피와 Soft Dice 손실의 조합이다: L_total = (1-α) L_CE + α(1 - SoftDice)으로 α = 0.5이다.

실험 결과

연구 질문

• RQ1 3D U-KAN 및 UKAN-SE가 BraTS 2024 Task 1에서 기존 기준선(U-Net, Attention U-Net, Swin UNETR)과 비교해 어떤 성능을 보이는가?
• RQ2 전통적인 U-Net 기반 모델에 비해 U-KAN/UKAN-SE를 사용할 때 파라미터 수 및 학습 시간 측면에서의 효율성 향상은 무엇인가?
• RQ3 SE 주의(attention) 도입(UKAN-SE)이 분할 정확도와 경계 구분에서 U-KAN에 비해 측정 가능한 이점을 제공하는가?
• RQ4 중간 크기의 KAN 구성([128,160,256])이 3D 다중 모달 MRI 뇌종양 분할에 미치는 영향은 무엇인가?

주요 결과

• UKAN-SE는 ET 및 RC를 특히 포함한 대부분의 종양 하위 영역에서 병변별 및 전체 영상 Dice 점수에서 최상의 성능을 달성한다.
• UKAN-SE는 일반적으로 U-KAN보다 우수한 성능을 보이며 SE 기반 글로벌 어텐션의 추가 가치를 시사한다.
• U-KAN과 UKAN-SE는 약 10.6M 파라미터로 강력한 성능을 제공하며 U-Net, Attention U-Net, Swin UNETR에 비해 학습 시간이 크게 더 빠르다.
• U-KAN은 에포크당 학습 시간이 훨씬 짧아(약 803초) U-Net(약 3322초)보다 경쟁력 있는 결과를 얻는다.
• UKAN-SE는 U-KAN에 비해 파라미터 수와 학습 시간의 증가가 제한적이면서도 분할 정확도를 향상시킨다.