Contrastive Flow Matching

TL;DR

Contrastive Flow Matching enhances conditional diffusion models by explicitly enforcing flow uniqueness across conditions, leading to faster training, fewer denoising steps, and improved image quality.

Contribution

It introduces a contrastive objective to flow matching, improving condition separation and overall model performance in conditional diffusion tasks.

Findings

Training speed improved up to 9x

Fewer denoising steps needed (up to 5x fewer)

Lower FID scores by up to 8.9

Abstract

Unconditional flow-matching trains diffusion models to transport samples from a source distribution to a target distribution by enforcing that the flows between sample pairs are unique. However, in conditional settings (e.g., class-conditioned models), this uniqueness is no longer guaranteed--flows from different conditions may overlap, leading to more ambiguous generations. We introduce Contrastive Flow Matching, an extension to the flow matching objective that explicitly enforces uniqueness across all conditional flows, enhancing condition separation. Our approach adds a contrastive objective that maximizes dissimilarities between predicted flows from arbitrary sample pairs. We validate Contrastive Flow Matching by conducting extensive experiments across varying model architectures on both class-conditioned (ImageNet-1k) and text-to-image (CC3M) benchmarks. Notably, we find that training models with Contrastive Flow Matching (1) improves training speed by a factor of up to 9x, (2) requires up to 5x fewer de-noising steps and (3) lowers FID by up to 8.9 compared to training the same models with flow matching. We release our code at: https://github.com/gstoica27/DeltaFM.git.