🤖 AI Summary
To address the high computational cost and memory overhead of Transformer models in medical image semantic segmentation, this work systematically investigates the application of key-value (KV)-only attention mechanisms in both pure Transformer and CNN-Transformer hybrid architectures. By eliminating the query (Q) branch and retaining only the Key-Value pathways, self-attention computation is significantly simplified. Experiments across multiple medical segmentation benchmarks demonstrate that the KV variant reduces model parameters by over 40% and multiply-accumulate operations (MACs) by approximately 35%, while maintaining mIoU performance comparable to standard QKV-based models. This study not only validates the effectiveness of KV attention for pixel-wise dense prediction tasks but also reveals its practical advantages in the accuracy-efficiency trade-off. The proposed approach establishes a novel paradigm for lightweight medical image segmentation, offering a principled path toward resource-efficient yet accurate deep models.
📝 Abstract
While CNNs were long considered state of the art for image processing, the introduction of Transformer architectures has challenged this position. While achieving excellent results in image classification and segmentation, Transformers remain inherently reliant on large training datasets and remain computationally expensive. A newly introduced Transformer derivative named KV Transformer shows promising results in synthetic, NLP, and image classification tasks, while reducing complexity and memory usage. This is especially conducive to use cases where local inference is required, such as medical screening applications. We endeavoured to further evaluate the merit of KV Transformers on semantic segmentation tasks, specifically in the domain of medical imaging. By directly comparing traditional and KV variants of the same base architectures, we provide further insight into the practical tradeoffs of reduced model complexity. We observe a notable reduction in parameter count and multiply accumulate operations, while achieving similar performance from most of the KV variant models when directly compared to their QKV implementation.