U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation

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 UKAN 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.

한 줄 요약

U-KAN은 KAN을 U-Net에 통합하여 의료 이미지 분할 및 생성에서 뛰어난 성능과 해석성을 제공하는 새로운 백본을 제안한다.

핵심 기여도

U-Net 파이프라인에 KAN 레이어를 통합한 **tokenized KAN 블록**을 제안.
**의료 이미지 분할 벤치마크**에서 높은 정확도와 낮은 계산 비용을 달성.
U-KAN을 **확산 모델의 노이즈 예측기**로 활용 가능성 검증.
**Small, Large 모델**을 통해 성능-비용 균형을 실험적으로 분석.

핵심 아이디어

기존 U-Net 및 변형 모델은 선형 커널 기반의 패턴 모델링과 해석성 부족이라는 한계를 가진다. 이를 해결하기 위해, 저자들은 **Kolmogorov-Arnold Networks (KANs)**의 비선형 학습 능력과 해석성을 활용한다. KAN은 Kolmogorov-Arnold 표현 정리에 기반한 **비선형 학습 가능한 활성화 함수 스택**을 통해 신경망을 재구성한다. 본 연구는 U-Net의 중간 표현에 KAN 레이어를 통합한 **tokenized KAN 블록**을 도입하여, 비선형 패턴을 효과적으로 추출하고, 해석 가능한 모델 구조를 구축한다. 이는 기존 U-Net의 선형성과 불확실성을 극복하는 핵심 아이디어이다.

기술적 접근법

**U-KAN 아키텍처**: U-Net의 인코더-디코더 구조에 **tokenized KAN 블록**을 추가.
**tokenized KAN 블록**: 중간 특징을 토큰으로 변환한 후, KAN 연산자를 적용하여 비선형 패턴을 추출.
**모델 크기 실험**: Small (64-96-128 채널), Large (256-320-512 채널) 모델을 비교.
**기본 모델**: 128-160-256 채널 설정으로 성능과 비용의 균형을 맞춤.
**하이퍼파라미터**: 채널 수를 제외하면 U-Net의 기존 설정을 유지.

주요 결과

**의료 이미지 분할**에서 U-KAN은 기존 U-Net 대비 **높은 정확도**와 **낮은 계산 비용**을 달성.
**Small 모델**은 64-96-128 채널로도 뛰어난 성능을 보임.
**Large 모델**은 256-320-512 채널로 성능 향상이 관찰됨.
**확산 모델**에서 U-KAN을 노이즈 예측기로 사용할 경우, 기존 U-Net보다 **더 나은 생성 성능**을 보임.

의의 및 한계

U-KAN은 KAN의 해석성과 비선형 모델링 능력을 U-Net에 결합하여, 의료 이미지 분할 및 생성 분야에서 높은 정확도와 효율성을 동시에 달성한다. 특히, **tokenized KAN 블록**은 기존 U-Net의 선형성과 불확실성을 개선하는 데 기여하며, 의료 영상 분석에서 신뢰성 있는 모델 설계를 가능하게 한다. 그러나, KAN 기반 모델의 **대규모 데이터셋에 대한 일반화 능력**이나 **복잡한 비주얼 태스크에서의 확장성**은 추가 연구가 필요하다. 또한, KAN의 수학적 구조가 복잡하기 때문에, **실제 임상 환경에서의 적용 가능성**도 검증이 필요하다.

실용적 활용

U-KAN은 의료 이미지 분할, 이미지 생성, 확산 모델 등 다양한 시각적 작업에서 활용 가능하다. 특히, **해석성이 필요한 의료 진단 시스템**이나 **자원이 제한된 포인트 오브 케어 환경**에서 유용하게 사용될 수 있다. 또한, **의료 AI 모델의 신뢰성 향상**을 위한 기초 연구로도 활용 가능하다.