[논문 리뷰] U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation

Chenxin Li, Xinyu Liu외 6명|arXiv (Cornell University)|2024. 06. 05.

Brain Tumor Detection and Classification인용 수 43

한 줄 요약

U-KAN은 Kolmogorov-Arnold Networks (KAN)을 U-Net 백본에 통합하여 의료 영상 분할 및 확산 기반 영상 생성 성능을 향상시키며, 해석 가능성이 개선된 높은 정확도와 효율성을 달성합니다.

ABSTRACT

U-Net has become a cornerstone in various visual applications such as image segmentation and diffusion probability models. While numerous innovative designs and improvements have been introduced by incorporating transformers or MLPs, the networks are still limited to linearly modeling patterns as well as the deficient interpretability. To address these challenges, our intuition is inspired by the impressive results of the Kolmogorov-Arnold Networks (KANs) in terms of accuracy and interpretability, which reshape the neural network learning via the stack of non-linear learnable activation functions derived from the Kolmogorov-Anold representation theorem. Specifically, in this paper, we explore the untapped potential of KANs in improving backbones for vision tasks. We investigate, modify and re-design the established U-Net pipeline by integrating the dedicated KAN layers on the tokenized intermediate representation, termed U-KAN. Rigorous medical image segmentation benchmarks verify the superiority of U-KAN by higher accuracy even with less computation cost. We further delved into the potential of U-KAN as an alternative U-Net noise predictor in diffusion models, demonstrating its applicability in generating task-oriented model architectures. These endeavours unveil valuable insights and sheds light on the prospect that with U-KAN, you can make strong backbone for medical image segmentation and generation. Project page:\url{https://yes-u-kan.github.io/}.

연구 동기 및 목표

• 더 나은 비선형 모델링과 해석 가능성을 갖춘 의학 영상 분할용 백본 개선의 필요성을 제시한다.
• KAN을 U-Net 파이프라인에 통합하기 위한 토큰화된 KAN 블록을 도입한다.
• 다양한 의학 데이터셋에서 우수한 분할 정확도를 시연한다.
• 확산 기반 영상 생성으로 아키텍처를 확장하고 생성 능력을 평가한다.

제안 방법

• 두 구간 아키텍처를 제안한다: 컨볼루션 구절 다음에 토큰화된 KAN 구절이 오고, 그다음 스킵 연결이 있는 U자형 디코더가 따른다.
• 컨볼루션 특징의 토큰화를 사용하여 패치를 만들고 이를 잔차 연결과 층 정규화를 갖춘 다중 KAN 층으로 처리한다.
• 중간 특징을 D차원 임베딩에 투영하고 K개의 층(Phi_i)으로 구성된 KAN으로 처리한 후 DwConv와 BN을 거치는 토큰화된 KAN 블록을 구현한다.
• 시간 임베딩을 주입하고 DDPM 프레임워크에서 잡음 epsilon_t를 예측하도록 학습하여 확산 U-KAN으로 확장하며, 시간 조건이 있는 KAN과 MSE 손실(L_diff)을 사용한다.
• 세 개의 의학 데이터셋에서 이진 교차 엔트로피와 Dice 손실의 조합으로 분할 모델을 학습하고, IoU와 F1 점수로 평가하며, 효율성(Gflops, Params)을 보고한다.
• FID와 IS 지표를 사용하여 표준 U-Net 변형에 비해 확산 백본을 벤치마크한다.

실험 결과

연구 질문

• RQ1 Kolmogorov-Arnold Networks를 U-Net 백본에 통합하면 큰 계산 부담 없이 의학 영상의 분할 정확도가 향상될 수 있는가?
• RQ2 전통적인 CNN/MLP/트랜스포머 백본에 비해 토큰화된 KAN 블록이 의학 영상에서 해석 가능성과 성능상의 이점을 제공하는가?
• RQ3 U-KAN이 분할 외의 태스크 지향적 영상 생성을 위한 확산 모델 백본으로 효과적으로 작동할 수 있는가?

주요 결과

• Seg-U-KAN은 여러 베이스라인보다 높은 분할 정확도를 달성한다(예: BUSI, GlaS, CVC 데이터셋에서의 IoU/F1).
• Seg. U-KAN의 평균 분할 지표는 IoU 78.69와 F1 87.22에 도달하며, 데이터셋 전반에서 경쟁력 있는 효율성(Gflops 14.02, Params 6.35)을 보인다.
• 변수 시험 결과, 세 개의 KAN 층이 테스트 구성 중 최상의 분할 성능(IoU 66.65, F1 79.75)을 보인다.
• KAN 층을 MLP로 대체하면 성능이 저하되며, 이 작업에서 KAN 블록의 효능이 강조된다.
• 확산 U-KAN은 표준 확산 U-Net 변형들보다 생성 지표를 개선하며, BUSI, GlaS, CVC 데이터셋에서 더 나은 FID/IS를 보인다(예: KANBlock을 포함한 Diffusion U-KAN이 우수한 FID/IS를 보임).
• 세 개의 KAN 층 구성과 시간 임베딩 확장은 분할 및 확산 생성 성능을 모두 향상시킨다.
• 모델 규모 확장 연구는 더 큰 U-KAN 변형이 IoU/F1에서 추가 이점을 제공하나 Gflops 증가 비용이 있음을 시사한다.