Solving a Special Type of Optimal Transport Problem by a Modified Hungarian Algorithm
YXIE 350@ GATECH . EDU
YLUO 373@ GATECH . EDU HUO @ GATECH . EDU
Abstract
We observe that computing empirical Wasserstein distance in the independence test is an optimal transport (OT) problem with a special structure.
This observation inspires us to study a special type of OT problem and propose a modified Hungarian algorithm to solve it exactly.
For an OT problem involving two marginals with m and n atoms ( m ≥ n ), respectively, the computational complexity of the proposed algorithm is O ( m 2 n ) .
Computing the empirical Wasserstein distance in the independence test requires solving this special type of OT problem, where we have m = n 2 .
The associated computational complexity of our algorithm is O ( n 5 ) , while the order of applying the classic Hungarian algorithm is O ( n 6 ) .
Numerical experiments validate our theoretical results.
Broader applications of the proposed algorithm are discussed at the end.
1. Introduction
One appealing application of optimal transport (OT) and Wasserstein distance [16, 18] is the independence test.
The Wasserstein distance between two distributions µ 1 , µ 2 ∈ Z is defined as:
W ( µ 1 , µ 2 ) :
= inf (cid:26)(cid:90)
Z 2 d ( z, z ′ ) dγ
( z, z ′ ) :
γ ( µ 1 , µ 2 ) is a distribution with marginals µ 1 and µ(cid:27) , where ( Z, d ) is a metric space (w.l.o.g., we consider 1-Wasserstein distance in this paper).
Wasserstein distance is a metric on probability measures [18].
To test the independence between variables Y ∼ νand Z ∼ ν 2 , people utilize the Wasserstein distance between the joint distribution of Y, Z and the product distribution of Y, Z , i.e., W p ( γ ( ν 1 , ν 2 ) , ν 1 ⊗ ν 2 ) .
While the statistical properties of this approach have been intensely investigated [11, 14, 19], there is no existing literature focusing on the computational aspect.
In this paper, we will consider ‘ How to calculate the empirical Wasserstein distance in the independence test? ’
In practice, given n i.i.d.
samples { ( y 1 , z 1 ) , ..., ( y n , z n ) } generated from ( Y, Z ) , we build the statistic— W p ( γ ( (cid:98) ν 1 , (cid:98) ν 2 ) , (cid:98) ν 1 ⊗ (cid:98) ν 2 ) , where (cid:98) ν denotes the corresponding empirical distribution—to test the independence.
We will show later that computing this statistic is equivalent to solving the following optimization problem:
[Equation]
where the metric could be chosen as d (( y i , z j ) , ( y k , z l )) =
∥
y i − y k ∥ p + ∥ z j − z l ∥ p , and ∥ · ∥ p denotes the l p norm. The problem (1) is an OT problem involving two marginals.
One marginal is an n -dimension uniform ( (cid:80) ni,j =1 X ◦ ij ; k = 1 /n, ∀ k ), and the other marginal is an n 2 -dimension uniform ( (cid:80) nk =1 X ◦ ij ; k = 1 /n 2 , ∀ i, j ). Per Birkhoff’s theorem [3], the solution to (1) is a vertex (whose coordinates are zeros and ones).
Consequently, we study the following special type of OT problem:
[Equation]
where 0 < n ≤ m , i = 1 , ..., m , j = 1 , ..., n , and m j are positive integers.
In essence, problem (1) is a special case of problem (2) where m
j = n, m = n 2 , seeing more details in Section 3.
Then, we focus on the question: ‘How to solve problem (2) ?’
One may recall the assignment problem, which aligns with the permutation matrix.
X ′ ∈ U ′ is similar but different from the permutation matrix: X ′ ∈
U ′ is an m × n matrix instead of a square matrix and has multiple entries of 1 in each column instead of only one entry.
In this way, we are not able to directly apply algorithms for the assignment problem like the Hungarian algorithm [12] to solve the problem (2).
An approach to obtain the precise solution to (2) is first to duplicate the columns of C and X ′ , and then apply the Hungarian algorithm.
The computational complexity of this approach is O ( m 3 ) .
In this paper, we will propose a modified Hungarian algorithm.
The algorithm specializes in solving the special type of OT problem (2) with a provable lower order— O ( m 2 n ) .
Back to the Wasserstein-distance-based independence test problem (1), the resulting computational complexity of applying the proposed algorithm is O ( n 5 ) while the order of applying the classic Hungarian algorithm is O ( n 6 ) .
Related work: Notice that our modified Hungarian algorithm is an exact OT solver.
There are also a bunch of approximation algorithms [4, 5, 8, 20].
Note that some scenarios may require the precise solution of OT problems, and paper [6] demonstrates the favorable numerical performance of the exact solutions over the approximate solutions.
Therefore, developing efficient exact algorithms is meaningful.
Our contributions: Firstly, we find out that computing the Wasserstein distance in the independence test (1) belongs to a special type of OT problem (2).
Secondly, we propose a modified
Hungarian algorithm to solve problem (2).
The algorithm modifies the Hungarian algorithm to deal with the scenario where two marginals are of different dimensions and atoms in one of the marginals have multiple assignments.
Thirdly, we adopt our algorithm to the independence test problem (1).
We further discuss other applications of the proposed algorithm, seeing Section 5.
2. Modified Hungarian algorithm
In this section, we propose a modified Hungarian algorithm to solve the problem (2). [Figure] D A C
B D [Figure] We first review the basics of the Hungarian algorithm [12, 17].
Suppose that a bipartite graph
G =
( V, E ) ,
where V = V 1 ∪ V 2 , E ⊂ V 1 × V 2 , has weight w ( · ) for each edge e ∈ E .
M ⊂ E is a matching if every node of G coincides with at most one edge of M .
Kuhn-Munkres theorem [12] illustrates that solving the assignment problem is equivalent to finding a perfect matching on the equality graph associated with some feasible labeling.
(The feasible labeling is defined as a function l : V → R satisfying l ( v 1 )
+ l ( v 2 ) ≥ w ( v 1 , v 2 ) , ∀ v 1 ∈ V 1 , v 2 ∈ V 2 , and the equality graph w r.t. labeling l is G ′
=
( V, E l ) , where E l = { ( v 1 , v 2 ) :
l ( v 1 ) + l ( v 2 ) =
w ( v 1 , v 2 ) } .)
Hence, Hungarian algorithm solves the assignment problem by looking for a perfect matching on some equality graph.
In problem (2), X ′ ∈
U ′ has one entry of 1 in each row, multiple entries of 1 in each column and 0 ′ s elsewhere.
Since a permutation matrix corresponds to a (perfect) matching in the bipartite graph, we define ‘pseudo-matching’ in the bipartite graph to describe X ′ :
In the bipartite graph G , where | V 1 | = m, | V 2 | = n .
P M ⊂ E is a pseudo-matching if every node of V 1 coincides with at most one edge of P M and j th node of V 2 coincides with at most m j edges of P M , where (cid:80) nj =
1 m j = m .
Further, if every node of V 1 coincides with one edge of P M and j th node of V 2 coincides with m j edges of P M , P M is called a perfect pseudo-matching .
See examples in Figure 1.
We develop a modified Kuhn-Munkres theorem (the theorem and its proof can be found in the
Appendix).
It demonstrates that we only need to find a perfect pseudo-matching on some equality graph to solve the problem (2).
In this way, we design a modified Hungarian algorithm (Algorithm 1).
It improves either the feasible labeling (adding edges to the associated equality graph) or the pseudo-matching until the pseudo-matching is perfect on some equality graph w.r.t.
some feasible labeling.
The algorithm improves the pseudo-matching by generating pseudo-augmenting paths and then exchanging the edge status along the paths.
This process is called the pseudo-augmenting process.
Also, we force the pseudo-augmenting paths emanating from V 2 , which have a lower order of nodes.
See the relevant definitions in Definition 1 and examples in Figure 1 .
Definition 1 (Free, matched, pseudo-matched, pseudo-alternating path, pseudo-augmenting path)
Let P M be a pseudo-matching of G =
( V, E ) .
• If the node v is in V 1 , it is pseudo-matched if it is an endpoint of some edge in P M ; if the node v is the j th node in V 2 , it is pseudo-matched if it is an endpoint of m j edges in P M .
Otherwise, the node is free.
• If the node v ∈ V , we say it is matched if it is an endpoint of some edge in P M .
• A path is pseudo-alternating if its edge alternates between
P M and E − P M .
A pseudo-alternating path is pseudo-augmenting if both its endpoints are free.
We now analyze the computational complexity of Algorithm 1.
Similar to the Hungarian algorithm [17], we keep track of slack v 1 = min v 2 ∈ S { l ( v 1 )
+ l ( v 2 ) − w ( v 1 , v 2 ) } , ∀ v 1 ̸∈ T .
The number of edges of the pseudo-matching increases by 1 after one loop, so we need O ( m ) loops to
Generate an initial feasible labeling l : ∀ v 2 ∈ V 2 , l ( v 2 ) = 0; ∀ v 1 ∈ V 1 , l ( v 1 ) = max v 2 ∈ V 2 { w ( v 1 , v 2 ) } and initialize a pseudo-matching M in E l ; 1 if M is a perfect pseudo-matching then
Stop else
Pick up a free node v free ∈ V 2 .
Set S
= { v free } , T = ∅ ; for v 1 ∈ V 1 is matched to v free do
T = T ∪ vend end 2 if N l ( S ) − T = ∅ then update labeling such that forcing N l ( S ) − T ̸
= ∅ : α l = min v 1 ̸∈ T,v 2 ∈ S { l ( v 1 ) + l ( v 2 ) − w ( v 1 , v 2 ) }
; l ( v )
=
l ( v ) − α l , v ∈ S ; l ( v ) =
l ( v ) + α l , v ∈ T ; l ( v ) =
l ( v ) , otherwise .
end 3 if N l ( S ) − T ̸
= ∅ then pick v 1 ∈ N l ( S ) − T ; if v 1 is free then v free → v 1 is a pseudo-augmenting path.
Pseudo-augment the pseudo-matching M .
Go to Step 1; end if v 1 is pseudo-matched to z then extend the pseudo-matching tree: S = S ∪ { z } , T = T ∪ { v 1 } ; for v 1 ∈ V 1 is matched to z do
T = T ∪ v 1 ; end end
Go to Step 2.
end form a perfect pseudo-matching.
There are two subroutines in each loop: (1): updating feasible labeling (Step 2); (2): improving pseudo-matching (Step 3).
In procedure (1): Since there are n nodes in V 2 , improving the labeling occurs O ( n ) times to build a pseudo-alternating tree.
Computing α l , updating the slack s and calculating the labeling cost O ( m ) each time.
In procedure (2): When a new node has been added to S , it costs O ( m ) to update slack s, and O ( n ) nodes could be added. On the other hand, when a node has been added to T , we just remove the corresponding slack v 1 . From the above we conclude that each loop costs O ( mn ) , so the total computational complexity of Algorithm 1 to solve problem (2) is O ( m 2 n ) . Note that the adoption of the Hungarian algorithm has an order of O ( m 3 ) .
Hence, the proposed modified Hungarian algorithm will outperform, especially when m ≫ n .
3.5 ln ( sample size )6101418l n ( # o p e r a t i o n ) dependent case, p = 2.
3.5 ln ( sample size )6101418dependent case, p = 2.
3.5 ln ( sample size )6101418independent case, p = 2.
3.5 ln ( sample size )6101418independent case, p =
[Figure]
3. Independence test using the Wasserstein distance
In this section, we discuss the special case of the OT problem (2)—independence test using the
Wasserstein distance, and apply the modified Hungarian algorithm to solve it.
Suppose that we have n i.i.d.
samples { ( y 1 , z 1 ) , ..., ( y n , z n ) } , where ( y i , z j ) ∼ ( Y, Z ) .
If we utilize the statistic W p ( γ ( (cid:98) ν 1 , (cid:98) ν 2 ) , (cid:98) ν 1 ⊗ (cid:98) ν 2 ) to test the independence between Y and Z , we need to solve the following optimization problem:
[Equation]
where
[Equation]
It is worth noting that X ij ; kl = 0 , ∀ k ̸ = l .
If we further let X ◦ ij ; k := (cid:80) nl =1 X ij ; kl , problem (3) could be reformulated as problem (1) and is equivalent to the following optimization problem: (The proof and relevant discussions are relegated to the Appendix.) min
[Equation]
Problem (4) belongs to the special type of OT problem where m j = n, m = n 2 .
Adopting the
Hungarian algorithm to problem (4) costs O ( n 6 ) while adopting the proposed Hungarian algorithm directly costs O ( n 5 ) .
4. Numerical experiments
In this section, we carry out experiments on the independence test problem (1) on CIFAR10 1 [10] and Wisconsin breast cancer 2 [7] datasets.
The numerical results validate the computational complexity of the modified Hungarian algorithm and the favorability of applying the proposed algorithm over
3.5 ln ( sample size )6101418l n ( # o p e r a t i o n ) dependent case, p = 2.
3.5 ln ( sample size )6101418dependent case, p = 2. 3.5 ln ( sample size )6101418independent case, p = 2. 3.5 ln ( sample size )6101418independent case, p = [Figure]
applying the Hungarian algorithm.
The performance is evaluated by the number of numerical operations here.
We relegate the results of the running time to the Appendix.
Experiments details: We create one dependent case and one independent case with different sample sizes: CIFAR10 :
Each image is a 3072-dim vector.
Let X ∈ R 3072 be generated uniformly from the first 5 classes; Y ∈ R 3072 be generated uniformly from other classes.
We calculate the Wasserstein distance (1) independent case : between X and Y 1 (where Y 1 is the first 1536 coordinates of Y ); (2) dependent case : between X and Z (where Z = X 2 / 2 + Y 1 / 2 , X 2 is the lastcoordinates of X , Y 1 is the first 1536 coordinates of Y ).
Breast cancer data: Each instance is a 30-dim vector.
Let X ∈ R 30 be generated uniformly from the benign class, and Y ∈ R 30 be generated uniformly from the malignant class.
We calculate the Wasserstein distance (1) independent case : between X 1 and Y 2 (where X 1 is the first 5 coordinates of X , Y 2 the last 25 coordinates of Y ); (2) dependent case : between X and Z (where Z = X 1 ∗ Y 1 , X 1 is the first 5 coordinates of X , Y 1 is the first 5 coordinates of Y , ∗ means the coordinate-wise product).
We run the algorithms on each case 10 times and plot the worst, best and average number of numerical operations.
Experiments results: The results are presented in Figure 2, 3.
The figures illustrate that the proposed algorithm gains a factor n in computational complexity when solving problem (1).
To be more specific, notice that the slope of ln (number of numerical operations) over ln (sample size) indicates the order of the associated algorithm, and the slope of our algorithm is around 5 while the slope of the Hungarian algorithm is around 6.
This observation implies that the order of applying our algorithm is O ( n 5 ) while the order of applying the Hungarian algorithm is O ( n 6 ) .
It is consistent with our theoretical results.
5. Discussion
We propose a modified Hungarian algorithm to solve a special type of OT problem (2).
The proposed algorithm could help solve one-to-many assignment [21] and many-to-many assignment [22] problems.
In this way, our algorithm could be applied to related practical problems in service assignment problems [13], sensor networks [2], access control [1], etc.
In addition, we may apply our algorithm to some modern machine learning models, including clustering [9] and representation learning [15].
Appendix A. Modified Kuhn-Munkres theorem
Theorem 2 (Modified Kuhn-Munkres theorem)
If l is a feasible labeling on the weighted bipartite graph G =
( V, E ) and P M ⊂ E l is a perfect pseudo-matching on the corresponding equality graph G ′
=
( V, E l ) , then P M is a maximum weighted pseudo-matching.
Proof Denote the edge e ∈ E by e =
( e v 1 , e v 2 ) .
Let P M ′ be any perfect pseudo-matching in G (not necessarily in the equality graph E l ).
And v i 1 , i = 1 , ..., m ; v j 2 , j = 1 , ..., n are nodes from V 1 and
V 2 , respectively.
Since v i 1 ∈ V 1 is covered exactly once by P M ′ and v j 2 ∈ V 2 is covered exactly m j times by P M ′ , we have
[Equation]
where the first inequality comes from the definition of the feasible labeling.
Thus, (cid:80) mi =1 l ( v i 1 ) + (cid:80) nj =
m j l ( v j 2 ) is upper bound of the weight of any perfect pseudo-matching.
Then let P M be a perfect pseudo-matching in the equality graph E l , we have
[Equation]
Hence w ( P M ′ ) ≤ w ( P M ) and P M is the maximum weighted pseudo-matching.
A.1. Proof and relevant discussion for Section
Instead of direct proving (3) ⇐⇒ (4), we prove the more general scenario first [Equation]
Proposition
The optimization problem (2) is equivalent to the optimization problem (5) :
[Equation]
Proof First consider the following two optimization problems (6)-(7):
[Equation]
where C † is an m × m matrix generated by duplicating the j th column of C m j times:
[Equation]
where C 1 is an m × m matrix generated by duplicating the j th column of C m j times:
[Equation]
Then we denote the objective functions of the problems (2), (5), (6) and (7) by f ( X ) , f † ( X † ) , f ′ ( X ′ ) and f 1 ( X 1 ) , respectively.
Firstly, we prove (5) ⇐⇒ (7).
On the one hand, for any X 1 ∈ U 1 , if we let
[Equation]
then we have
[Equation]
Thus, X ∈ U .
For the objective functions
[Equation]
we have the following:
[Equation]
On the other hand
[Equation]
for any X ∈ U
[Equation]
if we let
[Equation]
then we have
[Equation]
Thus, X 1 ∈ U 1 .
For the objective functions [Equation] we have the following:
[Equation]
Hence, (5) ⇐⇒ (7).
By Birkhoff’s theorem, we know (6) ⇐⇒ (7).
Therefore, we have (5) ⇐⇒ (6).
Similarly, for any X † ∈ U † , if we let
[Equation]
then X ′ ∈
U ′ and f † ( X † ) =
f ′ ( X ′ ) .
For any X ′ ∈ U ′ , if we let
[Equation]
j ≥ [Equation]
then X † ∈ U † and f † ( X † ) =
f ′ ( X ′ ) .
Therefore, (6) ⇐⇒ (2).
In conclusion, we have (5) ⇐⇒ (6) ⇐⇒ (2).
Next, we apply Proposition 3 to get our main result.
Proposition 4 The optimization problem (1) is equivalent to the optimization problem (4) .
Proof Applying Proposition 3 to the optimization problem (1), the optimization problem (1) is equivalent to problem (8):
[Equation]
If we express the index h by two indices i, j and express t by k, l , one could check that (8) ⇐⇒ (4).
In conclusion, (1) ⇐⇒ (4)
[Figure]
Appendix B. Experimental results w.r
[Figure]
t. running time
[Figure]
[Figure]
We run the Hungarian algorithm code from package ‘scipy’ in Python and implement the modified
Hungarian algorithm.
We record the running time.
Figure 4, 5 are experiment results w.r.t.
the running time.
When the sample size is small, our algorithm may be slower than the Hungarian algorithm.
However, as the sample size increases, our algorithm becomes more efficient than the Hungarian algorithm.
In theory, our proposed algorithm will compare more favorably if both algorithms are implemented in a low-level language, such as C++.
Appendix C. Resources
All experiments in this paper are implemented using Python 3.8 on Macbook Pro 2019 with 16 GB memory.