Transformer UNet with super token attention for medical image segmentation

Accurate segmentation of medical images requires a careful balance between leveraging local features and addressing long-range dependencies. Transformers, through their self-attention mechanisms, effectively model global contexts and long-range dependencies, overcoming the spatial limitations of convolutional neural networks. However, a major drawback of Transformers is their quadratic computational complexity relative to the number of visual tokens, which creates significant challenges when processing high-fidelity medical image segmentation. In this study, we propose Super Token Attention UNet, a novel framework designed to address these issues. Our method introduces a super token space, where similar visual tokens are grouped, and each group is represented by a single super token that aggregates its information. This strategy allows efficient global context computation without sacrificing accuracy. However, the token aggregation process within the U-shaped encoder–decoder architecture may introduce semantic inconsistency between hierarchical features. To address this, we leverage semantic consistency regularization and internal feature distillation techniques to mitigate semantic loss and reduce feature redundancy, thereby ensuring stable and accurate feature representations across network layers. Experimental results across four medical imaging datasets demonstrate that our method achieves state-of-the-art segmentation performance while maintaining low model complexity, offering an efficient solution for high-fidelity medical image analysis.