Optimizing transformations for contrastive learning in a differentiable framework

TL;DR

This paper introduces a differentiable transformation network to optimize data augmentations for contrastive learning, significantly improving performance in low-data regimes without requiring generative models.

Contribution

It presents a novel framework for learning optimal transformations in contrastive learning using a differentiable network, enhancing accuracy and convergence speed with limited supervision.

Findings

Outperforms fully supervised models with only 10% labeled data

No generative model needed for transformation optimization

Improves convergence speed in low-data regimes

Abstract

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.