Towards Universal Wasserstein Barycenters through Flow Matching

Eduardo Fernandes Montesuma

cs.LGSep 29, 2026

Towards Universal Wasserstein Barycenters through Flow Matching

Authors: Eduardo Fernandes Montesuma

Organizations: Sigma Nova Paris, France

Abstract

Defining a weighted mean over probability measures under probability metrics is a central tool in probabilistic machine learning. Under the Wasserstein metric, these are called \emph{Wasserstein barycenters}. While most approaches compute barycenters for a fixed weight vector, approximating the whole family of barycenters over the simplex, which we call the \emph{Wasserstein simplex}, remains underexplored. We refer to this problem as \emph{Universal Barycenter Approximation}, and propose \texttt{BaryFM}, a flow matching model transporting the marginal measures into any barycenter in the Wasserstein simplex. Once trained, the network can draw samples from measures in the Wasserstein simplex through an ordinary differential equation. We validate our method on 4 downstream tasks: domain adaptation, generalization, Bayesian posterior aggregation and algorithmic fairness. \texttt{BaryFM} achieves the best average rank among 15 competing methods across 10 domain adaptation benchmarks, matching or surpassing non-universal solvers.

Figures & tables

Figure 1: Wasserstein simplex over the set of Gaussian measures. In this paper, we study how to approximate every barycenter in the family $\mathcal{W}(p_{1:K})$ .

Figure 2: BaryFM architecture. A shared encoder maps

(z,k,\lambda)

into a latent code

h

, which is used by the velocity field estimation branch, and the dual OT branch.

Figure 3: Illustration of the convergence of the fixed-point iterations using the neural barycentric map in equation 15 .

Panel of Figure 4 — Figure 4: Sampling of pairs $(z_{0},z_{1})$ for the CFM training. In each panel, $z_{0}\sim p_{k}$ , $k=1,\cdots,3$ respectively. The amber shaded region corresponds to $\gamma_{\epsilon,\psi}(\cdot|z_{0},k)$ , from which $z_{1}$ is sampled.

Figure 5: Inference. One ODE solve carries each

z\sim p_{k}

(coloured) along the paths shown to

\phi(z,1,k,\lambda)

(dark), landing on the barycenter support (gray).

Panel of Figure 6 — Figure 6: Unsupervised (a, $\mathcal{X}\times\mathcal{Y}=$ ✗), and supervised (b, $\mathcal{X}\times\mathcal{Y}=$ ✓) WBT methodologies.

Figure 7: Comparison of BaryFM with 15 methods, and ablation of BaryFM components. We report the average rank over 10 benchmarks. Detailed results are shown in Table 13 and Appendix G .

Figure 8: Data augmentation strategy for domain generalization. On the left, samples from different domains are encoded into a common latent space. In the center, synthetic samples are generated in the Wasserstein simplex of source domains. On the right, the latent synthetic samples are decoded back into raw data space.

Method	Office 31	Office Home	DomainNet	BCI-CIV-2a	SEED-VIG	Avg. Rank
Backbone	ResNet 50	ResNet 50	ResNet 101	CBraMod	CBraMod	-
Latent Space	VAE			PCA		-
ERM	84.38 ±0.53	69.74 ±0.57	49.57 ±0.24	50.63 ±3.45	58.30 ±2.44	4.60
MixUp	85.66 ±0.69	70.43 ±0.34	49.53 ±0.11	52.37 ±3.11	61.11 ±2.57	3.20
Domain MixUp	84.97 ±0.49	68.86 ±0.82	50.63 ±0.20	52.70 ±2.78	61.04 ±3.15	3.00
BaryFM ${}_{\text{Vertex}}$	84.89 ±0.58	70.55 ±0.47	49.69 ±0.34	54.18 ±3.48	59.25 ±5.63	2.80

Table 1: Downstream performance of data augmentation strategies on domain generalization benchmarks. Results denote

\mu_{\pm\sigma}

over seeds (c.f. Section H ).

Appendix figures & tables48 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Reference	Parametrization	Learned object	# nets	# params	Inner opt.	Amort. sampling	Universal	$\mathcal{X}\times\mathcal{Y}$
Free-support (empirical)
Discrete	( Cuturi & Doucet, 2014 )	Particles	support	—	$nd$	—	✗	✗	✗
DaDiL	( Montesuma et al., 2023 )	Particles	support (atoms)	—	$nd$	—	✗	✗	✓
WGF	( Montesuma et al., 2026a )	Particles	support	—	$nd$	—	✗	✗	✓
Restricted parametric family
GMM-WBT	( Montesuma et al., 2024b )	GMM	means, covariances	—	$md^{2}$	—	✓	✗	✓

Appendix

Table 2: Comparison of different approaches in the state-of-the-art along parametrization, objects, number of neural nets and parameters, optimization strategy, amortization, universality, and handling of joint measures.

		Input		Algorithmic					Output
Method	Reference	Mini-Batch OT	Parametric OT	Entropic OT	Unbalanced OT	Latent Pushforward	Constraint Relaxation	Fixed-Pt. Truncation	Time Discretization
Discrete	( Cuturi & Doucet, 2014 )	✗	✗	✓(✗ $\text{ }\epsilon=0$ )	✗	✗	✗	✓	✗
WGF	( Montesuma et al., 2026a )	✓	✗	✓(✗ $\text{ }\epsilon=0$ )	✗	✗	✗	✗	✓
GMM	( Montesuma et al., 2024b )	✗	✓	✗	✗	✗	✗	✓	✗
NWB	( Fan et al., 2020 )	✗	✗	✗	✗	✓	✗	✗	✗
NWBU	( Fan et al., 2020 , App. B)	✗	✗	✗	✗	✓	✗	✗	✗

Appendix

Table 3: Taxonomy of barycenter solvers biases, along input, algorithmic and output biases.

Method	Conditional Flow Matching	Neural Dual Estimation	Balanced
Discrete Barycenter	✗	✗	✓
NeuralFP	✗	✓	✓
Empirical BaryFM	✓	✗	✓
BaryFM	✓	✓	✓
Semi-unbalanced BaryFM	✓	✓	✗

Appendix

Table 4: Summary of different methods proposed in this paper. The discrete barycenter corresponds to the algorithm of Cuturi & Doucet (2014) , which does not use neither CFM nor the dual OT head.

Panel of Figure 9 — Figure 9: The $K=3$ Swiss-roll marginals.

Figure 10: Wasserstein simplex

\mathcal{W}_{p_{1:3}}

with 4 measures highlighted. Due the location-scatter structure, all measures in

\mathcal{W}(p_{1:3})

are known in closed-form.

Figure 11: BW2-UVP% over the simplex for WGF , NWBU , and BaryFM (ours) in the Swiss Roll toy example. Blue points represent the calculated barycenter by each method, gray points represent points from the ground-truth.

Figure 12: BW2-UVP% as a function of # of velocity field evaluations and inference time.

		Homogeneous				Heterogeneous
Method	Universal	BW2-UVP%	Rank	SW 2	Rank	BW2-UVP%	Rank	SW 2	Rank
Discrete	✗	17.70 ±0.02	16	0.310 ±0.004	17	18.92 ±0.08	17	0.317 ±0.007	17
WGF	✗	20.42 ±0.02	18	0.333 ±0.004	18	20.93 ±0.00	18	0.338 ±0.004	18
TDSB	✗	14.22 ±1.73	14	0.234 ±0.019	15	17.03 ±1.46	15	0.256 ±0.011	14
TSBM	✗	8.09 ±0.03	3	0.179 ±0.007	8	10.34 ±0.07	5	0.202 ±0.007	9
CW2B	✗	11.58 ±0.69	12	0.222 ±0.012	14	16.27 ±1.54	14	0.264 ±0.011	15

Appendix

Table 5: Bayesian Posterior Aggregation results. We report BW2-UVP% and SW 2 over seeds 42, 43 and 44 for each method.

Panel of Figure 13 — Figure 13: BW2-UVP% between $\hat{q}_{\lambda}^{(L_{\text{FP}})}$ , the $\lambda$ -weighted Wasserstein barycenter of marginals, against the ground truth full posterior $q_{0}$ over the simplex $\Delta_{3}$ .

Panel of Figure 14 — Figure 14: Different CFG strategies ( $k$ - CFG and $\lambda$ - CFG ) for Bayesian Posterior Aggregation. Through CFG ( $s\neq 1$ , especially $s\in(1,2]$ our proposed BaryFM improves over the best previous baseline, U-NOT ( Gazdieva et al., 2024 ) .

	BaryFM		BaryFM-U
	–	$+k$ -CFG	–	$+k$ -CFG
Semi-unbalanced penalty $\tau$	–	–	0.2	0.2
Batch size $B$	1 024	4 096	1 024	4 096
$k-$ dropout (Homogeneous/Heterogeneous)	0.0	0.8 / 0.1	0.0	0.7 / 0.2
Guidance scale $s$ (Homogeneous/Heterogeneous)	1.0	2.0 / 1.5	1.0	2.0 / 1.5
Entropic regularisation $\epsilon$	2

Appendix

Table 6: Hyperparameters of the universal BaryFM rows in Table 5 .

Panel of Figure 15 — Figure 15: Sensitivity to batch size $B$ , entropic regularization $\epsilon$ and # of fixed point iterations $L_{\text{FP}}$ .

Dataset	Task ( $Y{=}1$ )	# Samples	# Features	Protected Variable	Groups	# Samples per group	%	$\Pr(Y{=}1\mid S{=}s)$
Adult	income $>50$ K	45,222	5	gender	Male	30,527	67.5	0.312
Adult	income $>50$ K	45,222	5	gender	Female	14,695	32.5	0.114
COMPAS	two-year non-recidivism	5,278	8	race	Afr.-Am.	3,175	60.2	0.477
COMPAS	two-year non-recidivism	5,278	8	race	Caucasian	2,103	39.8	0.609
Credit	no payment default	30,000	22	gender	Female	18,112	60.4	0.792
Credit	no payment default	30,000	22	gender	Male	11,888	39.6	0.758

Appendix

Table 7: Fairness datasets. On each dataset, we highlight the task according to the positive outcome

Y=1

. We further list the demographic groups considered on each dataset, and denote in bold the favored group .

Method	Adult		COMPAS		Credit		NLSY
	Acc	DI	Acc	DI	Acc	DI	Acc	DI
Baseline	80.84 ±0.19	0.509 ±0.014	69.19 ±2.25	0.662 ±0.007	80.77 ±0.28	0.974 ±0.010	83.69 ±0.12	0.840 ±0.003
DWB	78.82 ±0.33	1.019 ±0.013	67.02 ±0.33	1.013 ±0.008	80.02 ±0.03	0.998 ±0.009	81.27 ±0.13	0.960 ±0.006
WGF	78.67 ±0.32	1.030 ±0.029	65.94 ±1.26	0.981 ±0.022	77.03 ±0.70	0.990 ±0.004	81.00 ±0.10	0.998 ±0.004
NormFlow	80.54 ±0.19	0.652 ±0.071	67.42 ±2.32	0.871 ±0.036	80.72 ±0.31	0.984 ±0.014	83.05 ±0.26	0.942 ±0.025
NOT	79.52 ±0.33	1.074 ±0.180	68.09 ±0.33	0.916 ±0.083	80.78 ±0.30	1.002 ±0.014	82.34 ±0.22	1.014 ±0.011

Appendix

Table 8: Performance of total repair, per barycenter solver, across 4 fairness datasets.

Figure 16: Fairness vs. classification accuracy tradeoff. For having a single view of methods performance across different datasets, we use the average rank on each axis of methods across datasets.

Panel of Figure 17 — Figure 17: Disparate Impact (DI) vs. classification accuracy as a function of flow time. At $t=0$ , we have the starting point of samples at each $p_{k}$ . As time evolves, these samples are progressively transported to the barycenter, which usually improves DI (models become more fair) at the cost of accuracy.

Figure 18: Fairness vs. classification accuracy tradeoff for BaryFM under different hyper-parameters.

Dataset	Adaptation Protocol	Backbone	# Features	# Classes	# Domains	# Samples	Seeds
Office 31	LODO	ResNet-50	2048	31	3	4,110	42, 43, 44
Office-Home	LODO	ResNet-101	2048	65	4	15,500	42, 43, 44
DomainNet	LODO	ResNet-101	2048	345	6	586,575	42, 43, 44
BCI-CIV-2a	Fixed Target	CBraMod	200	4	9	5,184	42, 43, 44
FACED	Fixed Target	CBraMod	200	9	123	10,332	42, 43, 44
SEED-VIG	Fixed Target	CBraMod	200	–	21	20,355	42, 43, 44

Appendix

Table 9: Summary of datasets used in our experiments.

Dataset	Hidden width	MLP embedding	$x$	$t$	$y$	$k$	$\lambda$	# Params
Office 31	512	✓	64	64	64	64	64	4.09M
Office-Home	512	✗	128	32	32	64	64	3.24M
DomainNet	512	✗	128	64	128	64	64	3.54M
BCI-CIV-2a	512	✓	64	64	64	64	64	2.18M
FACED	512	✓	128	64	64	64	64	2.36M
SEED-VIG	512	✓	128	64	32	128	64	2.38M

Appendix

Table 10: Architecture of the network used in domain adaptation. Columns

x

t

y

k

and

\lambda

give the embedding dimensions of the features, time, labels, domain index and barycentric weights. The last column gives the total number of trainable parameters, including the dual OT head.

Benchmark	$B$	$N$	$L_{\text{Dual}}$	$L_{\text{FP}}$	LR	Weight Decay	$\varepsilon$
Office 31	310	10 000	1	1	$10^{-3}$	0.01	$10^{-3}$
Office Home	1300	10 000	1	1	$10^{-4}$	0	$10^{-3}$
DomainNet	2048	30 000	1	1	$10^{-4}$	0	$10^{-3}$
BCI-CIV-2a	400	10 000	1	1	$10^{-3}$	0.01	$10^{-3}$
FACED	256	10 000	1	1	$10^{-4}$	0.01	$10^{-3}$
SEED-VIG	1024	10 000	5	5	$10^{-4}$	0.01	$10^{-3}$

Appendix

Table 11: Training hyperparameters of BaryFM across benchmarks.

	Office 31	Office Home	DomainNet	BCI-CIV-2a	FACED	SEED-VIG	Mumtaz	Physio	SHU-MI	TEP
$n_{\text{gen}}$	4 650	39 000	103 500	1 000	9 000	53 100	10 000	10 000	10 000	5 800
ODE steps	100	100	100	100	100	100	100	100	100	100
ODE Solver	Euler	Euler	Euler	Euler	Euler	Euler	Euler	Euler	Euler	Euler
Transport	emd	emd	cfm	emd	emd	emd	emd	emd	emd	emd

Appendix

Table 12: Inference hyperparameters of BaryFM across benchmarks.

Method	$\mathcal{X}\times\mathcal{Y}$	Universal	Office31	OfficeHome	DomainNet	BCI-CIV-2a	FACED	SEED-VIG	Mumtaz	Physio	SHU-MI	TEP	Avg. Rank
Backbone	-	-	ResNet-50	ResNet-101		CBraMod						CNN	-
Source-Only	-	-	86.58 ±0.07	76.93 ±0.09	51.23 ±0.01	52.72 ±0.05	54.87 ±0.27	44.72 ±2.15	91.25 ±0.36	64.10 ±0.08	62.68 ±0.41	79.18 ±0.18	12.10
Discrete ✗	✗	✗	85.76 ±0.56	75.95 ±0.05	50.46 ±0.20	62.62 ±0.41	54.47 ±0.32	34.66 ±0.50	91.30 ±0.13	64.95 ±0.29	55.89 ±3.26	84.54 ±1.47	11.45
Discrete ✓	✓	✗	85.46 ±0.36	76.09 ±0.36	51.01 ±0.10	62.50 ±0.63	54.43 ±0.50	34.76 ±0.60	91.30 ±0.13	64.97 ±0.27	53.80 ±4.03	85.23 ±1.26	11.40
WGF ✗	✗	✗	86.69 ±0.33	77.09 ±0.28	49.48 ±0.04	63.02 ±0.23	58.06 ±0.24	52.29 ±0.62	92.79 ±0.41	65.12 ±0.40	62.87 ±0.90	85.52 ±1.48	5.40
WGF ✓	✓	✗	88.90 ±0.74	78.03 ±0.13	52.08 ±0.14	62.56 ±0.13	56.99 ±0.24	51.74 ±0.44	91.77 ±0.56	65.54 ±0.23	63.70 ±0.11	86.87 ±1.27	4.50

Appendix

Table 13: Downstream performance of WBT per barycenter solver. Results report

\mu_{\pm\sigma}

over 3 seeds. Avg. Rank averages the per-dataset ranks.

CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	ResNet-50	ResNet-101		CBraMod						CNN	-
				Office31	OfficeHome	DomainNet	BCI-CIV-2a	FACED	SEED-VIG	Mumtaz	Physio	SHU-MI	TEP	Avg. Rank
✗	✓	✓	✓	88.53 ±0.11	78.61 ±0.02	52.18 ±0.09	64.93 ±0.19	56.78 ±0.10	47.65 ±3.33	91.94 ±0.17	65.15 ±0.09	63.63 ±0.49	87.07 ±1.32	3.50
✓	✗	✓	✓	89.09 ±0.33	78.68 ±0.04	52.54 ±0.12	64.47 ±0.28	59.27 ±0.16	53.52 ±0.49	93.12 ±0.27	64.98 ±0.21	63.61 ±0.66	87.26 ±1.15	1.90
✓	✓	✗	✓	89.03 ±0.18	78.60 ±0.07	52.70 ±0.04	65.22 ±0.53	57.95 ±0.39	52.15 ±0.65	92.05 ±0.48	65.00 ±0.17	63.61 ±0.30	87.32 ±1.20	2.35
✓	✓	✓	✗	87.88 ±0.27	78.65 ±0.09	52.45 ±0.04	64.79 ±0.18	57.82 ±0.23	49.99 ±0.06	92.35 ±0.99	64.89 ±0.17	63.47 ±0.28	87.14 ±1.22	3.50
✓	✓	✓	✓	89.74 ±0.23	78.72 ±0.12	52.56 ±0.07	64.53 ±0.65	58.25 ±0.03	57.42 ±0.94	93.12 ±0.59	65.00 ±0.11	63.70 ±0.15	87.30 ±1.39	1.75

Appendix

Table 14: Architectural ablation of BaryFM . Results report

\mu_{\pm\sigma}

over 3 seeds. The rank of each method is computed by replacing BaryFM in Table 13 ; Avg. Rank averages the per-dataset ranks.

Panel of Figure 19 — Figure 19: Hyper-parameter sensitivity to $n_{\text{gen}}$ , ODE solver, # of fixed-point iterations ( $L_{\text{FP}}$ ), and # of ODE steps ( $T$ ) on Office 31 and TEP benchmarks.

Method	Office31	OfficeHome	DomainNet	BCI-IV-2a	FACED	SEED-VIG	Mumtaz	Physio	SHU-MI	TEP
CW2B	3,957,504	5,936,256	9,893,760	1,023,360	16,373,760	2,660,736	1,023,360	14,327,040	3,070,080	677,760
WIN	3,979,140	5,688,966	9,108,618	1,518,882	23,033,832	3,813,810	1,518,882	20,165,172	4,387,542	1,223,178
NOT	1,643,778	2,465,667	4,109,445	552,045	8,832,720	1,435,317	552,045	7,728,630	1,656,135	413,445
U-NOT	1,643,778	2,465,667	4,109,445	552,045	8,832,720	1,435,317	552,045	7,728,630	1,656,135	413,445
TDSB	2,453,760	3,680,640	6,134,400	2,567,760	41,084,160	6,676,176	2,567,760	35,948,640	7,703,280	2,428,800
TSBM	8,946,688	13,420,032	22,366,720	6,941,520	111,064,320	18,047,952	6,941,520	97,181,280	20,824,560	3,168,000

Appendix

Table 15: Number of trainable parameters of neural barycenter solvers.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	Amazon	dSLR	Webcam	Avg.	Rank
Source-Only	-	-	-	-	64.23 ±0.36	99.41 ±0.51	96.10 ±0.68	86.58 ±0.07	10
Discrete ✗	-	-	✗	✗	70.67 ±0.40	92.26 ±0.52	94.35 ±0.89	85.76 ±0.56	13
Discrete ✓	-	-	✓	✗	69.98 ±0.61	91.67 ±1.36	94.74 ±1.01	85.46 ±0.36	14
WGF ✗	-	-	✗	✗	69.63 ±0.56	93.75 ±0.89	96.69 ±0.34	86.69 ±0.33	9
WGF ✓	-	-	✓	✗	70.33 ±0.44	98.51 ±1.36	97.86 ±0.68	88.90 ±0.74	2
CW2B	-	-	✗	✗	69.05 ±1.22	95.83 ±2.73	96.49 ±0.58	87.12 ±1.27	5

Appendix

Table 16: Per-domain Office 31 results, along average over domains, and rank.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	Art	Clipart	Product	Real World	Avg.	Rank
Source-Only	-	-	-	-	75.55 ±0.25	63.21 ±0.15	83.57 ±0.01	85.40 ±0.07	76.93 ±0.09	4
Discrete ✗	-	-	✗	✗	69.63 ±0.86	65.29 ±0.51	84.12 ±0.64	84.78 ±0.50	75.95 ±0.05	9
Discrete ✓	-	-	✓	✗	70.15 ±1.28	65.73 ±0.22	83.28 ±0.11	85.19 ±0.39	76.09 ±0.36	8
WGF ✗	-	-	✗	✗	75.24 ±1.24	64.76 ±0.89	83.33 ±0.23	85.02 ±0.99	77.09 ±0.28	3
WGF ✓	-	-	✓	✗	75.45 ±0.43	65.06 ±0.11	85.32 ±0.33	86.28 ±0.24	78.03 ±0.13	2
CW2B	-	-	✗	✗	71.07 ±0.82	64.95 ±0.62	82.81 ±1.02	84.89 ±0.63	75.93 ±0.19	10

Appendix

Table 17: Per-domain Office-Home results, along average over domains, and rank.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	Clipart	Infograph	Painting	Quickdraw	Real	Sketch	Avg.	Rank
Source-Only	-	-	-	-	71.75 ±0.02	28.63 ±0.03	58.73 ±0.00	17.34 ±0.01	71.07 ±0.00	59.86 ±0.00	51.23 ±0.01	7
Discrete ✗	-	-	✗	✗	69.61 ±0.38	27.33 ±0.32	55.97 ±0.32	21.33 ±0.63	70.84 ±0.29	57.72 ±0.36	50.46 ±0.20	9
Discrete ✓	-	-	✓	✗	70.81 ±0.58	27.19 ±0.23	56.64 ±0.41	21.28 ±0.41	72.06 ±0.17	58.06 ±0.21	51.01 ±0.10	8
WGF ✗	-	-	✗	✗	67.88 ±0.25	27.44 ±0.06	55.28 ±0.46	20.70 ±0.23	68.72 ±0.04	56.85 ±0.19	49.48 ±0.04	12
WGF ✓	-	-	✓	✗	71.28 ±0.29	29.46 ±0.42	58.02 ±0.45	22.41 ±0.36	72.46 ±0.19	58.88 ±0.09	52.08 ±0.14	5
CW2B	-	-	✗	✗	69.01 ±1.52	26.10 ±0.36	55.56 ±1.17	16.98 ±0.79	69.63 ±0.35	57.56 ±0.74	49.14 ±0.69	13

Appendix

Table 18: Per-domain DomainNet results, along average over domains, and rank.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	S8	S9	Pool.	Rank
Source-Only	-	-	-	-	54.98 ±0.10	50.46 ±0.10	52.72 ±0.05	16
Discrete ✗	-	-	✗	✗	66.72 ±0.44	58.51 ±0.46	62.62 ±0.41	8
Discrete ✓	-	-	✓	✗	66.67 ±0.76	58.33 ±0.97	62.50 ±0.63	10
WGF ✗	-	-	✗	✗	66.90 ±0.20	59.14 ±0.66	63.02 ±0.23	2
WGF ✓	-	-	✓	✗	66.15 ±0.17	58.97 ±0.20	62.56 ±0.13	9
CW2B	-	-	✗	✗	66.38 ±0.27	59.03 ±0.63	62.70 ±0.35	4

Appendix

Table 19: Per-domain BCI-IV-2a cross-subject results, along pooled performance, and rank.

Figure 20: Source-only accuracy vs. method accuracy, for BaryFM

{}_{\text{Emp}}

, BaryFM , WIN and CW2B , per test subject. The bisecting line defines the no-improvement region. Points under this line (e.g., 119), in blue, show subjects with no improved performance. Subjects on top of that line (e.g., 115) show subjects with improved performance.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	S18	S19	S20	S21	Pool.	Rank
Source-Only	-	-	-	-	87.19 ±0.56	7.13 ±7.81	27.26 ±9.24	70.33 ±2.16	44.72 ±2.15	9
Discrete ✗	-	-	✗	✗	82.93 ±0.48	-4.43 ±0.11	40.98 ±0.68	64.49 ±0.29	34.66 ±0.50	16
Discrete ✓	-	-	✓	✗	82.39 ±0.22	-8.97 ±0.44	39.70 ±0.81	64.94 ±0.85	34.76 ±0.60	15
WGF ✗	-	-	✗	✗	78.83 ±0.56	13.64 ±1.34	41.61 ±1.42	70.10 ±0.09	52.29 ±0.62	3
WGF ✓	-	-	✓	✗	78.09 ±1.52	12.28 ±0.14	38.78 ±2.26	70.89 ±1.00	51.74 ±0.44	4
CW2B	-	-	✗	✗	83.97 ±0.48	-20.12 ±1.32	43.87 ±0.50	72.42 ±0.38	51.06 ±0.74	5

Appendix

Table 21: Per-domain SEED-VIG results, along pooled performance, and rank.

Domain	#H	#MDD	# Samples	Healthy subjects	MDD subjects
$p_{1}$	5	5	1135	H 21, 24, 25, 27, 28	MDD 15, 17, 18, 27, 28
$p_{2}$	4	5	990	H 13, 16, 19, 23	MDD 10, 11, 12, 22, 25
$p_{3}$	4	4	906	H 1, 17, 22, 26	MDD 1, 2, 16, 19
$p_{4}$	4	4	941	H 2, 12, 14, 15	MDD 13, 21, 23, 26
$p_{5}$	4	4	919	H 10, 11, 18, 20	MDD 14, 20, 24, 29
Sources (total)	21	22	4891

Appendix

Table 22: Grouping of source domain subjects into different subjects.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	Pool.	Rank
Source-Only	-	-	-	-	91.25 ±0.36	15
Discrete ✗	-	-	✗	✗	91.30 ±0.13	13
Discrete ✓	-	-	✓	✗	91.30 ±0.13	13
WGF ✗	-	-	✗	✗	92.79 ±0.41	2
WGF ✓	-	-	✓	✗	91.77 ±0.56	7
CW2B	-	-	✗	✗	91.88 ±0.19	4

Appendix

Table 23: Mumtaz2016 results, along pooled performance, and rank.

Figure 21: Source-only accuracy vs. method accuracy for BaryFM (5th), WGF (1st), GMM (1st) and NWB (4th). The bisecting line defines the no-improvement region. Points under this line (e.g., S98) denote subjects for which performance is worse than the source-only. Points above it denote subjects for which performance improved over that baseline.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	S21	S22	S23	S24	S25	Pool.	Rank
Source-Only	-	-	-	-	76.97 ±0.21	67.50 ±0.42	67.97 ±1.46	52.86 ±1.43	46.46 ±1.19	62.68 ±0.41	14
Discrete ✗	-	-	✗	✗	57.95 ±6.86	61.53 ±3.34	61.52 ±7.94	48.36 ±3.81	49.06 ±2.61	55.89 ±3.26	15
Discrete ✓	-	-	✓	✗	53.32 ±4.50	59.10 ±4.90	58.57 ±10.66	49.05 ±3.50	48.27 ±1.63	53.80 ±4.03	16
WGF ✗	-	-	✗	✗	76.69 ±1.04	68.12 ±2.29	67.35 ±0.63	53.47 ±0.95	47.19 ±1.15	62.87 ±0.90	13
WGF ✓	-	-	✓	✗	78.08 ±0.24	68.75 ±0.62	69.14 ±0.00	53.70 ±0.95	47.19 ±0.57	63.70 ±0.11	1
CW2B	-	-	✗	✗	76.90 ±0.73	67.43 ±0.24	69.07 ±1.32	53.32 ±0.79	46.90 ±0.98	63.05 ±0.20	7

Appendix

Table 25: Per-domain SHU-MI results, along pooled performance, and rank.

Method	CFM	Dual OT	$\mathcal{X}\times\mathcal{Y}$	Universal	Mode 1	Mode 2	Mode 3	Mode 4	Mode 5	Mode 6	Avg.	Rank
Source-Only	-	-	-	-	82.40 ±0.07	67.03 ±1.03	88.52 ±0.07	78.30 ±0.15	73.51 ±0.13	85.34 ±0.07	79.18 ±0.18	16
Discrete ✗	-	-	✗	✗	89.66 ±0.75	73.53 ±1.35	84.96 ±1.95	88.34 ±0.75	84.91 ±0.89	85.81 ±2.34	84.54 ±1.47	6
Discrete ✓	-	-	✓	✗	89.97 ±1.25	73.11 ±1.54	88.48 ±1.76	88.52 ±0.83	84.46 ±0.84	86.85 ±1.07	85.23 ±1.26	5
WGF ✗	-	-	✗	✗	91.52 ±1.32	73.78 ±1.25	87.44 ±1.17	89.39 ±0.62	84.88 ±1.76	86.09 ±2.23	85.52 ±1.48	3
WGF ✓	-	-	✓	✗	92.34 ±1.05	76.10 ±1.57	89.79 ±0.90	89.32 ±1.04	85.92 ±1.36	87.75 ±1.52	86.87 ±1.27	2
CW2B	-	-	✗	✗	88.76 ±1.51	68.37 ±0.85	85.34 ±1.48	88.17 ±1.00	81.44 ±2.71	85.92 ±0.95	83.00 ±1.55	13

Appendix

Table 26: Per-mode TEP results, along average over modes, and rank.

Figure 22: Visualization of the linear mixing of multiple domain data via Domain MixUp of Wang et al. (2020) . In the first row, we show the simplex, and the picked mixing weights $\lambda$ . In the second row, we show the mixed samples obtained by equation 36 , and decoded back into image space. In the third row, we show the t-SNE visualization ( Maaten & Hinton, 2008 ) of the VAE embeddings of the mixed samples.

Figure 23: Visualization of nonlinear mixing of multiple domain data via BaryFM (ours). In the first row, we show the simplex, and picked mixing weights $\lambda$ . In the second row, we show the decoded images of samples from $\text{Bar}(\lambda,p_{1:K})$ . In the third row, we show a t-SNE visualization ( Maaten & Hinton, 2008 ) of the VAE embeddings of the samples from the barycenter.

Panel of Figure 24 — Figure 24: The effect of autoencoding images in the Office31, Office Home, and Domain Net, using the VAE of Rombach et al. (2022) . On each sub-figure, we show in the left the original images in the dataset $x$ , and the reconstructed images from the latent codes, $\hat{x}=\text{Dec}(\text{Enc}(x))$ .

Figure 25: Architecture of our BaryFM used in the domain generalization experiments.

Dataset	Latent shape	$K$	$C$	Channels	Attention res.	# Params
Office 31	$4\times 32\times 32$	2	31	$(128,256,512,512)$	$\{16,8,4\}$	117.23M
Office-Home	$4\times 32\times 32$	3	65	$(128,256,512,512)$	$\{16,8,4\}$	117.25M
DomainNet	$4\times 32\times 32$	5	345	$(128,256,512,512)$	$\{16,8,4\}$	117.39M

Appendix

Table 27: Architecture of the UNet used for generating the augmented pools in domain generalization.

K

is the number of source domains in each leave-one-domain-out fold, and

C

the number of classes. All networks use 2 residual blocks per level, 4 attention heads, and a conditioning dimension

\text{dim}_{c}=256

. The last column gives the total number of trainable parameters, including the label and dual OT heads.

Panel of Figure 26 — Figure 26: Channel-wise BW2-UVP%, log-spectral distance, MSE and downstream performance as a function of number of principal components for BCI-IV-2a (a – d) and SEED-VIG (e – h).

Hyper-parameter	Office 31	Office-Home	DomainNet	BCI-CIV-2a	SEED-VIG
Latent space	VAE , $4\times 32\times 32$	VAE , $4\times 32\times 32$	VAE , $4\times 32\times 32$	PCA , $2048$	PCA , $4096$
Architecture	UNet	UNet	UNet	MLP	MLP
Optimizer	Adam	Adam	Adam	AdamW	AdamW
Learning rate	$10^{-4}$	$10^{-4}$	$10^{-4}$	$10^{-4}$	$10^{-4}$
Weight decay	–	–	–	0.05	0.05
Learning rate schedule	Cosine, $\text{LR}_{\text{min}}=10^{-6}$	Cosine, $\text{LR}_{\text{min}}=10^{-6}$	Cosine, $\text{LR}_{\text{min}}=10^{-6}$	Cosine, $\text{LR}_{\text{min}}=10^{-6}$	Cosine, $\text{LR}_{\text{min}}=10^{-6}$

Appendix

Table 28: Training and inference hyper-parameters used in our domain generalization experiments.

Method	Amazon	dSLR	Webcam	Avg.	Rank
ERM (ResNet 50)	60.39 ±0.74	98.40 ±0.55	94.34 ±1.33	84.38 ±0.53	5
MixUp	60.74 ±1.08	99.00 ±0.00	97.23 ±1.14	85.66 ±0.69	2
Domain MixUp	58.72 ±1.49	99.20 ±0.84	96.98 ±1.03	84.97 ±0.49	3
BaryFM ${}_{\text{Vertex}}$	60.04 ±1.08	98.80 ±0.45	95.85 ±1.14	84.89 ±0.58	4
BaryFM ${}_{\text{Random}}$	62.20 ±1.31	99.40 ±0.89	96.23 ±1.18	85.94 ±0.66	1

Appendix

Table 29: Per domain Office 31 domain generalization results. Results denote classification accuracy (in %) over 5 seeds.

Method	Art	Clipart	Product	Real World	Avg.	Rank
ERM (ResNet 50)	65.68 ±1.19	53.91 ±1.10	79.37 ±1.16	80.00 ±0.69	69.74 ±0.57	4
MixUp	65.56 ±0.94	57.54 ±1.43	78.99 ±0.69	79.64 ±0.55	70.43 ±0.34	3
Domain MixUp	63.12 ±2.06	55.25 ±1.83	78.30 ±1.16	78.78 ±0.75	68.86 ±0.82	5
BaryFM ${}_{\text{Vertex}}$	66.43 ±1.56	55.74 ±1.00	80.15 ±0.79	79.87 ±0.82	70.55 ±0.47	2
BaryFM ${}_{\text{Random}}$	67.18 ±1.34	57.49 ±1.11	80.17 ±0.92	79.82 ±0.38	71.17 ±0.76	1

Appendix

Table 30: Per domain Office Home domain generalization results. denote classification accuracy (in %) over 5 seeds.

Method	Clipart	Infograph	Painting	Quickdraw	Real	Sketch	Avg.	Rank
ERM (ResNet 101)	69.97 ±0.33	28.34 ±0.94	58.34 ±1.20	13.61 ±0.07	70.13 ±0.34	57.05 ±0.71	49.57 ±0.24	4
MixUp	69.71 ±0.21	28.24 ±0.45	58.99 ±0.66	13.50 ±0.59	69.59 ±0.60	57.19 ±0.34	49.53 ±0.11	5
Domain MixUp	71.00 ±0.27	28.91 ±0.54	59.53 ±1.01	14.97 ±0.30	70.93 ±0.60	58.43 ±0.40	50.63 ±0.20	1
BaryFM ${}_{\text{Vertex}}$	69.72 ±0.39	28.25 ±0.81	58.79 ±0.82	13.43 ±0.09	70.57 ±0.68	57.40 ±1.18	49.69 ±0.34	3
BaryFM ${}_{\text{Random}}$	70.57 ±0.18	28.68 ±0.47	59.83 ±0.84	13.83 ±0.56	70.25 ±0.72	57.99 ±0.47	50.19 ±0.21	2

Appendix

Table 31: Per domain DomainNet domain generalization results. Results denote classification accuracy (in %) over 3 seeds.

Method	S8	S9	Pool.	Rank
ERM (CBraMod)	50.06 ±4.84	51.20 ±2.78	50.63 ±3.45	5
MixUp	51.04 ±3.30	53.70 ±3.36	52.37 ±3.11	4
Domain MixUp	50.89 ±3.02	54.51 ±2.90	52.70 ±2.78	3
BaryFM ${}_{\text{Vertex}}$	53.75 ±4.06	54.62 ±3.76	54.18 ±3.48	1
BaryFM ${}_{\text{Random}}$	52.71 ±3.70	54.76 ±3.91	53.73 ±3.43	2

Appendix

Table 32: Per domain BCI-IV-2a domain generalization results. denote classification accuracy (in %) over 5 seeds.

Method	S14	S15	S16	S17	Pool.	Rank
ERM (CBraMod)	55.57 ±4.05	33.58 ±5.75	80.43 ±3.16	62.29 ±7.41	58.30 ±2.44	5
MixUp	53.46 ±2.91	39.40 ±4.83	80.79 ±2.19	59.13 ±6.94	61.11 ±2.57	2
Domain MixUp	57.46 ±2.61	40.36 ±5.49	82.77 ±2.28	68.12 ±6.10	61.04 ±3.15	3
BaryFM ${}_{\text{Vertex}}$	53.62 ±2.43	47.04 ±4.94	79.04 ±2.09	59.13 ±5.07	59.25 ±5.63	4
BaryFM ${}_{\text{Random}}$	54.93 ±2.84	43.18 ±5.95	80.68 ±2.33	63.45 ±6.31	62.08 ±2.38	1

Appendix

Table 33: Per domain SEED-VIG domain generalization results. Results denote Pearson correlation (in %) over 5 seeds.

Explore similar work

May 13, 2026cs.LG

Path-independent Flow Matching for Multi-parameter Generative Dynamics

Flow Matching is a powerful framework for learning transport maps between probability distributions. Yet its standard single-parameter formulation is not designed to capture multi-parameter variations where the resulting transport should be path-independent. Path independence is crucial because it ensures that transformations depend only on the initial and target distributions, not on the specific path. In this work, we introduce Path-independent Flow Matching (PiFM), a method for learning vector fields whose induced flows yield path-independent transport between distributions. We show that PiFM generalizes Flow Matching to higher-dimensional parameter domains while enforcing structural conditions that ensure consistency of composed transformations. In addition, we show that, under suitable assumptions, PiFM approximates the Wasserstein barycenter, linking the framework to a notion of distributional interpolation. To enable practical training, we propose a tractable, simulation-free objective that regresses onto multi-parameter conditional probability paths. We showcase empirically that PiFM outperforms other approaches on both synthetic and real world data in interpolating path-independent trajectories and generating desired out of distribution samples.

Francisco Téllez, AmirHossein Zamani, Philippe Martin +5

Université de Montréal, Mila · Concordia University, Mila · Concordia University +2

Conditional Flow Matching Latent Transport Vector Fields Interpolation

Ask PDF

May 8, 2026cs.LG

Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once

Flow matching has recently emerged as a flexible and efficient framework for generative modelling by learning deterministic transport dynamics between probability measures. In this work, we extend flow matching to the space of probability measures over probability measures, introducing a Wasserstein-on-Wasserstein (WoW) formulation. Leveraging the nested Wasserstein geometry, we show that measures over transport plans naturally induce velocity fields that realize metameasure flows. This yields a principled generalization of Wasserstein flow matching via coupled outer and inner transport plans. To address the substantial computational cost of WoW transport, we propose scalable approximations based on sliced and linear Wasserstein distances, enabling efficient training while promoting numerically stable, near-straight trajectories. Our framework unifies and extends existing approaches to point cloud and set generation, providing a practical and theoretically grounded method for generative modelling in WoW spaces.

Moritz Piening, Richard Duong, Gabriele Steidl

Institut für Mathematik, Technische Universität Berlin, Germany

Velocity Field Generative Models Point Clouds Transport

Ask PDF

Aug 6, 2026cs.LG

Potential Matching Optimal Transport: Continuous Normalizing Flows for Exact $p$ -Wasserstein Dynamics

We introduce Potential Matching Optimal Transport (PMOT), a potential-flow framework for general

p

-cost optimal transport with

c_p(x,y)=\|x-y\|^p

. PMOT parameterizes the CNF velocity field with a scalar potential in the generalized Benamou--Brenier form for the chosen exponent

p

. It trains the potential gradient with a self-induced matching loss along straight bridges determined by the model's own endpoints, while allowing flexible terminal distribution matching. Our main result establishes zero-loss exactness: under the stated regularity, exact terminal matching, and uniqueness assumptions, any zero-loss solution satisfies the generalized Benamou--Brenier optimality system and recovers the corresponding

p

-optimal transport map and dynamics. On synthetic benchmarks, PMOT learns

p

-specific maps that agree with the corresponding

p

-matched OT references. It also remains competitive as a likelihood-based density model on high-dimensional tabular data, and an MMD-based color transformation experiment demonstrates flexible sample-based terminal matching.

Lishuo Zhang, Ruizhi Huang, Yang Yu +1

School of Mathematical Sciences, Shanghai Jiao Tong University · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University

Normalizing Flows Velocity Field Exact

Ask PDF

Abstract

Explore similar work

Path-independent Flow Matching for Multi-parameter Generative Dynamics

Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once

Potential Matching Optimal Transport: Continuous Normalizing Flows for Exact ppp-Wasserstein Dynamics

Potential Matching Optimal Transport: Continuous Normalizing Flows for Exact $p$ -Wasserstein Dynamics