One-Shot Data Selection for Medical Image Classification via Graph Coverage

2026-06-20

AI Summary

This work addresses the inefficiency of existing active learning and geometric selection methods in medical image classification under limited annotation budgets, which often suffer from repeated model retraining or neglect of local manifold structure. The authors propose a training-free, one-shot sample selection approach that constructs a k-nearest neighbor graph using embeddings from a frozen pretrained model and introduces a biterm heat diffusion coverage kernel incorporating both direct and two-hop neighborhood relationships. A greedy facility location algorithm then selects a high-quality subset that simultaneously ensures class balance and manifold coverage. Requiring only a single hyperparameter and leveraging sparse matrices for efficient full-spectrum heat kernel approximation, the method achieves state-of-the-art balanced accuracy on five MedMNIST datasets across nine out of ten sampling ratios, with particularly pronounced gains over dynamic-training and geometric baselines in class-imbalanced scenarios.

Abstract

Training medical image classifiers on entire datasets is wasteful when annotation budgets are limited: not all samples contribute equally, yet acquiring expert labels is expensive. Active learning reduces annotation cost through iterative querying, but assumes repeated access to an oracle and requires multiple rounds of model training. One-shot geometry-based methods such as facility location avoid retraining but operate on pairwise distances that ignore the local structure of the data manifold. We propose a graph-based one-shot selection method that operates entirely on frozen foundation model embeddings. Given embeddings from a pretrained encoder, we construct a k-nearest neighbor graph over all training samples and derive a two-term coverage kernel from the heat diffusion kernel, capturing both direct and two-hop neighborhood relationships. Greedy facility location on this kernel selects class-balanced subsets that maximize coverage of the data manifold. The two-term kernel matches the full spectral heat kernel in selection behavior while reducing computation to sparse matrix operations with a single hyperparameter. We evaluate on five MedMNIST datasets spanning histopathology, radiology, and microscopy, comparing against both training-dynamics and geometry-based baselines. Our method achieves the highest balanced accuracy on nine of ten dataset-ratio conditions, with the largest gains on class-imbalanced datasets where global graph construction captures cross-class structure that per-class methods miss, all without any model training during selection. Code is available at https://github.com/zahiriddin-rustamov/graph-coverage-selection.

Problem

• one-shot selection
• medical image classification
• data selection
• annotation budget
• class imbalance

Innovation

• graph-based selection
• one-shot data selection
• heat diffusion kernel
• foundation model embeddings
• medical image classification