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

Reweighted Flow Matching via Unbalanced Optimal Transport for Long-tailed Generation Hyunsoo Song National Institute for Mathematical Sciences [email protected] Minjung Gim National Institute for Mathematical Sciences [email protected] Jaewoong Choi∗ Sungkyunkwan University, MODULABS [email protected] Abstract Flow matching has recently emerged as a powerful framework for continuoustime 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 UOTReweighted 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. 1 Introduction Generative modeling addresses the problem of approximating a target data distribution. Deep generative models have achieved remarkable progress in recent years, such as GANs [ 1 , 15 ], optimal transport maps [ 7 , 9 , 28 ], and diffusion models [ 17 , 30 ]. Among them, flow matching models [ 22 ] have emerged as a promising approach for continuous-time generative models. Flow matching learns a continuous normalizing flow [ 4 ], i.e., a vector field describing the dynamics between an initial prior distribution and the target distribution, while avoiding costly numerical likelihood estimation. Flow matching model is trained through regression to conditional vector field, constructed from conditional probability path between prior and target samples. Despite these advances, flow matching models remain vulnerable to majority bias when trained on long-tailed distributions. In many real-world datasets, data often follow long-tailed or imbalanced distributions, where a few classes dominate while minority classes are severely underrepresented [ 3 , 27 , 33 ]. In such cases, standard flow matching tends to overfit the majority mode due to its regression-based learning nature, while undersampling or exhibiting low fidelity for the minority mode. This limitation reduces diversity and degrades the quality of samples from rare classes. To overcome these challenges, we propose the flow matching model based on the Unbalanced Optimal Transport (UOT) [ 6 , 21 ]. A key property of UOT is that it naturally produces a density ratio between the target distribution and the UOT marginal, which we call the majority score. Building on this, we propose UOT-Reweighted Flow Matching (UOT-RFM), which corrects bias via inverse weighting ∗Corresponding Author 39th Conference on Neural Information Processing Systems (NeurIPS 2025) Workshop: Reliable ML from Unreliable Data. and allows high-order corrections ( k > 1 ) to further emphasize tail samples. Our method provides an unsupervised bias-correction mechanism and enhances coverage of long-tailed data. Our experiments on the CIFAR-10-LT benchmark demonstrate that our method outperforms existing flow matching baselines. Our contributions can be summarized as follows: • We propose UOT-RFM, which leverages UOT couplings and the majority score for unsupervised bias correction. • We establish a bias correction theorem, proving that first-order inverse weighting with the majority score recovers the true target distribution. • Our experiments show that our method achieves improved performance on long-tailed data generation and offers a principled trade-off between majority and minority emphasis with higher-order correction. 2 Preliminaries Flow Matching Continuous Normalizing Flows (CNFs) [ 4 , 22 ] model the dynamics of the probability densities through a probability density path p(t, x) : [0,1] ×Rd7→ R≥0 which transports the initial or source distribution (e.g., Gaussian distribution) p0 to the target data distribution p1 . Specifically, the CNF model is defined by the following Ordinary Differential Equation (ODE), governed by a parameterized vector field vθ: [0,1] ×Rd7→ Rd , i.e., dx dt=vθ t(x) , where we use the notation vt(x) interchangeably with v(t, x) . Then, the associated flow map ϕt(x) denotes the solution of this ODE with initial condition ϕ0(x) = xand the density at time tis given by pt= (ϕt)#p0. Lipman et al. [22] proposed flow matching, a scalable method for training CNFs. The idea is to train the CNF by minimizing a regression loss LFM(θ) between the parameterized vector field vθ t and the ground-truth vector field ut that generates the probability path pt . However, a major challenge is that the marginal ground-truth vector field utis intractable. LFM(θ) = Et∼U,xt∼pt(xt)∥vθ(t, xt)−ut(xt)∥2 2.(1) To overcome this, the flow matching [ 22 , 31 ] introduces a conditional flow matching. Instead of matching ut , the model is trained against the tractable conditional vector field ut(xt|z) , which generates a conditional probability path pt(xt|z) , where z denotes sample pairs ( x0,x1 ). The sample pairs (x0,x1) follow the joint distribution (couplings) of π(z) = π(x0,x1) . The training objectives are given by LCFM(θ) = Et∼U,z∼π(z),xt∼pt(xt|z)∥vθ(t, xt)−ut|z(xt|z)∥2 2.(2) CFM replaces the intractable marginal vector field with a tractable conditional one based on couplings. In particular, the conditional probability path pt(xt|z) and the associated conditional vector field ut(xt|z)can be defined as follows [31]: pt(xt|z) = Nx|tx1+ (1 −t)x0|σ2, ut(xt|z) = x1−x0(3) for some bandwidth hyperparameter σ > 0 . In this case, the marginal probability path and the marginal vector field that generates this path are given by pt(xt) = Zpt(xt|z)π(z)dz, ut(xt) := Eπ(z) ut(x|z)pt(x|z) pt(x)=Ept(z|xt)[ut(xt|z)] (4) Initial Coupling in Flow Matching A key component in the training flow matching model is the choice of initial sample couplings π(z) = π(x0,x1) .This coupling determines how the flow matching model is trained, because the obtained model vt(xt)≈ut(xt) relies on aggregating the conditional vector field over paired samples pt(z|xt) (Eq. 4). The original flow matching framework [ 22 ] employs an independent coupling between the source and target distributions. However, such independence often leads to curved trajectories, which arises from the mean-shift phenomenon due to the flow crossing problem [ 20 , 24 ]. These curved trajectories result in increased numerical errors in ODE simulation and thereby high computational costs in sampling [23]. To improve couplings, recent works adopted the Optimal Transport (OT) approaches between minibatches [26, 31]. Note that the Kantorovich formulation of the Optimal Transport is given by Cot(µ, ν) := inf π∈Π(µ,ν)ZX×Y c(x,y)dπ(x,y).(5) 2 Here, the optimal coupling π⋆ is defined as the minimizer of the transport cost c(x, y) between empirical measures of minibatches from the source samples x0 and target samples x1 . Alternatively, Rectified flow [ 23 ] proposed that leverages pretrained flow models to improve couplings. In this approach, the trajectories are iteratively refined using the previous model as the initial coupling, resulting in straighter paths. 3 Method In this section, we present our model, called UOT-Reweighted Flow Matching (UOT-RFM), that addresses the majority bias of existing flow matching approaches on long-tailed distributions. Our model leverages minibatch Unbalanced Optimal Transport coupling, which naturally provides a majority score for each sample. Intuitively, we compensate for majority bias by over-correcting each target data utilizing this score. In Sec 3.1, we introduce the Uabalanced Optimal Transport problem. In Sec 3.2, we introduce our UOT-RFM model. 3.1 Unbalanced Optimal Transport We introduce the Unbalanced Optimal Transport problem [ 6 , 21 ] and its key properties, which will be leveraged in our approach. The standard OT problem (Eq. 5) enforces exact transport between the source and target distributions, i.e., π0=µ, π1=ν . However, this strict marginal constraint makes OT sensitive to outliers [ 2 , 7 , 14 , 29 ]. To address these issues, the Unbalanced Optimal Transport problem relaxes this constraint and introduces the divergence penalties on the marginal distributions. Cuot(µ, ν) = inf π∈M+(X×Y)ZX ×Y c(x, y)dπ(x, y) + τ1DΨ1(π0∥µ) + τ2DΨ2(π1∥ν),(6) where we assume c(x, y) = 1 2∥x−y∥2 2 and τ1, τ2>0 control the strength of the marginal matching penalties. Here, M+(X × Y) indicates the set of positive Radon measures on X × Y . The terms DΨ1(π0∥µ) and DΨ2(π1∥ν) are two f -divergences that penalizes deviations of the marginals π0, π1 from the source µand target ν, respectively. Therefore, the optimal UOT coupling πu softly matches µ and ν , i.e., πu 0≈µ and πu 1≈ν . Moreover, the UOT problem can represent exact matching of one marginal by setting the divergence penalty appropriately. Specifically, if Ψi is the convex indicator function ι at {1} , then Dι(πi∥η) = 0 if πi= ηa.s., and ∞otherwise. For example, if Ψ1=ι , we obtain at the source-fixed UOT problem where πu 0=µand πu 1≈ν. 3.2 Proposed Method Majority Score Our method leverages the mini-batch UOT coupling πu and the resulting majority score sτ(·) = (dπu 1/dν) . This score is utilized to address the majority oversampling bias of flow matching models on long-tailed distributions by inversely reweighting each target sample. Intuitively, the optimal UOT coupling πu exhibits distribution error whenever a small increase in DΨ leads to a large decrease in transport cost c(x, y) (Eq. 6). As a result, πu prioritizes matching the majority modes, while down-weighting outlier modes with small mass and large cost. This property explains the robustness of UOT to outliers, as the UOT effectively reduces the influence of outliers [ 2 , 7 , 29 ]. Based on this property, we define the majority score sτ(·) = (dπu 1/dν) as the density ratio in the target space under the source-fixed UOT problem. sτ(y) := dπu,⋆ 1 dν (y)>0(7) Here, τ1 is irrelevant, so we simply set τ=τ2 . Intuitively, the majority score measures how strongly each target sample is emphasized by the UOT coupling. sτ>1 indicates emphasized majority samples, while sτ≪1 correspond to down-weighted outlier samples. Importantly, this weighting is entirely unsupervised, arising from the intrinsic geometry of probability distributions (see Appendix A for details). Proposed Method Our corrected conditional flow matching objective with correction order k≥1 is defined as follows (Algorithm 1): Lours,k(θ) = Et∼U,z∼πu(z),xt∼pt(xt|z)sτ(x1)−k∥vθ(t, xt)−ut|z(xt|z)∥2 2.(8) 3 Table 1: Quantitative results on CIFAR-10-LT. In both settings, the training data is long-tailed. The test data is either long-tailed (LT→LT) or balanced (LT→Balanced). Model LT→LT LT→Balanced FID (↓) Prec (↑) Recall (↑) F1 (↑) FID (↓) Prec (↑) Recall (↑) F1 (↑) I-CFM 14.57 0.67 0.28 0.39 25.46 0.60 0.22 0.32 OT-CFM 17.31 0.71 0.24 0.36 27.51 0.63 0.16 0.26 UOT-CFM 14.25 0.67 0.29 0.41 24.94 0.59 0.23 0.33 Ours 11.03 0.61 0.41 0.49 24.06 0.55 0.38 0.45 Table 2: Ablation study on the correction order k when τ= 2.0 . Reported values are FID scores. Training→Test Correction order kUOT-CFM (k= 0) 1.0 2.0 4.0 6.0 8.0 10.0 LT→LT 13.77 13.41 12.42 11.72 11.37 11.04 14.25 LT→Balanced 25.02 24.60 24.54 24.70 24.76 25.35 24.94 where the conditioning variable z= (x0,x1) . Compared with standard flow matching (Eq. 3), our formulation employs the UOT coupling πu for pairing z and introduce an additional weighting factor sτ(x1)−k that rebalances majority and minority samples. Our method is motivated by the following bias correction theorem (see Appendix B for formal statements and proof): Theorem 3.1 (Informal).Flow matching with UOT coupling generates a biased distribution p1= πu τ,1=ν , which overweights majority modes. UOT-RFM corrects this bias by reweighting with the majority score: with correction order k , it generates p1∝s−k τπu τ,1 . In particular, k= 1 exactly recovers the true target distribution ν. Theorem 3.1 shows that when training a flow matching model with UOT coupling (UOT-CFM, [ 8 ]), the generated distribution p1 is biased, i.e., p1=πu 1=ν . In particular, the distribution πu 1 magnifies the majority modes while suppressing the tail modes. This bias can be corrected by applying inverse weighting with the majority score sτ . Our method extends this idea with over-correction ( k > 1 ), further emphasizing tail samples with sτ(·)<1 . Unlike OT-CFM [ 31 ], which relies on mini-batch OT coupling, our approach provides an unsupervised estimate of the majority score, without requiring class labels [27, 33]. 4 Experiments We evaluate our model on long-tailed distributions. In each experiment, our model is compared with several flow matching baselines: independent coupling (I-CFM), OT coupling (OT-CFM, [ 25 , 31 ]), and UOT coupling (UOT-CFM, [8]). Long-Tailed CIFAR-10 We evaluate our model on CIFAR-10 under two settings, using long-tailed CIFAR-10 (CIFAR-10-LT) [ 3 ] as the training data in both cases. In the first setting (LT → LT), the test set is also CIFAR-10-LT, assessing how well each model fits the long-tailed distribution. In the second setting (LT → Balanced), the test set is the original balanced CIFAR-10 [ 18 ], evaluating whether a model trained on imbalanced data can recover the balanced distribution. This evaluation setup is often adopted in supervised long-tailed learning [ 27 , 33 ]. Performance is measured using FID [16], Precision, Recall, and F1-score [19]. Table 1 reports the quantitative results (see Appendix E for qualitative examples). In both settings, our model outperforms all flow matching baselines. In particular, our model achieves significant improvement in the Recall metric, demonstrating improved coverage of minority modes. Although OTCFM achieves the best precision metric, our model achieves the best F1-score, which comprehensively evaluates the Precision and Recall metrics. Moreover, note that the additional training cost is minimal: UOT-RFM requires only about 7% more time than OT-CFM. 4 Correction Order We conduct an ablation study on the correction order k to examine its impact on performance. Interestingly, the best FID scores are achieved when k > 1 , rather than with the exact correction k= 1 . Moreover, compared to UOT-CFM (UOT-RFM without correction), introducing correction generally improves FID scores. Overall, UOT-RFM remains robust to correction order, outperforming other baseline models for all moderate correction orders 2≤k≤8. Acknowledgements Hyunsoo Song and Minjung Gim was supported by National Institute for Mathematical Sciences (NIMS) grant funded by the Korea government (MSIT) (No. B25810000). Jaewoong was supported by the National Research Foundation of Korea(NRF) grant funded by the Korea government(MSIT) [RS2024-00349646]. We thank the Center for Advanced Computation in KIAS for providing computing resources. This research was supported by Brian Impact Foundation, a non-profit organization dedicated to the advancement of science and technology for all. 5 References [1] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214–223. PMLR, 2017. [2] Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems, 33:12934–12944, 2020. [3] Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-distribution-aware margin loss. In Advances in Neural Information Processing Systems, volume 32, 2019. [4] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018. [5] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816, 2016. [6] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274 (11):3090–3123, 2018. [7] Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semidual formulation of unbalanced optimal transport. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. [8] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024. URL https: //openreview.net/forum?id=2UnCj3jeao. [9] Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. ISSN 28358856. URL https://openreview.net/forum?id=2mZSlQscj3. Featured Certification. [10] Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International conference on machine learning, pages 3186–3197. PMLR, 2021. [11] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021. URL http://jmlr.org/papers/v22/20-451.html. [12] Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. Advances in neural information processing systems, 28, 2015. [13] Thomas Gallouët, Roberta Ghezzi, and François-Xavier Vialard. Regularity theory and geometry of unbalanced optimal transport. arXiv preprint arXiv:2112.11056, 2021. [14] Milena Gazdieva, Jaemoo Choi, Alexander Kolesov, Jaewoong Choi, Petr Mokrov, and Alexander Korotin. Robust barycenter estimation using semi-unbalanced neural optimal transport. In The Thirteenth International Conference on Learning Representations, 2025. [15] Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014. [16] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017. 6 [17] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. [18] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. [19] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019. [20] Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ODE-based generative models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 18957–18973. PMLR, 23–29 Jul 2023. [21] Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new hellinger–kantorovich distance between positive measures. Inventiones mathematicae, 211(3):969–1117, 2018. [22] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2022. [23] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=XVjTT1nw5z. [24] Dogyun Park, Sojin Lee, Sihyeon Kim, Taehoon Lee, Youngjoon Hong, and Hyunwoo J Kim. Constant acceleration flow. Advances in Neural Information Processing Systems, 37:90030– 90060, 2024. [25] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 28100–28127. PMLR, 23–29 Jul 2023. [26] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 28100–28127. PMLR, 23–29 Jul 2023. [27] Yiming Qin, Huangjie Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang. Class-balancing diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18434–18443, 2023. [28] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022. [29] Thibault Séjourné, Gabriel Peyré, and François-Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:2211.08775, 2022. [30] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. [31] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. 7 [32] Adrien Vacher and François-Xavier Vialard. Semi-dual unbalanced quadratic optimal transport: fast statistical rates and convergent algorithm. In International Conference on Machine Learning, pages 34734–34758. PMLR, 2023. [33] Tianjiao Zhang, Huangjie Zheng, Jiangchao Yao, Xiangfeng Wang, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed diffusion models with oriented calibration. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview. net/forum?id=NW2s5XXwXU. 8 A Unbalanced Optimal Transport The classical OT problem assumes an exact transport between two distributions µ and ν , i.e., π0= µ, π1=ν . However, this exact matching constraint results in sensitivity to outliers [ 2 , 29 ] and vulnerability to class imbalance in the OT problem [ 8 ]. To mitigate this issue, a new variation of the OT problem is introduced, called Unbalanced Optimal Transport (UOT) [ 6 , 21 ]. Formally, the UOT problem is expressed as follows: Cuot(µ, ν) = inf π∈M+(X×Y)ZX ×Y c(x, y)dπ(x, y) + DΨ1(π0|µ) + DΨ2(π1|ν),(9) where M+(X × Y) denotes the set of positive Radon measures on X × Y . DΨ1 and DΨ2 represents two f -divergences generated by convex functions Ψi , and are defined as DΨi(πj|η) = RΨidπj(x) dη(x)dη(x) . These f -divergences penalize the discrepancies between the marginal distributions π0, π1 and µ, ν , respectively. Hence, in the UOT problem, the two marginal distributions are softly matched to µ, ν , i.e., π0≈µ and π1≈ν . Intuitively, the UOT problem can be seen as the OT problem between π0≈µ and π1≈ν , rather than between the exact distributions µ and ν [ 7 ]. This flexibility offers robustness to outliers [ 2 ] and adaptability to class imbalance problem between µand ν[8] to the UOT problem. Similar to the standard OT problem, the UOT problem also admits a dual formulation [7, 13, 32]: Cuot(µ, ν) = sup u(x)+v(y)≤c(x,y)ZX −Ψ∗ 1(−u(x))dµ(x) + ZY −Ψ∗ 2(−v(y))dν(y),(10) with u∈ C(X) , v∈ C(Y) where C denotes a set of continuous functions over its domain. Here, f∗ denotes the convex conjugate of f , i.e., f∗(y) = supx∈R{⟨x, y⟩ − f(x)} for f:R→[−∞,∞] . Note that this dual problem conducts maximization over two continuous functions u and v . This dual problem can be simplified into a semi-dual formulation by eliminating u via the optimality condition: Cuot(µ, ν) = sup v∈C ZX −Ψ∗ 1(−vc(x))) dµ(x) + ZY −Ψ∗ 2(−v(y))dν(y),(11) where the c -transform of v is defined as vc(x) = inf y∈Y (c(x, y)−v(y)) . Here, vc corresponds to the optimal ugiven v. Finally, the relationship between the marginals of the optimal UOT plan πu,⋆ and the original source and target distributions can be expressed using the optimal UOT potential v⋆ from the semi-dual problem: Theorem A.1 ([ 7 , 13 , 32 ]).Let v⋆ be a solution of the dual formulation of the UOT problem between the source distribution µ and the target distribution ν . Then, the marginal distributions of the optimal UOT plan πu,⋆ satisfy dπu,⋆ 0(x)=Ψ∗ 1 ′(−v⋆c(x))dµ(x)and dπu,⋆ 1(y)=Ψ∗ 2 ′(−v⋆(y))dν(y)(12) B Proofs of theorem In this section, we provide the proof of our bias correction theorem (Theorem ?? ) from the main text. Our proof builds on three key lemmas for the standard flow matching model, originally established in [22, 31], which we restate here for completeness. Lemma B.1 ([ 31 ], Theorem 3.1).The marginal vector field ut generates the probability path pt(xt) from initial conditions p0(x0). pt(xt) = Zpt(xt|z)π(z)dz, ut(xt) := Eπ(z) ut(x|z)pt(x|z) pt(x)=Ept(z|xt)[ut(xt|z)] (13) Lemma B.2 ([ 31 ], Theorem 3.2).If pt(xt)>0 for all xt∈Rd and t∈[0,1] , then, up to a constant independent of θ,LCFM (Eq. 2) and LFM (Eq.1) are equal, and hence ∇θLFM(θ) = ∇θLCFM(θ).(14) 9