Retinal vessel segmentation using Kolmogorov-Arnold networks and test-time adaptation
摘要
Recently, TransUNet has established itself as a state-of-the-art architecture for medical image segmentation. Although attention mechanisms have yielded remarkable advances, its Multi-Layer Perceptron (MLP)-based design still suffers from limitations in capturing the complex non-linear representations essential for segmenting intricate anatomical structures. Additionally, its static inference paradigm lacks robustness when confronted with domain shifts, a common challenge in medical image data. To address these inherent limitations, we propose a dual-optimized segmentation framework, drawing inspiration from the superior non-linear approximation capability of Kolmogorov–Arnold Networks (KAN) and the adaptive advantage of Test-Time Adaptation (TTA). Specifically, we redesign the feature transformation pathway within the Transformer encoder by integrating learnable spline-based mappings derived from KAN, which significantly enhances the model’s representation capacity. Furthermore, we reformulate the inference process by incorporating a TTA strategy that enables online parameter adaptation through the minimization of prediction entropy during deployment. Extensive experiments conducted on retinal vessel segmentation benchmarks demonstrate that the proposed method consistently outperforms state-of-the-art approaches, achieving substantial improvements in both segmentation accuracy and robustness against domain variations. Our implementation is available at: https://github.com/plo404/KAN-TransUnet.