Optimizing transformations for contrastive learning in a differentiable framework
This work addresses the need for efficient contrastive learning in medical imaging with limited labeled data, offering an incremental improvement over previous methods by eliminating the need for generative models.
The paper tackles the problem of finding optimal transformations for contrastive learning by proposing a differentiable transformation network, which improves performance in low annotated data regimes. On brain MRI and chest X-ray datasets with 10% labeled data, the model outperforms fully supervised models using 100% labels.
Current contrastive learning methods use random transformations sampled from a large list of transformations, with fixed hyperparameters, to learn invariance from an unannotated database. Following previous works that introduce a small amount of supervision, we propose a framework to find optimal transformations for contrastive learning using a differentiable transformation network. Our method increases performances at low annotated data regime both in supervision accuracy and in convergence speed. In contrast to previous work, no generative model is needed for transformation optimization. Transformed images keep relevant information to solve the supervised task, here classification. Experiments were performed on 34000 2D slices of brain Magnetic Resonance Images and 11200 chest X-ray images. On both datasets, with 10% of labeled data, our model achieves better performances than a fully supervised model with 100% labels.