Exploring the Integration of Key-Value Attention Into Pure and Hybrid Transformers for Semantic Segmentation
This work addresses the need for efficient models in medical imaging, but it is incremental as it applies an existing KV Transformer method to a new domain.
The paper tackled the computational expense of Transformers in semantic segmentation by evaluating KV Transformers, which reduced parameter count and operations while achieving similar performance to traditional QKV Transformers in medical imaging tasks.
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.