Contrastive Beckmann Maps Improve One-step Generation

Beckmann Transport Models (BTM) recently opened a promising avenue for learning one-step generative maps end-to-end, bypassing the need for distillation, intermediate targets, or external features. However, its quality has so far lagged behind state-of-the-art methods, which largely depend on condition strengthening techniques, such as classifier-free guidance for flows and diffusions: guidance combines velocity or score predictions, and a native map has neither. We address this limitation by introducing Contrastive Beckmann Maps (CBM). We observe that the object defining a BTM, its probability current, is linear in the target distribution, so guidance can be applied to the current itself. CBM subtracts the current toward mismatched image–label pairs, obtained by pairing each label with another image in the batch. Under the assumptions of our analysis, the ideal map avoids the negative region of the contrastive target, while sampling still requires one network call. Contrary to CFG, CBM is a theoretically grounded training technique tailored to the BTM framework, repelling correct samples from negatively drawn samples during the training phase. On ImageNet $256\times 256$, fine-tuning a converged BTM model with CBM for around 15% of an epoch lowers one-step FID from 20.8 to 6.7 and 5.5 with postprocessing. Trained from scratch, CBM reaches the same FID of an optimized BTM with $3.5\times$ fewer epochs. On CIFAR-10, CBM improves one-step FID from 7.5 to 5.3.