Ali Farhadi2 · Ranjay Krishna2 · Judy Hoffman1
1 Georgia Tech 2University of Washington
*Equal Contribution
[arXiv]
Summary: We propose Contrastive Flow Matching, a simple extension to the flow matching objective which reduces training iterations and denoising steps while improving generation quality.
This repository is heavily built off of the REPA repository. Nearly all our files are the same as theirs, with the exception of a few scripts & modifications.
conda env create -f environment.yml
conda activate repaPlease follow the setup described in REPA. We mirror their preprocessing procedures.
Training with Contrastive Flow Matching is straightforward.
accelerate launch train.py \
--report-to=wandb \
--allow-tf32 \
--mixed-precision=fp16 \
--seed=0 \
--path-type=linear \
--prediction=v \
--weighting=uniform \
--enc-type=dinov2-vit-b \
--output-dir=[CHECKPOINT_DIR] \
--data-dir=[YOUR_DATA_PATH] \
--model=SiT-XL/2 \
--enc-type=dinov2-vit-b \
--encoder-depth=8 \
--batch-size 256 \
--proj-coeff=0.5 \
--loss-type "contrastive" \ # Specify the contrastive objective. Use "mean" to specify the standard flow matching loss
--contrastive-weight=0.05 \ # Specify the penalty (referred to as lambda in our work) to apply on the contrastive term
--resolution 256This script will automatically generate an experiment name based on the arguments passed, and save logs and checkpoints to the "output-dir". All remaining arguments match those of REPA. To run ImageNet512x512, please follow the procedure in REPA, and simply assign the "contrastive" argument to "--loss-type"
Our generation procedure is nearly equivalent to that of REPA, with a couple minor changes. Specifically, we alter the generation script to account for interference between CFG and Contrastive Flow Matching. To evaluate our models with CFG, please use the following script:
torchrun --nnodes=1 --nproc_per_node=8 --master-port 29501 generate.py \
--model SiT-XL/2 \
--num-fid-samples 50000 \
--ckpt [MODEL_CHECKPOINT_PATH] \
--path-type=linear \
--encoder-depth=8 \
--projector-embed-dims=768 \
--per-proc-batch-size=64 \
--mode=sde \
--num-steps=50 \
--cfg-scale=1.85 \
--guidance-high=0.65 \
--sample-dir [SAMPLES_SAVE_DIR] \
--interference-path /interference_vectors/imnet256_interference_vector.pt \ # path to computed interference vector following Section 5.4 of our paper
--interference-weight 0.05 \ # Proportion of interference to remove from model
--velocity-weight 0.95 \ # Rescaling weight for flow matching model
--reduce-interference # Apply interference reduction. Remove, to not apply itTo evaluate without CFG, you may use:
torchrun --nnodes=1 --nproc_per_node=8 --master-port 29501 generate.py \
--model SiT-XL/2 \
--num-fid-samples 50000 \
--ckpt [MODEL_CHECKPOINT_PATH] \
--path-type=linear \
--encoder-depth=8 \
--projector-embed-dims=768 \
--per-proc-batch-size=64 \
--mode=sde \
--num-steps=50 \
--cfg-scale=1.0 \
--sample-dir [SAMPLES_SAVE_DIR]We have uploaded our REPA SiT-XL/2 models at 256x256 and 512x512 resolutions here.
@misc{stoica2025contrastive,
title={Contrastive Flow Matching},
author={George Stoica and Vivek Ramanujan and Xiang Fan and Ali Farhadi and Ranjay Krishna and Judy Hoffman},
year={2025},
eprint={2506.05350},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2506.05350},
}