Reweighted Flow Matching via Unbalanced Optimal Transport for Long-tailed Generation
Hyunsoo Song · Minjung Gim · Jaewoong Choi
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.
Chat is not available.
Successful Page Load