Foundation Model Makes Clustering A Better Initialization For Cold-Start Active Learning

TL;DR

Cold-start initialization in active learning for medical imaging is fragile when using random sampling or naive clustering. The paper proposes a foundation-model-based clustering approach that uses self-supervised embeddings from TXRV, CXRF, and REMEDIS (with ImageNet DenseNet as a baseline) to select informative initial samples, followed by uncertainty-driven subsequent learning. Experiments on ChestX-Det pneumothorax classification and segmentation show that foundation-model embeddings can improve initialization and improve downstream performance compared with baselines, though performance varies by model. This work highlights a practical paradigm for robust cold-start active learning in clinical imaging and emphasizes the importance of embedding quality for clustering-based initialization.

Abstract

Active learning selects the most informative samples from the unlabelled dataset to annotate in the context of a limited annotation budget. While numerous methods have been proposed for subsequent sample selection based on an initialized model, scant attention has been paid to the indispensable phase of active learning: selecting samples for model cold-start initialization. Most of the previous studies resort to random sampling or naive clustering. However, random sampling is prone to fluctuation, and naive clustering suffers from convergence speed, particularly when dealing with high-dimensional data such as imaging data. In this work, we propose to integrate foundation models with clustering methods to select samples for cold-start active learning initialization. Foundation models refer to those trained on massive datasets by the self-supervised paradigm and capable of generating informative and compacted embeddings for various downstream tasks. Leveraging these embeddings to replace raw features such as pixel values, clustering quickly converges and identifies better initial samples. For a comprehensive comparison, we included a classic ImageNet-supervised model to acquire embeddings. Experiments on two clinical tasks of image classification and segmentation demonstrated that foundation model-based clustering efficiently pinpointed informative initial samples, leading to models showcasing enhanced performance than the baseline methods. We envisage that this study provides an effective paradigm for future cold-start active learning.

Figures (1)

• Figure 1: Schematic comparison between the proposed and two baseline initialization approaches for cold-start active learning. The green blocks depict the shared modules while the blue blocks highlight the difference.