Reweighted Flow Matching via Unbalanced Optimal Transport for Long-tailed Generation

Hyunsoo Song, Minjung Gim, Jaewoong Choi

Keywords: Flow Matching, Long-tailed Distribution, Unbalanced Optimal Transport, Optimal Transport

TL;DR: We propose a flow matching model based on Unbalanced Optimal Transport for Long-tailed distribution.

Abstract: Flow matching has recently emerged as a powerful framework for continuous-time generative modeling. However, when applied to long-tailed distributions, standard flow matching suffers from majority bias, oversampling majority modes while generating minority modes with low fidelity. In this work, we propose UOT-Reweighted Flow Matching (UOT-RFM), which leverages Unbalanced Optimal Transport (UOT) to estimate an unsupervised majority score for each target data. Using this score, we correct bias via inverse weighting and introduce higher-order corrections ($k>1$) to further emphasize minority modes. We establish a bias correction theorem, showing that first-order weighting exactly recovers the target distribution. We show that UOT-RFM outperforms existing flow-matching baselines by improving diversity and fidelity on synthetic long-tail data and CIFAR-10-LT.

Submission Number: 141