CONSULT: Contrastive Self-Supervised Learning for Few-shot Tumor Detection

TL;DR

A self-supervised training scheme for anomaly detection, enhancing model performance and data reliability, and a contrastive loss, Tritanh Loss, stabilizes learning by offering a unique solution all while enhancing gradient flow.

Abstract

Artificial intelligence aids in brain tumor detection via MRI scans, enhancing the accuracy and reducing the workload of medical professionals. However, in scenarios with extremely limited medical images, traditional deep learning approaches tend to fail due to the absence of anomalous images. Anomaly detection also suffers from ineffective feature extraction due to vague training process. Our work introduces a novel two-stage anomaly detection algorithm called CONSULT (CONtrastive Self-sUpervised Learning for few-shot Tumor detection). The first stage of CONSULT fine-tunes a pre-trained feature extractor specifically for MRI brain images, using a synthetic data generation pipeline to create tumor-like data. This process overcomes the lack of anomaly samples and enables the integration of attention mechanisms to focus on anomalous image segments. The first stage is to overcome the shortcomings of current anomaly detection in extracting features in high-variation data by incorporating Context-Aware Contrastive Learning and Self-supervised Feature Adversarial Learning. The second stage of CONSULT uses PatchCore for conventional feature extraction via the fine-tuned weights from the first stage. To summarize, we propose a self-supervised training scheme for anomaly detection, enhancing model performance and data reliability. Furthermore, our proposed contrastive loss, Tritanh Loss, stabilizes learning by offering a unique solution all while enhancing gradient flow. Finally, CONSULT achieves superior performance in few-shot brain tumor detection, demonstrating significant improvements over PatchCore by 9.4%, 12.9%, 10.2%, and 6.0% for 2, 4, 6, and 8 shots, respectively, while training exclusively on healthy images.

Figures (4)

• Figure 1: Heatmap Comparison of Unhealthy Images. Different columns show unhealthy MRI images of different positions, sizes, and textures. The first row is the raw image. The second row is the heatmap predicted with PatchCore. The third raw is the heatmap predicted with the proposed method.
• Figure 2: Overall Algorithm for Training Stage. CONSULT takes in $K$ training images and passes them into an augmentation pipeline to generate different representations of healthy and pseudo-unhealthy images. Next, CONSULT trains the backbone through context-aware contrastive learning and self-supervised feature adversarial learning. We use the features extracted from the first, second, and third layers for training. The contrastive learning pipeline helps the model focus on abnormal features by pushing away from the normal features. Furthermore, the normal features are pulled towards each other. The self-supervised feature adversarial learning maximizes the self-similarity of the anchor feature by pulling them closer to each other. We also employ a modified KoLeo Loss to prevent catastrophic collapse to ensure the healthy representation is well spread.
• Figure 3: Architecture of the CONSULT Model. CONSULT uses a ResNet backbone as the feature extractor. We incorporate context-aware attention modules after each ResBlocks and replace the activation function with Leaky ReLU. We choose the first, second, and third feature maps to perform context-aware contrastive learning. Each layer feature is aggregated through resize and concatenation. The context-aware contrastive learning trains the attention module to focus on anomalous regions iteratively, enhancing the robustness of feature extraction. An illustrative attention mechanism visualization is provided, showcasing the model's focus on real anomalous regions within MRI images, derived by averaging the first dimension of the attention maps. On the other hand, the features of the anchor image are passed to the self-supervised feature adversarial learning to maximize the similarity between the normal features while maintaining representative features.
• Figure 4: Comparison between Anchor Loss and Tritanh Loss