3D U-KAN Implementation for Multi-modal MRI Brain Tumor Segmentation

2024-08-01

arXiv.org

AI Summary

To address the challenges of high intra-tumoral heterogeneity in gliomas and the difficulty of modeling 3D multimodal MRI data, this paper proposes UKAN-SE—the first U-Net variant adapted for 3D medical image segmentation by integrating the Kolmogorov–Arnold Network (KAN). It innovatively incorporates a Squeeze-and-Excitation (SE) module to model channel-wise global attention. The architecture employs 3D convolutions and multimodal feature fusion. Evaluated on BraTS 2024, UKAN-SE significantly outperforms baseline models including U-Net, Attention U-Net, and Swin UNETR in segmentation accuracy. With only 10.6 million parameters, it achieves superior computational efficiency: training time is reduced to one-quarter that of U-Net and one-sixth that of Swin UNETR. Thus, UKAN-SE delivers both state-of-the-art performance and exceptional parameter- and time-efficiency for 3D glioma segmentation.

Technology Category

• Machine Learning: Multimodal Learning
• Computer Vision: Segmentation
• Search and Optimization: Combinatorial Optimization

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

We explore the application of U-KAN, a U-Net based network enhanced with Kolmogorov-Arnold Network (KAN) layers, for 3D brain tumor segmentation using multi-modal MRI data. We adapt the original 2D U-KAN model to the 3D task, and introduce a variant called UKAN-SE, which incorporates Squeeze-and-Excitation modules for global attention. We compare the performance of U-KAN and UKAN-SE against existing methods such as U-Net, Attention U-Net, and Swin UNETR, using the BraTS 2024 dataset. Our results show that U-KAN and UKAN-SE, with approximately 10.6 million parameters, achieve exceptional efficiency, requiring only about 1/4 of the training time of U-Net and Attention U-Net, and 1/6 that of Swin UNETR, while surpassing these models across most evaluation metrics. Notably, UKAN-SE slightly outperforms U-KAN.

Problem

• Improving 3D brain tumor segmentation in multi-modal MRI
• Addressing modality variability and computational complexity challenges
• Enhancing feature fusion and multi-scale representation in U-KAN

Innovation

• Extends U-KAN with 3D for MRI segmentation
• Uses Efficient Channel Attention for feature fusion
• Implements Pyramid Feature Aggregation for multi-scale