Self-Supervised Learning for Pre-training Capsule Networks: Overcoming Medical Imaging Dataset Challenges

TL;DR

The paper tackles the challenge of limited high-quality medical imaging data for training capsule networks by investigating self-supervised pre-training on polyp classification. It introduces a modified CapsNet architecture and two SSL auxiliary tasks—colourisation and contrastive learning with in-painting—evaluated on the PICCOLO dataset. Results show that SSL contrastive pre-training with in-painting can achieve performance comparable to or exceeding ImageNet pre-training, with a reported accuracy improvement of about 5.26% over some initialisation strategies, while colourisation underperforms. The work highlights the slow training dynamics and sensitivity of CapsNets to initialisation, suggesting longer training and targeted learning-rate strategies to maximize SSL gains in medical imaging contexts.

Abstract

Deep learning techniques are increasingly being adopted in diagnostic medical imaging. However, the limited availability of high-quality, large-scale medical datasets presents a significant challenge, often necessitating the use of transfer learning approaches. This study investigates self-supervised learning methods for pre-training capsule networks in polyp diagnostics for colon cancer. We used the PICCOLO dataset, comprising 3,433 samples, which exemplifies typical challenges in medical datasets: small size, class imbalance, and distribution shifts between data splits. Capsule networks offer inherent interpretability due to their architecture and inter-layer information routing mechanism. However, their limited native implementation in mainstream deep learning frameworks and the lack of pre-trained versions pose a significant challenge. This is particularly true if aiming to train them on small medical datasets, where leveraging pre-trained weights as initial parameters would be beneficial. We explored two auxiliary self-supervised learning tasks, colourisation and contrastive learning, for capsule network pre-training. We compared self-supervised pre-trained models against alternative initialisation strategies. Our findings suggest that contrastive learning and in-painting techniques are suitable auxiliary tasks for self-supervised learning in the medical domain. These techniques helped guide the model to capture important visual features that are beneficial for the downstream task of polyp classification, increasing its accuracy by 5.26% compared to other weight initialisation methods.

Figures (4)

• Figure 1: Modified capsule network
• Figure 2: Outputs from our modified CapsNet architecture trained by combining SSL contrastive and in-painting auxiliary tasks. The top and middle rows show pairs of differently augmented versions of the same images. The bottom row demonstrates the network's in-painting capability after 200 epochs of training.
• Figure 3: t-SNE visualisation of feature embeddings from our modified CapsNet model after SSL contrastive learning and in-painting auxiliary task. Each dot represents an image projected into 2D space, with matching numeric IDs and colours indicating pairs of the same image.
• Figure 4: Outputs of our modified CapsNet after 200 epochs of training on SSL colourise task