Flow based distribution matching speeds up self supervised image representation learning

Learning a Flow to Self-Supervised Representations

Computer Vision and Pattern RecognitionMachine Learning

Summary

Learning useful representations of images without labeled data is challenging because it requires structuring the way computers understand image features. The authors propose a new method called Flow-Based Distribution Matching (FBDM), which avoids complex adversarial training by using a flow-based model to match distributions. This method assigns multiple augmented views of the same image to shared geometric targets, speeding up training with only a small increase in memory use. Their experiments show that FBDM performs nearly as well as previous methods but trains significantly faster. The paper also explains why representations learned this way can help in downstream classification tasks.

What this means in practice

  • For computer vision engineers: Train image recognition models faster by using FBDM to learn representations without labels but with structured geometric references.
  • For machine learning infrastructure teams: Optimize GPU usage and reduce training time for self-supervised models by replacing adversarial distribution matching with flow-based methods like FBDM.

Authors

Yuling Jiao, Wensen Ma, Houduo Qi, Defeng Sun

Abstract

Explicit geometric references offer a direct way to structure self-supervised representations. Existing adversarial distribution-matching formulations, however, require costly encoder-critic optimization. We introduce Flow-Based Distribution Matching (FBDM), a non-adversarial framework that learns this reference-directed geometry through spherical conditional velocity regression. An ETF-inspired reference allows its number of components K' to exceed the auxiliary flow dimension d* while retaining structured geometric separation. We assign both augmented views of each image to the same target, while limiting how many images each reference center can receive. An explicit alignment loss further pulls the two views' representations closer together. Experiments across benchmarks ranging from CIFAR to ImageNet show that FBDM achieves performance nearly on par with DM and remains competitive with existing SSL methods. Matched training-cost comparisons show a 1.48- to 1.83-fold speedup over DM with a negligible increase in GPU memory usage. We also provide a theoretical explanation for the usefulness of the learned representations: under stated conditions, we bound the downstream misclassification rate in terms of the FBDM pretraining loss.
View PDFOpen arXiv