[논문 리뷰] TransUKAN:Computing-Efficient Hybrid KAN-Transformer for Enhanced Medical Image Segmentation

Yanlin Wu, Tao Li외 3명|arXiv (Cornell University)|2024. 09. 23.

Medical Image Segmentation Techniques인용 수 4

한 줄 요약

TransUKAN은 의료 영상 분할을 위한 계산 효율적인 하이브리드 KAN-Transformer 아키텍처를 제안하며, 기존의 MLP와 QKV 투영을 개선된 파rameter 효율적인 KAN(EfficientKAN)으로 대체하여 비선형 모델링을 향상시키고 계산 부담을 줄입니다. 이 모델은 20.85M개의 파라미터로 다수의 의료 데이터셋에서 최신 기술 수준의 성능을 달성하여 뛰어난 정확도, 효율성 및 강인성을 입증합니다.

ABSTRACT

U-Net is currently the most widely used architecture for medical image segmentation. Benefiting from its unique encoder-decoder architecture and skip connections, it can effectively extract features from input images to segment target regions. The commonly used U-Net is typically based on convolutional operations or Transformers, modeling the dependencies between local or global information to accomplish medical image analysis tasks. However, convolutional layers, fully connected layers, and attention mechanisms used in this process introduce a significant number of parameters, often requiring the stacking of network layers to model complex nonlinear relationships, which can impact the training process. To address these issues, we propose TransUKAN. Specifically, we have improved the KAN to reduce memory usage and computational load. On this basis, we explored an effective combination of KAN, Transformer, and U-Net structures. This approach enhances the model's capability to capture nonlinear relationships by introducing only a small number of additional parameters and compensates for the Transformer structure's deficiency in local information extraction. We validated TransUKAN on multiple medical image segmentation tasks. Experimental results demonstrate that TransUKAN achieves excellent performance with significantly reduced parameters. The code will be available athttps://github.com/wuyanlin-wyl/TransUKAN.

연구 동기 및 목표

• 기존 U-Net 및 Transformer 기반 모델의 높은 파라미터 수와 계산 비용 문제를 해결하기 위해.
• 표준 KAN이 영상 처리에서 가지는 한계를 극복하기 위해 메모리와 계산을 줄이는 스파arsed, 효율적인 변형(EfficientKAN)을 도입하기 위해.
• U-Net 프레임워크 내에서 KAN을 Transformer 메커니즘에 통합하여 국소적 특징 모델링과 전역 의존성 학습을 향상시키기 위해.
• 모델 크기와 학습 복잡성을 크게 줄이며 동시에 최신 기술 수준의 분할 정확도를 달성하기 위해.

제안 방법

• Eff율리브KAN은 평균 풀링을 통해 키 활성화 값만 통합하여 행렬 계산 복잡도를 낮춘 KAN의 스파arsed 버전입니다.
• 멀티헤드 자기주의 메커니즘의 완전 연결 층과 QKV 투영을 EfficientKAN으로 대체하여 더 적은 파라미터로 비선형 적합성을 향상시킵니다.
• 스킵 연결을 갖춘 U-Net 유사 인코더-디코더 아키텍처에 향상된 KAN 모듈을 통합하여 공간적 세부 정보를 유지합니다.
• CNN, U-Net, Transformer, KAN을 융합한 하이브리드 설계를 통해 국소적 및 전역적 특징 추출의 균형을 맞춥니다.
• KAN에 학습 가능한 활성화 함수를 도입하여 표준 MLP보다 더 효율적으로 복잡한 비선형 관계를 모델링합니다.
• 교차 검증과 아블레이션 연구를 통해 다양한 의료 영상 데이터셋에서 아키텍처의 유효성을 검증합니다.

실험 결과

연구 질문

• RQ1 Transformer의 MLP와 QKV 투영을 KAN 기반으로 대체함으로써 파라미터 수와 계산량을 줄이면서도 분할 성능을 향상시킬 수 있는가?
• RQ2 EfficientKAN은 표준 KAN과 ReLUKAN에 비해 의료 영상 분할 작업에서 파라미터 효율성과 성능 면에서 어떻게 비교되는가?
• RQ3 KAN을 U-Net-Transformer 하이브리드 아키텍처에 통합함으로써 국소 세부 정보 추출과 전역 맥락 모델링이 모두 향상되는가?
• RQ4 제안된 TransUKAN 모델이 다양한 모odal리티와 해부학적 구조를 가진 다양한 의료 영상 데이터셋에 얼마나 잘 일반화되는가?
• RQ5 기존 최신 기술 수준의 방법들보다 크게 줄어든 파라미터 수로도 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• ISIC 데이터셋에서 TransUKAN은 91.17%의 Dice 스코어를 기록하여 베이스라인 TransUNet(86.30%) 및 기타 최신 기술 수준 모델을 초월합니다.
• Kvasir 데이터셋에서 TransUKAN은 87.75%의 Dice를 기록하여 다음으로 우수한 모델(86.30%)보다 1.45%p 높습니다.
• TransUNet의 105.32M에 비해 파라미터를 20.85M로 크게 줄였으며, 모든 평가된 데이터셋에서 성능 향상을 이룹니다.
• 아블레이션 연구 결과, EfficientKAN은 파라미터 수와 학습 시간을 감소시키면서 정확도를 향상시키는 반면, ReLUKAN은 파라미터 수를 늘리고 성능을 떨어뜨리는 것으로 확인되었습니다.
• TransUKAN은 다양한 해부학적 구조에서 강력한 일반화 능력을 유지하여 NKUT 데이터셋에서 90.29%의 Dice를 기록하고 MWT에서 89.09%를 달성했습니다.
• BUSI, ISIC, Kvasir, SAM 포함 총 일곱 개의 평가된 데이터셋에서 모두 최신 기술 수준의 성능을 달성하였으며, 이전 방법들에 비해 일관된 향상을 보였습니다.