[論文レビュー] UKAN-EP: Enhancing U-KAN with Efficient Attention and Pyramid Aggregation for 3D Multi-Modal MRI Brain Tumor Segmentation

Yanbing Chen, Tianze Tang+2 more|arXiv (Cornell University)|Aug 1, 2024

Medical Image Segmentation Techniques被引用数 6

ひとこと要約

本論文は U-KAN を 3D の多 modality MRI 脳腫瘍セグメンテーションへ適応し、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.

研究の動機と目的

• 多 modality MRI 脳腫瘍セグメンテーションのために、3D U-Net フレームワーク内で Kolmogorov-Arnold Networks (KAN) の利用を提案する。
• 2D の U-KAN モデルを 3D に適応し、グローバル注意機構のために Squeeze-and-Excitation (SE) モジュールを組み込んだ UKAN-SE を提案する。
• BraTS 2024 Task 1 データセット上で、U-KAN および UKAN-SE を U-Net、Attention U-Net、Swin UNETR と比較して評価する。
• パラメータ数とトレーニング時間の効率性を、セグメンテーション性能を維持または向上させつつ強調する。

提案手法

• Convolution Phase に続く Tokenized KAN (Tok-KAN) Phase を備えた 3D 適応版 U-KAN を用いる。
• パターンモデリングを改善するため、学習可能な活性化関数を備えた KAN 層を組み込む。
• 各畳み込みブロックの後に SE モジュールを配置したグローバル注意機構を備える U-KAN の変種 UKAN-SE を導入する。
• BraTS 2024 Task 1 で four MRI modalities (T1, T1Gd, T2, FLAIR) を用いて、5 モデル (U-Net, Attention U-Net, Swin UNETR, U-KAN, UKAN-SE) を訓練・評価する。
• 損失関数はクロスエントロピーとソフト Dice ロスの組み合わせであり： L_total = (1-α) L_CE + α(1 - SoftDice) with α = 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注意機構の組み込み（UKAN-SE）は、セグメンテーション精度と境界の delineation において、U-KAN より測定可能な利得を提供するか？
• RQ4 中規模の KAN 設定（[128,160,256]）が 3D 多 modality MRI 脳腫瘍セグメンテーションに与える影響は何か？

主な発見

• UKAN-SE は、多くの腫瘍サブ領域で病変ごとの Dice スコアと全画像 Dice スコアの双方で最良を達成し、特に ET と RC で顕著である。
• UKAN-SE は一般に U-KAN を上回り、SE ベースのグローバル注意機構の付加価値を示している。
• U-KAN および UKAN-SE は約 10.6M パラメータで強力な性能を示し、U-Net、Attention U-Net、Swin UNETR と比較してトレーニング時間が大幅に短い。
• U-KAN は U-Net（約 3322s）よりはるかに短い1エポックあたりのトレーニング時間（約 803s）で競争力のある結果を達成する。
• UKAN-SE は、U-KAN に比べてパラメータ数とトレーニング時間の増加が控えめでありながら、セグメンテーション精度を向上させる。