A lightweight vision transformer with weighted global average pooling: implications for IoMT applications

Abstract

Vision Transformers (ViTs) have garnered significant interest for analysing medical images in Internet of Medical Things (IoMT) systems due to their ability to capture global context. However, deploying ViTs in resource-constrained IoMT environments requires addressing the challenge of adapting these computationally intensive models to meet device limitations while maintaining efficiency. To tackle this issue, we introduce LightAMViT, a lightweight attention mechanism-enhanced ViT, which incorporates K-means clustering layers to reduce the computational complexity of the self-attention matrix, along with an optimized global average pooling layer that leverages all stacked attention block outputs, each weighted by learnable parameters. Additionally, it employs an adaptive learning strategy that facilitates faster convergence by dynamically adjusting the learning rate. We evaluate the proposed technique on two medical image datasets: BUSI and ISIC2020. Our model outperforms conventional CNNs and demonstrates competitive performance compared to the original ViTs, showcasing improvements in both accuracy and computational efficiency. These findings indicate the model’s robustness and generalisation across various medical image analysis tasks, thereby enhancing the applicability of ViTs in resource-limited IoMT devices.