Clustering, factor discovery and optimal transport
Hongkang Yang Esteban G. Tabak

All-Against-One Linear-Quadratic
 Differential Games
    
	The authors are with the Department of Electrical and Computer Engineering, University of Central Florida, Orlando, FL. Emails: shahriar@ece.ucf.edu, simaan@ucf.edu, qu@ucf.edu.
	This work was supported in part by the US National Science Foundation under grants ECCS-1308928.
Shahriar Talebi, Student Member, IEEE, Marwan A. Simaan, Life Fellow, IEEE,  Zhihua Qu, Fellow, IEEE,


    
=================================================================================================================================================================================================================================================================================================================================================================================================




The clustering problem, and more generally, latent factor discovery –or latent space inference– is formulated in terms of the Wasserstein barycenter problem from optimal transport. The objective proposed is the maximization of the variability attributable to class, further characterized as the minimization of the variance of the Wasserstein barycenter. Existing theory, which constrains the transport maps to rigid translations, is extended to affine transformations. The resulting non-parametric clustering algorithms include k-means as a special case and exhibit more robust performance. A continuous version of these algorithms discovers continuous latent variables and generalizes principal curves. The strength of these algorithms is demonstrated by tests on both artificial and real-world data sets.


Keywords: Clustering, optimal transport, Wasserstein barycenter, factor discovery, explanation of variability, principal curve



AMS Subject classification: 62H30, 62H25, 49K30



§ INTRODUCTION


In order to analyze a given data set, it is often instructive to consider a factor model:

    X = G(Z_1, … Z_n),

where the data is modelled as a random variable X generated from unknown factors, the random variables Z_i, through an unknown generating process, the function G. The given data is thought of as a sample set drawn from a probability measure ρ, and our task is to construct a factor model (<ref>) such that the law of X equals ρ.

This task can be addressed in at least two different ways, leading to two distinct problems:


  * Generative modeling and density estimation: The factors Z_i are fixed, for instance as Gaussians, and the emphasis is on constructing the generator G, which helps us generate more samples from ρ or estimate its density function ρ(x). This approach is exemplified by variational autoencoders <cit.>, generative adversarial networks <cit.> and normalizing flows <cit.>.


  * Factor discovery: The emphasis is instead on identifying reasonable and interpretable factors Z. An informal criterion for choosing Z is that it can satisfactorily explain the data ρ. This approach is exemplified by clustering, principal components and principal surfaces <cit.>.

This paper concerns the second problem. Most of the existing works focus on either of two directions:


  * Effective objective functions that characterize a satisfactory factor Z.

  * Effective modelling assumptions on Z and generating process G.

To make the discussion more concrete, let us consider an elementary setting:

    X = G(Z) + Y

where Z is the factor we seek and Y is a random variable that takes care of the component of X unexplained by Z. Define the objective function as the mean-square error

    L(Z,G) = inf_X ∼ρ[ X-G(Z)^2 ].


Different modeling assumptions on Z,G could lead to diverse solutions. Suppose that Z is a random variable on some latent space  with law ν:


  * Assumption 1: The latent space is finite, ={1,… K}. Then (<ref>) assigns X to the nearest point G(k) and thus the problem reduces to k-means clustering. The probability measure ν becomes the weights of the clusters ρ_k, and the unexplained Y conditioned on Z=k is the k-th cluster with its mean removed.


  * Assumption 2: The latent space is low-dimensional Euclidean =^k and the generator G is linear. Then, the problem reduces to finding a k-dimensional principal hyperplane, ν is the projection of the data ρ onto the hyperplane, and Y|Z are the residuals of X perpendicular to the plane.


  * Assumption 3: =^k and G is smooth. Then, the problem reduces to principal curve/surface/hypersurface, and Y|Z are the residuals of X perpendicular to the surface.


Hence, even with the simplest mean-square objective (<ref>), the factor discovery problem contains rich algorithms and applications. The aim of this paper is to design an objective function that is more suitable than (<ref>) for the task of factor discovery, and explore what clustering and principal surface algorithms can be derived from it.



The factor discovery problem is not an end in itself. Ultimately, we aim at fully understanding the data ρ, identifying in particular all hidden factors that participate in the data-generating process (<ref>). Once a factor Z_1 is found by a factor discovery algorithm, we may filter out the influence of Z_1 in the data X and proceed to look for other, less prominent factors Z_2,Z_3,… in the filtered data.

As a concrete example of filtering, consider the task of removing batch effects <cit.>, common in biostatistics. Biomedical datasets such as gene expression microarrays are highly vulnerable to batch effects: the data are usually processed and collected from multiple small batches, inevitably introducing a large amount of batch-related variance into the dataset, which corrupts the data and prevents us from discovering meaningful factors other than the batches.

One solution is to study separately each conditional distribution ρ(· |z) (the law of X|Z=z, z∈). A problem with this is that each conditional distributions could correspond to very few data points and, in examples other than batch removal, when  can be continuous, no data at all in our finite dataset. In addition, any additionaL factors discovered could be inconsistent across different z ∈. Instead, The procedure that we propose consolidates all conditional distributions ρ(· |z) into one representative distribution, their Wasserstein or optimal transport barycenter μ. This will be described in section <ref>, but see Figure <ref> for a demonstration.

The intuition is that the barycenter is the z-independent distribution that is closest to all conditional distributions ρ(· |z), so mapping each ρ(· |z) to the berycenter is a way to consolidate all batches into one dataset with minimum distortion. Thus, the barycenter should preserve all meaningful patterns in the data X that are not explainable by Z, allowing us to perform factor discovery on this filtered data to uncover additional hidden factors.



To see how this procedure affects the choice of objective function, recall the mean-square objective (<ref>):

    min_G L(Z,G) = min_G inf_X ∼ρ[ [ X-G(Z)^2 | Z=z ] ] = inf_X ∼ρ[ Var(X| Z=z) ] ]

where the optimal G(z) is the mean of X|Z=z. Then, the remaining problem, i.e. factor discovery, is equivalent to minimizing the weighted average of the variances of all conditionals ρ(· |z):

    min_Zinf_X ∼ρ[ Var(X| Z=z) ] ] = min_ρ(·|z),ν{[Var(X|z)]  |  X|z ∼ρ(· |z),  ρ = ∫_ρ(·|z) dν(z) }.

For factor discovery, we replace the conditionals ρ(·|z) by their barycenter μ, which is more suitable for conducting downstream analysis. So it is natural to define a new objective function by the variance of this barycenter:

    min_ρ(·|z),ν{ Var(Y)  |  Y ∼μ,  μ is the barycenter of ρ(·|z)ν(z),  ρ = ∫ρ(·|z) dν(z) }.

Thus, a minimizer produces the maximum reduction in variance when we go from the raw data ρ to the processed data (the barycenter) μ, which intuitively corresponds to the maximum gain in information.
Finally, replace ρ(·|z) and ν by the assignment distribution ν(·|x), which softly assigns latent values z to the sample x:

    min_ν(·|x){ Var(Y)  |  Y ∼μ,  μ is the barycenter of ρ(·|z)ν(z),  ρ(·|z) ν = ν(·|x)ρ}





The application of Wasserstein barycenter to the factor discovery problem was first proposed in <cit.>. The main difficulty in solving (<ref>) is to compute the barycenter's variance (given ρ(·|z) and ν(z)), for which there are currently two approaches:


  * General distributions: Compute the barycenter for general conditionals ρ(·|z), which involves finding the optimal transport maps, and optimize the assignment distribution ν(·|x) on top of this computation. This is the approach in <cit.>, which models and optimizes these functions through neural networks.


  * Special distributions: Assume that the conditionals ρ(·|z) belong to certain distribution families that admit a tractable solution for barycenters. The setting considered in <cit.> is that all ρ(·|z) are equivalent up to rigid translations, and thus the barycenter μ corrresponds to a z-dependent rigid translation of each ρ(·|z). It follows that the objective (<ref>) reduces to the mean-square objective (<ref>), and from it we can recover classical algorithms such as k-means and principal components.


This article follows the second approach and generalizes the assumption of translation equivalence in <cit.> to the less stringent condition that the conditionals ρ(·|z) are equivalent up to affine transforms. By taking into consideration the second moments, this generalized objective function leads to clustering and principal surface algorithms that are more robust to the clusters' diversity. In particular, we introduce two practical algorithms: barycentric k-means (Algorithm <ref>) and affine factor discovery (Algorithm <ref>). 
Barycentric k-means is a clustering algorithm almost identical to k-means, except for a simple change in the update rule: k-means assigns each sample x_i to cluster k_i by

    k_i argmin_k ||x_i-x_k||^2

while barycentric k-means assigns x_i by

    k_i argmin_k(||x_i - x_k||^2/σ_k + σ_k )

where x_k,σ_k are the means and standard deviations of cluster k. Affine factor discovery is a principal curve/surface algorithm analogous to barycentric k-means.


Existing literature includes several clustering algorithms that utilize the second moment information, which we will introduce and compare with our algorithms in Section <ref>. The purpose of this article is not so much to build one new method for clustering or principal surfaces, but rather to develop a paradigm that conceptualizes these algorithms as a procedure for the general factor discovery problem (<ref>) and to show that optimal transport and Wasserstein barycenter are useful tools for uncovering and assessing hidden factors.


The plan of this article is as follows. After this introduction, Section <ref> gives a quick overview of optimal transport, Wasserstein barycenter and its affine formulation. Section <ref> optimizes our factor discovery objective function (<ref>) with affine maps and derives clustering algorithms based on k-means. Section <ref> considers the continuous setting when = and derives a more general principal surface algorithm. Section <ref> experiments with these algorithms on synthetic and real-world data. Section <ref> summarizes the work and sketches some directions of current research.






§ BACKGROUND ON THE WASSERSTEIN BARYCENTER



All data distributions in this paper live in ^d, while the factors z belong to a measurable space  to be specified. Denote the raw data by ρ. For each data point x, the assignment distribution ν(z|x) indicates the likelihood that factor z corresponds to x. One can define the joint distribution π


    π(x,z) = ρ(x) ν(z|x) = ρ(x|z) ν(z),


where ν is the latent distribution of the factors and the conditional distributions ρ(·|z) are commonly referred to as clusters. We say that a joint distribution π=ρ(·|z) ν is a disintegration or clustering plan of ρ if

    ρ = ∫_ρ(·|z) dν(z) = ∫_π dz .


Denote by (^d),_2(^d),_ac(^d) the space of Borel probability measures in ^d, those with finite second moments and those that are absolutely continuous respectively.

A random variable X with distribution law ρ is denoted by X ∼ρ. The notation T#ρ = μ indicates that a measurable map T transports (or pushes-forward) a probability measure ρ to another measure μ: If X∼ρ, then μ = law(T(x)). Equivalently, for all measurable subset A ⊆ℝ^d,

    μ(A) = ρ(T^-1(A)).


The Wasserstein metric W_2 or optimal transport cost with square distance <cit.> between probability measures ρ and μ is given by


    W_2^2(ρ,μ) = inf_X∼ρ,Y∼μ[X-Y^2].


By Brennier's theorem <cit.>, whenever ρ∈_2,ac(^d), there exists a unique map T, the optimal transport map, such that T#ρ=μ and

    W_2(ρ,μ) = _X∼ρ[X-T(X)^2].


Given a disintegration (or clustering plan) π = ρ(·|z)ν, the Wasserstein barycenter μ is a minimizer of the average transport cost

    min_μ∈𝒫_2(ℝ^d)∫ W_2^2(ρ(·|z),μ)  dν(z)

The barycenter always exists by Theorem 1 of <cit.>. The Wasserstein metric (<ref>) can be considered as a measurement of pointwise distortion of the data ρ, and thus the barycenter is a way to summarize all conditionals ρ(·|z) into one distribution while minimizing their distortion. Hence, one should expect that the barycenter can reliably preserve the underlying patterns common to all ρ(·|z).



In order to work with the objective function (<ref>), we need to make tractable the computation and optimization of the barycenter's variance Var(Y),Y∼μ. As linear problems are often amenable to closed-form solutions and efficient computations, the Wasserstein barycenter would become much simpler if all ρ(·|z) and μ are equivalent up to affine transformations, that is, T#ρ(·|z_1) = ρ(·|z_2) for an affine map T.

Intuitively, a sufficient condition for this property is that the conditionals ρ(·|z) have similar shapes (e.g. they are all Gaussians), which implies that the barycenter μ as their representative should also have that shape. Consequently, the map T_k only needs to translate and dilate the ρ_k to transform them into μ. This is confirmed by Theorem 5 of <cit.>:



Given any measureable space Z and a joint distribution (<ref>), assume that ρ has finite second moment and each conditional ρ(·|z) is a Gaussian distribution

    ρ(·|z) = (x(z),Σ(z))

with mean m(z) and covariance matrix Σ(z).
Then, there exists a barycenter μ, which is a Gaussian with mean and covariance


    y   = ∫x(z) dν(z) = 𝔼_X∼ρ[X],
    Σ_y    = ∫(Σ_y^1/2·Σ(z) ·Σ_y^1/2)^1/2 dν(z),


where Σ^1/2 is the principal matrix square root.





If we further assume that each ρ(·|z) is isotropic: Σ(z) = σ^2(z) · I_d, then the unique barycenter is also isotropic, with standard deviation

    σ = ∫σ(z) dν(z).




If we are only concerned with the setting of clustering with a finite latent space Z={1,… K}, then the above result has been given by Corollary 4.5 of <cit.>. Nevertheless, Section <ref> studies the continuous case with Z=ℝ, and we need the full strength of Theorem <ref>.


Moreover, the optimal transport maps T between these distributions are also affine. By Theorem 2.1 of <cit.>, the optimal transport map T_z that transports ρ(·|z) to μ is given by

    T_z(x)    = A_z x + b_z
    
    A_z    = Σ(z)^-1/2(Σ(z)^1/2Σ_yΣ(z)^1/2)^1/2Σ(z)^-1/2,  
    b_z = y - A_z x(z)


From now on, we will always model the conditionals ρ(·|z) as Gaussian distributions. In fact, all our arguments and algorithms apply to the slightly more general case of location-scale families <cit.>, which by definition consist of distributions that are equivalently up to affine transformations.
The restriction to Gaussians conditionals (or location-scale families) is not a stringent requirement in practice. Common clustering methods such as EM algorithm <cit.> often model the clusters as Gaussians, while the “standard data" for k-means consists of spherical clusters with equal radii <cit.>.










§ CLUSTERING WITH AFFINE BARYCENTER



This section studies the setting of discrete factors ={1,… K} and derives clustering algorithms based on the objective function (<ref>) and the assumption of Gaussian conditionals ρ(·|z).

In practice, we are given the following empirical distribution


    ρ^(n) = 1/N∑_i=1^N δ_x_i,


where {x_i}_i=1^N are i.i.d. samples from ρ. Their weights are uniformly set to 1/N in this paper, but one could also consider a more general case with more/less weights attached, for instance, to outliers or to measurements performed under more or less controlled circumstances.

Any joint distribution π from (<ref>) can be expressed as an N× K matrix. For convenience, denote the conditionals ρ(· | k) by ρ_k for k ∈{1,… K}, represent the latent factor distribution ν by the weights P_k, and represent the assignment distribution ν(·|x) by the following matrix:

    P = [P_k^i],  P_k^i = ν(k|x_i) = N ·π(x_i,k).


The probability vector P^i representing ν(·|x_i) is a soft class assignment for x_i, which becomes hard when P^i is given by a basis vector _k_i. Each P^i ranges in the K-dimensional simplex Δ^K and the matrix P ranges in the compact convex set ∏_i=1^N Δ^K. The domain of hard assignments is the set of extremal points of ∏_i Δ^K <cit.>.

Once a class assignment P=N π is given, the clusters ρ_k(x) can be estimated by formula (<ref>),

    ρ_k(x_i)    = ν(k|x_i)ρ(x_i)/ν(k) = P_k^i N^-1/P_k = P_k^i/∑_j=1^N P_k^j.

Then, the cluster means x_k and covariances Σ_k can be estimated by

    x_k = ∑_i P^i_k x_i /∑_i P^i_k,  Σ_k = ∑_i P^i_k (x_i - x_k) · (x_i - x_k)^T /∑_i P^i_k .

By Theorem <ref>, the covariance of the barycenter μ is given by the discrete version of (<ref>):

    Σ_y = ∑_k=1^K P_k (Σ_y^1/2Σ_kΣ_y^1/2)^1/2.


This is a non-linear matrix equation that admits a unique positive-definite solution Σ_y if at least one of P_k Σ_k is positive-definite <cit.>. Then, Σ_y can be calculated through the following iteration scheme:

    Σ(n+1) Σ(n)^-1/2(∑_k=1^K P_kΣ(n)^1/2Σ_kΣ(n)^1/2)^2 Σ(n)^-1/2,


where the initialization Σ(0) is an arbitrary positive-definite matrix. This iteration is guaranteed to converge to the correct Σ_y (when one of P_k Σ_k is positive-definite) <cit.>.

Hence, the factor discovery objective (<ref>) can be specialized into

    min_P Tr[Σ_y],


with Σ_y defined by (<ref>) and (<ref>).




To gain insight into the functioning of (<ref>), we can analyze its behavior in the much simplified setting, when the Gaussian clusters ρ_k differ only in their means x_k. Then Theorem 2.1 of <cit.> implies that the optimal transport maps T_k are translations by y-x_k. In that case, <cit.> showed that the barycenter's variance is reduced to the sum of within-cluster variances (equivalently, sum of squared errors, SSE)

    Tr[Σ_y] = ∑_k=1^K ∑_i=1^N P^i_k ||x_i - x_k||^2

which is exactly the objective function of k-means. Hence, by upgrading to general Gaussian clusters, one develops a generalization of k-means that takes advantage of second moments.






 §.§ Gradient descent solution



Our approach is to perform gradient descent on P∈∏_i Δ^K to solve for the minimization (<ref>). Even though the implicit nonlinear matrix equation (<ref>) determines Σ_y uniquely, it is not clear a priori whether the solution Σ_y is differentiable. Thus, we prove in Appendix <ref> that the partial derivatives ∂Σ_y / ∂ P_k^i always exist. Then, in Appendix <ref>, we derive explicit formulae for these derivatives.

The gradient of Tr[Σ_y] with respect to each sample point x_i's probability vector P^i is given by

    ∇_P^i Tr[Σ_y] = ∑_k=1^K vec(I)^T · W_k · vec[(x_i-x_k)· (x_i-x_k)^T + Σ_k]_k

where _k is the basis vector, vec is vectorization (which stacks the columns of the input matrix into one long vector), and the W_k are weight matrices

    W_k    := (Σ_y^1/2⊗Σ_y^1/2) [∑_h=1^K P_h (U_h ⊗ U_h) (D_h^1/2⊗ I + I ⊗ D_h^1/2)^-1(D_h^1/2⊗ D_h^1/2) (U_h^T ⊗ U_h^T)]^-1
         [(U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1 (U_k^T ⊗ U_k^T)](Σ_y^1/2⊗Σ_y^1/2).

Here ⊗ is the Kronecker product, and the U are the orthonormal and D the diagonal matrices in the eigendecompositions

    Σ_y^1/2Σ_kΣ_y^1/2 = U_kD_kU_k^T  and Σ_y = U_yD_yU_y^T.




The update rule at each time of a gradient descent step t is a projected gradient descent:

    P(t+1) = Proj_∏_iΔ^K(P(t) -η·∇_PTr[Σ_y])

where η is the learning rate and Proj is the projection onto the closest stochastic matrix in ∏_iΔ^K, which can be computed efficiently as described in <cit.>.
The step size η can be either fixed at a small value, or determined at each step via backtracking line search <cit.>, using a threshold α∈ (0,1/2) and a shortening rate β∈(0,1), and reducing η into βη if the amount of descent is not enough:


    Tr[Σ_y](P(t+1)) - Tr[Σ_y](P(t)) > α vec(∇_PTr[Σ_y])^T · vec[P(t+1)-P(t)].


This descent-based clustering algorithm is summarized below, with initialization based on that of k-means.





It might appear at first sight that Algorithm <ref> performs an alternating descent, similarly to k-means, alternating between optimizing the assignment P and updating the means x_k. Nevertheless, the derivation of (<ref>) in Appendix <ref> does not treat x_k as constants. Instead, it directly solves for the gradient ∇_P Tr[Σ_y], incorporating the derivatives ∇_P x_k. Hence, Algorithm <ref> is simply a gradient descent on the objective (<ref>), which with sufficiently small step size η necessarily converges to a critical point.




Recall that k-means minimizes the sum of squared errors (<ref>), whose partial derivatives are simply

    ∂_P^i_k SSE = ||x_i-x_k||^2 + 2∑_j P^j_k (x_j-x_k)∂x_k/∂ P^i_k = ||x_i-x_k||^2

Instead of performing gradient descent, k-means directly assigns x_i to the closest cluster k_i, that is,

    k_i= argmin_k ||x_i-x_k||^2 = argmin_k ∂_P^i_k SSE

We can interpret this hard assignment as equivalent to a gradient descent on P^i with arbitrarily large step size, followed by the projection Proj_Δ^K, so that P^i arrives at an extremal point of Δ^K, which is a one-hot vector.

In exactly the same way, we can simplify Algorithm <ref> into a hard assignment algorithm, such that each x_i is assigned to the cluster k_i with the smallest gradient term in ∂_P^iTr[Σ_y].




One common issue for clustering algorithms that utilize the second moments is that numerical errors could arise when some cluster's covariance becomes non-invertible, or equivalently, the cluster lies on a low-dimensional plane. Specifically such problem could arise for the D_k terms in the weight matrices (<ref>), though it has not been noticed in our experiments in Section <ref>. In practice, one can easily avoid this problem by making the covariances Σ_k strictly positive-definite: e.g.

    Σ_k Σ_k + ϵ I_d

for some ϵ≪ 1.





 §.§ Relation to k-means



We show next that the barycentric clustering algorithms reduce to k-means in the latter's setting. The “standard data” for k-mean consist of spherical clusters with identical radii and proportions <cit.>, which implies that P_1 =… =P_K = 1/K and Σ_1=…=Σ_K = σ^2/d I for some common variance σ^2. Then the gradient (<ref>) simplifies into


    ∂_P_k^i Tr[Σ_y] = σ^2/d Tr[(x_i-x_k)· (x_i-x_k)^T + Σ_k ] = σ^2/d ( ||x_i - x_k||^2 + σ^2).


Since each P^i lies in the simplex Δ^K, the direction of gradient descent must be parallel to Δ^K. Hence the term σ^2, shared by all entries of ∇_P^iTr[Σ_y], is eliminated by the projection map Proj_∏_iΔ^K of Algorithm <ref>, while for Algorithm <ref>, it is eliminated by the argmin_k step. The resulting gradient 
    ∇_P^iTr[Σ_y] = ∑_k=1^K ||x_i - x_k||^2 _k
 is precisely the gradient of the sum of squared errors (<ref>), the objective function of k-means. It is straightforward to check that Algorithm <ref> reduces to k-means, and thus k-means can be seen as a special case of barycentric clustering.










 §.§ Isotropic solution



Section <ref> will demonstrate that the barycentric clustering algorithms can recognize clusters that deviate from the “standard data”, for which k-means and fuzzy k-means would fail, but this robustness comes at the expense of the complexity of  gradients in (<ref>) and (<ref>).
Here we explore a situation in between, making hypotheses weaker than the “standard data”, yet strong enough to yield solutions that, while more robust than k-means, are at the same time simpler than (<ref>) and easier to interpret.

Since the complexity of (<ref>) results mostly from the non-commutativity of the matrix product, we can impose the assumption that all covariances are of the form


    Σ_k = σ^2_k/d I

where σ^2_k is the variance of cluster ρ_k. 

This assumption can be seen as a generalization of the “standard data”'s requirement that all clusters be radial with equal variances.

From Corollary <ref>, the barycenter's covariance becomes


    Σ_y = σ^2_y/d I,  σ_y = ∑_k=1^K P_k σ_k


and the gradient (<ref>) reduces to

    ∂_P_k^i Tr[Σ_y] = σ_y^3/d(||x_i-x_k||^2/σ_k+σ_k).

Since the algorithms are only concerned with the gradient's direction, the gradient is effectively

    ∇_P^i Tr[Σ_y] = ∑_k (||x_i-x_k||^2/σ_k+σ_k) _k.




Alternatively, we can obtain the gradient (<ref>) directly differentiating the weighted sum of standard deviations (<ref>). Note that the standard deviation can also be calculated via


    σ_k = √(Var(ρ_k)) = ( 1/2∬ ||x-y||^2 dρ_k(x) dρ_k(y) )^1/2 = (∑_i,j=1^N P_i^k P_j^k ||x_i-x_j||^2)^1/2/√(2)∑_i=1^N P_i^k.

Then the computation of the gradient

    ∂σ_y/∂ P_i^k = ∑_j=1^N P_j^k ||x_i-x_j||^2 /(2∑_i,j=1^N P^j_k P^l_k ||x_j-x_l||^2)^1/2

involves the samples only through the pairwise distances ||x_i-x_j||^2, which can be computed at the onset of the algorithm.
This is helpful when the data space ℝ^d has very large dimension, so that computing the means x_k and distances x_i-x_k^2 of (<ref>) at each iteration becomes prohibitive.
Moreover, we can replace x_i-x_j^2 with any “dissimilarity measure" such as Riemannian distance, graph distance or kernel functions (even though naive substitution might not be justified by our barycenter model).
Nevertheless, computing (<ref>) takes O(N^2) time, while (<ref>) takes O(N· d), so the latter is more efficient for large sample sets of low-dimensional data.




In the isotropic scenario, barycentric clustering (Algorithms <ref> and <ref>) can be modified via (<ref>) into the following:








We name Algorithm <ref> “Barycentric k-means", as it closely resembles k-means.

As discussed in Section <ref>, numerical errors could arise if some cluster degenerates into a point and its standard deviation σ_k vanishes. In practice, one can replace the denominator σ_k in the gradient (<ref>) by σ_k + ϵ for some small ϵ.



Section <ref> will confirm the expectation that these algorithms are more robust than k-means under varying proportions and radii (P_k and σ^2_k), but are more vulnerable to non-isotropy (Σ_k not of the form σ^2_k I) than Algorithms <ref> and <ref>.





 §.§ Relation to Mahalanobis distance



Barycentric clustering is not the first clustering algorithm that deals with non-isotropic clusters using second moment information. A series of clustering methods <cit.> based on k-means measure the distance between sample points and clusters by the Mahalanobis distance:


    d^2(x_i,x_k) = (x_i-x_k)^T Σ_k^-1(x_i-x_k),

which reduces the distance along the directions corresponding to the large eigenvalues of the covariance Σ_k. However, as pointed out in <cit.>,  applying (<ref>) to k-means has the problem that the objective function (<ref>) becomes trivial:


    SSE = ∑_i,k P_k^i (x_i-x_k)^T Σ_k^-1(x_i-x_k) = ∑_k P_k Tr[Σ_kΣ_k^-1] ≡ Tr[I].

The Gustafson–Kessel algorithm <cit.> remedies this problem by modifying (<ref>) into

    d^2(x_i,x_k) = det(Σ_k)^1/d (x_i-x_k)^T Σ_k^-1(x_i-x_k).



To compare barycentric clustering and the Mahalanobis distance-based algorithms, note that (<ref>) is dimensionless, in the sense that any shrinkage or dilation of cluster ρ_k (with respect to the mean x_k) would be completely cancelled out in (<ref>), which explains how its objective function becomes trivial. Meanwhile, both the squared Euclidean distance and the modified Mahalanobis distance (<ref>) have dimension [l]^2, where [l] denotes the unit of length. For our algorithms, if we assume that the weight P_k is small so that the influence of Σ_k on Σ_y is small, then it is routine to check that the gradient (<ref>, <ref>) has dimension [l].

This comparison becomes explicit in the isotropic setting:

    ||x_i-x_k||^2/σ^2_k,  ||x_i-x_k||^2,  ||x_i-x_k||^2/σ_k+σ_k

The first term is the Mahalanobis distance (<ref>), the second term is squared Euclidean distance or equivalently the modified Mahalanobis (<ref>), and the third term is the gradient (<ref>), the isotropic version of (<ref>). Hence, our algorithms can be seen as a balanced solution between the Euclidean case that completely ignores second moment information and the Mahalanobis case where too much normalization nullifies the problem.

As a side remark, the EM algorithm <cit.> is also a clustering method that can recognize non-isotropic clusters, typically modeling ρ_k as Gaussian distributions. The Gaussians are essentially an exponential family built from the Mahalanobis distance (<ref>), indicating a possible connection between EM and our gradients (<ref>, <ref>), and thus a connection between maximum-likelihood-based method and the Wasserstein barycenter framework.




 §.§ An alternative approach



We describe briefly here a different approach to solving the clustering problem (<ref>), which avoids directly computing the barycenter μ. By the Variance Decomposition Theorem of <cit.>, we have the identity


    Tr[Σ_x] = Tr[Σ_y] + ∫ W_2^2(ρ(·|z),μ) dν(z) ,

where Σ_x,Σ_y are the covariances of the data ρ and the barycenter μ. Since Σ_x is fixed, the clustering problem (<ref>) that minimizes Tr[Σ_y] is equivalent to the following maximization problem

    max_ν(·|x)∫ W_2^2(ρ(·|z),μ) dν(z) = max_P∑_k=1^K P_k W_2^2(ρ_k,μ)

where as usual, P = [P_i^k] denotes the assignment distribution P^i_k = ν(k|x_i) and the clusters ρ_k are determined by (<ref>).

Now we seek to remove μ from (<ref>). It is an interesting observation that the term ∑_k P_k W_2^2(ρ_k,μ) appears like a “variance", the weighted mean square distance between each “point" ρ_k and their “mean" μ. Recall that we have the following identity:

    ∀ρ∈𝒫_2(ℝ^d),  ∫ ||x-x||^2 dρ(x) = 1/2∬ ||x-y||^2 dρ(x)dρ(y)

which computes the variance without involving the mean x. Similarly, the following result allows us to compute (<ref>) without μ.



For any measurable space Z and joint distribution π from (<ref>), assume that ρ has finite second moments and all conditionals ρ(·|z) are Gaussian. Let μ be any Wasserstein barycenter of π. Then, we have the following identity:

    ∫ W_2^2(ρ(·|z),μ) dν(z) = 1/2∬ W_2^2(ρ(·|z_1),ρ(·|z_2)) dν(z_1)dν(z_2)




See Appendix <ref>.


It follows that the clustering problem (<ref>) is equivalent to

    max_P1/2∑_k,h=1^K P_k P_h W_2^2(ρ_k,ρ_h).


By Theorem 2.1 of <cit.>, as the clusters ρ_k are assumed Gaussian, the cost terms W_2^2(ρ_k,ρ_h) can be computed via

    W_2^2(ρ_k,ρ_h) = ||x_k-x_h||^2 + Tr[Σ_k]+Tr[Σ_h]-2Tr[(Σ_h^1/2Σ_kΣ_h^1/2)^1/2]


Lemma 2.4 of <cit.> provides a formula for the partial derivatives of (<ref>) with respect to x_k,Σ_k. Hence one could optimize problem (<ref>) using gradient descent.

Although this approach avoids computing the iteration (<ref>), a drawback is that its objective function contains O(K^2) terms as we are solving pairwise distances.




§ CONTINUOUS EXTENSION



Having studied the simplest case with discrete latent space = {1,… K} and the resulting clustering problem (<ref>), this section extends the analysis to continuous  and derives novel algorithms that construct principal curves and surfaces.

For brevity, we focus on the simple case when = and all clusters ρ(·|z) are isotropic, as it is straightforward to extend the discussion to more general cases.

Hence, given a joint distribution π=ρ(·|z)ν, we assume that all conditionals are isotropic Gaussians ρ(·|z) = (x(z),σ^2(z)I_d) with means x(z) and variance σ^2(z) (in each dimension). Then, their barycenter μ is given by Corollary <ref>.

In practice, we are only given a finite sample set {x_i}_i=1^N ⊆^d. If one uses hard assignment ν(·|x_i) = δ_z_i and model each label z_i ∈ as an independent variable to be optimized, it would be impossible to evaluate the barycenter's variance (<ref>): since the sample set is finite whereas there are infinitely many z ∈, almost all conditionals ρ(·|z) will be represented by zero or at most one sample point.

Our solution is inspired by human vision. For the image below, it is evident that ρ(·|z_1) has greater variance than ρ(·|z_0), even though there is no sample point whose assignment is exactly z_0 or z_1. The key is that we can estimate ρ(·|z_1) using the points nearby, {x_i | z_i ≈ z_1}.



Hence, it is natural to use a soft assignment ν(·|x_i) on =, which is concentrated around some z_i and decays for z far away from z_i. Effectively, the latent distribution {z_i} is smoothed into

    ν = 1/N∑_i=1^N ν(·|x_i)

Given any z, the conditional density ρ(x|z) can be estimated using Bayes' formula,

    ρ(x_i|z)    = ν(z|x_i)ρ(x_i)/ν(z) = ν(z|x_i)/∑_j=1^N ν(z|x_j),


and Corollary <ref> implies that the barycenter's standard deviation σ is given by


    σ = ∫σ(z) dν(z) = ∫[∑_i=1^N ||x_i-x||^2 ν(z|x_i)/∑_i=1^N ν(z|x_i)]^1/2 dν(z).


The objective of clustering (<ref>) is now equivalent to

    min_ν(·|x)σ = min_ν(·|x)∫[∑_i=1^N ||x_i-x||^2 ν(z|x_i)/∑_i=1^N ν(z|x_i)]^1/2 dν(z)


For simplicity, we parameterize the assignment distribution of each sample x_i by a one-dimensional Gaussian, ν(·|x_i) = (z_i,ϵ^2), where z_i are the means and ϵ^2 is the common variance. The means z_i are parameters by which we minimize (<ref>), and we define the vector z = [z_i]. When the sample set {x_i} is large, rather than having an independent variable z_i for each x_i, one can replace z_i by a parameterized function z(x_i), for instance a neural network.

Note that ϵ^2 should adapt to the set {z_i}. Otherwise, a fixed ϵ^2 would lead to the trivial solution where the z_i are arbitrarily far apart, so that by (<ref>) each conditional ρ(·|z) would be concentrated at some x_i and σ(z) would go to zero.

Hence, we choose ϵ^2 so as to make the distributions ν(·|x_i) close to each other for nearby x_i. Intuitively, for some fixed 0 < α < 1 (e.g. α = 10%), we want that for each ρ(·|z), roughly a fraction α of the {x_i} participate in ρ(·|z). The trivial solution corresponds to α≈ 0, and one should not set α≈ 1 either, for otherwise each x_i would have significant presence in each ρ(·|z), contrary to the goal of clustering as a partition of {x_i}. Hence we set the following objective for ϵ^2, based on maximum likelihood:


    max_ϵ^2∏_i,j=1^N (z_j  | z_i,(ϵ/α)^2)






with optimal solution given by


    ϵ^2 = α^2/2N^2∑_i,j^N ||z_i-z_j||^2 = α^2/N∑_i=1^N ||z_i - z||^2 = α^2 Var({z_i}),


where z is the sample mean. This choice of ϵ^2 dilates ν(·|x_i) proportionally to the spread of the {z_i}, thus preventing the trivial solution.

To fix z_i, we can further require that their mean should not drift away from 0. Adding this an extra term z^2 to the penalty yields the simpler formula


    ϵ^2 = α^2 (Var({z_i}) + z^2) = α^2 z^2/N,




so we propose


    ν(·|x_i) = 𝒩( z_i, α^2 ||z||^2/N).


Next, we derive the gradient of the barycenter's standard deviation σ. By Appendix <ref>, we can differentiate under the integral sign in (<ref>) to obtain the following gradient:


    ∂σ/∂z_i   = 𝔼_v[G_i(z)]
    
    G_i(z)    = 1/2N·ν(z)[ C(z) z_i + z-z_i/ϵ^2ν(z|x_i) (σ(z)+||x_i-x(z)||^2/σ(z)) ] 
    
    C(z)    = 1/||z||^2∑_j=1^N [σ(z)+||x_j-x(z)||^2/σ(z)] ·[ ||z-z_j||^2/ϵ^2 - 1/||z||] ρ(z|x_j)


The computation of the integrand G_i(z) for any z takes linear time O(N).

Estimating the expectation (<ref>) by random sampling from the latent distribution ν, we obtain the following stochastic gradient descent algorithm,
which will be tested in Section <ref>:





Whenever we obtain a joint distribution π = ν(·|x)ρ, the conditional mean x(z) of ρ(·|z) can be seen as a curve, parameterized by z∈ℝ, that summarizes the data ρ. If, as in Remark <ref>, we make the simplifying assumption that all conditional ρ(·|z) are equivalent up to translations, then the variance of the barycenter can be computed by

    ∀ z,  σ^2 = σ^2(z) ⟶σ^2 = ∫σ^2(z) dν(z) ≈1/N∑_i=1^N ||x_i-x(z_i)||^2

Or, in terms of soft assignments from (<ref>),

    σ^2 = 1/N∑_i=1^N ∫ ||x_i-x(z_i)||^2 dν(z|x_i)

Since the factor discovery (<ref>) is to minimize σ^2, an immediate corollary is that, given any curve x(z), the sample x_i should be assigned to the closest point on x(z):

    z_i = argmin_z ||x_i - x(z) ||^2

or the soft assignment ν(·|x_i) should be concentrated around this z_i. It follows that our formulation of clustering is reduced to an alternating descent algorithm that alternates between updating the conditional means x(z) and reassigning z_i. Yet, this procedure is exactly the principal curve algorithm <cit.>. Hence, problem (<ref>) is a generalization of principal curves (and principal surfaces if we set =^k).








§ PERFORMANCE



Sections <ref> and <ref> compare k-means and the barycentric clustering algorithms on artificial data that deviate from the “standard data” of k-means, and Section <ref> tests them on real-world classification data sets. Finally, Section <ref> tests Affine factor discovery (Algorithm <ref>) on earthquake data to uncover meaningful continuous latent variables.







 §.§ Comparison of soft assignments



We design two families of artificial data. The “expansion test” is a collection of three spherical Gaussian distributions in ℝ^2:


100 samples from 𝒩([0,0]^T
, 1/10I ),

100(1+t) samples from 𝒩([0,2+t]^T, (1+t)^2/10I ),

100(1+2t) samples from 𝒩([t+1/t+2√(12(2t+1)), 2(1-t^2)/t+2]^T, (1+2t)^2/10I ).

The means and variances are designed so that, for all t≥ 0, the three samples are roughly contained in three pairwise adjacent balls of radii 1, 1+t and 1+2t. As t increases, the sample sizes and radii grow in distinct rates.

The “dilation test” is given by


100 samples from 𝒩([ 0; 1 ], 1/25[ (1+t)^2       0;       0       1 ]),

100 samples from 𝒩([ 0; 0 ], 1/25[ 1 0; 0 1 ]),

100 samples from 𝒩([  0; -1 ], 1/25[ (1+t)^2       0;       0       1 ]),


which are Gaussians stretched horizontally at different rates.

The expansion test challenges the “standard data” (Section <ref>) in its first two assumptions: similar radii and similar proportions, while the dilation test challenges the assumption of isotropy. In both cases, the amount of deviation from the standard data is parametrized by t≥ 0.


The performance of each algorithm is measured by its correctness rate, the percentage of overlap between the true labeling and the labeling produced by the algorithm, maximized over all identifications between the proposed clusters and the true clusters: given the true labeling {z_i} and either the labeling {k_i} or the stochastic matrix P produced by algorithm, we define the correctness rate as


    max_g ∈ S_K∑_i1_z_i = g(k_i) or max_g∈ S_K∑_iP^i_g(z_i)

where g ranges over the permutation group S_K.





We first compare the soft assignment algorithms: Fuzzy k-means, barycentric clustering (Algorithm <ref>), and isotropic barycentric clustering (Algorithm <ref>). Note that k-means' objective (<ref>) is approximately a linear function in the assignment P=[P_i^k], and thus the optimal solutions are the extremal points of Dom(P) = ∏_i Δ^K, which are hard assignments. Hence, in order to obtain a valid comparison among soft assignments, we use the fuzzy k-means algorithm <cit.>, which generalizes k-means, minimizing the following objective function:


    J_c(P,{x_k}) = ∑_i∑_k (P_i^k)^c ||x_i-x_k||^2.


This is a generalization of the sum of squared errors (<ref>), with an exponent c>1 that makes J_c strictly convex in P^i_k, and therefore yields soft assignments. Here we adopt the common choice c=2.



Each algorithm is initialized with random means x_k and assignment matrix P, and all algorithms share the same sample set (for each t). To stabilize performance, each algorithm is run 100 times over the same sample set and the result that minimizes this algorithm's objective function (J_c for fuzzy k-means, (<ref>) for barycentric clustering, and (<ref>) for isotropic barycentric clustering) is selected.

The experimental results are plotted below. The first row corresponds to the expansion test with t=2.2, and the second row to the dilation test with t=3.0. The class assignment displayed is given by the maximum probability, k_i argmax_k P_k^i.



For the expansion test, fuzzy k-means merged the two smaller clusters and split the largest one, whereas the barycentric clustering algorithms only made a few errors on the periphery. For dilation test, the correct clusters are produced only by Barycentric clustering, whereas fuzzy k-means and Isotropic barycentric clustering split the clusters. These results are not surprising, since Section <ref> shows that k-means is an approximation to Barycentric clustering that assumes clusters with identical sizes and radii, while Isotropic barycentric clustering, by design, assumes isotropic clusters.

Below are the plots of correct rates (<ref>) for t∈ [0,4]. Fuzzy k-means, with a steady decline, is dominated by the barycentric clustering algorithms, while the difference between the latter two is small. Eventually, as t→∞, all algorithms deviate from the true labeling, since for very large t the Gaussians become so disparate that the true labeling no longer minimizes Tr[Σ_y] or yields reasonable clusters that agree with human perception.


In fact, for the dilation test, the contrast can be seen well before t=3.0. The following is the result for t=1.6, with the shaded regions representing the convex hulls containing the “core points” of each class, defined as C_k = {x_i, P_k^i>1/3}. The soft clusters produced by fuzzy k-means exhibit significant overlap, indicating that many sample points are assigned with highly ambiguous probability vectors P^i.






 §.§ Comparison of hard assignments


We compare next the hard assignment algorithms: k-means <cit.>, Hard barycentric clustering (Algorithm <ref>), and Barycentric k-means (Algorithm <ref>). The results on the expansion test and dilation test are plotted below. Again, for each algorithm on each sample set, the objective-minimizing result over 100 trials is selected. The performance comparison is analogous to that of soft assignment.


Nevertheless, Hard barycentric clustering is a simplified version of Barycentric clustering, replacing the latter's gradient descent, which moves by small steps, by class reassignment, which hops among the extremal points of ∏_i Δ^K. The correctness rate curves of dilation test indicate that, whereas Barycentric k-means has similar performance as Isotropic barycentric clustering, Hard barycentric clustering is more unstable than Barycentric clustering.





 §.§ Clustering on real-world data


To compare the performance of k-means and our algorithms on real-world problems, we use data sets from the online UCI Machine Learning Repository <cit.>. These data sets, intended for classification, are provided with labels, which we use to calculate the correctness rates (<ref>). The “Wine” <cit.> data set classifies wines based on chemical compositions, “Seeds” <cit.> classifies wheats by the shapes of wheat kernels, “Breast cancer (original)” <cit.> classifies benign/malign cancers by the shapes of cell nuclei, while “Breast cancer (diagnostic)” <cit.> classifies them by other cell statistics, “Parkison's” <cit.> diagnoses the disease by the patients' voice and speech, and “E.coli” <cit.> classifies proteins by their sequences and structures.

Since the setting for our clustering problem is for data in ℝ^d, the data's categorical attributes as well as entries with missing data are removed. The samples are normalized along each dimension before clustering, since their attributes are often on disparate scales. Again, each algorithm is run 100 times on each sample set, and the objective-minimizing result is selected.

In the following table, for each sample set, the marked entries are the ones with maximum correctness rates among the hard and soft assignment groups.




One notable difference between our synthetic tests and these real-world data is that the latter have higher dimensions, which negatively influence fuzzy k-means' performance. Previous studies <cit.> have shown that, as dimension increases, the pairwise distances of the sample points become homogeneous, and fuzzy k-means tends to produce uniform assignments: P^i_k ≈ 1/K. Nevertheless, the barycentric clustering algorithms remain robust, as shown below.







 §.§ Continuous latent variable and seismic data



Finally, we test Affine factor discovery (Algorithm <ref>). As discussed in Remark <ref>, the continuous factor discovery problem (<ref>) generalizes principal curves, so it is natural to evaluate Algorithm <ref> in terms of its “principal curve", that is, the conditional mean x(z). Given data that appears to cluster around one or several curves, the curve x(z) should discover these patterns.

We use the earthquake data from <cit.>, which covers more than two thousand earthquakes in the 20th century in the Southeast Asia earthquake zone. The sample {x_i} is two dimensional, recording the latitude and longitude of the earthquakes. We apply Affine factor discovery with a fixed number of iterations T=50000, proportion constant α=2.5%, and learning rate η=5× 10^-1, and we initialize z_i to be proportional to longitude, which is evidently far from the optimal solution. The curve of conditional means x(z) is plotted below.



The Southeast Asia earthquake zone lies at the intersection of the Australian Plate, Eurasian Plate, Indian Plate, Philippine Plate, Yangtze Plate, Amur Plate, and numerous minor plates and microplates. The tectonic boundaries are complex and cannot be represented by a single curve. Affine factor discovery automatically solved this problem using piecewise principal curves.
Note that even though the latent space Z=ℝ is connected, the support of the latent distribution, suppν, consists of several disjoint clusters, giving rise to piecewise continuous x(z).







§ CONCLUSION




Factor discovery follows the human intuition of finding significant factors that explain the data.
To pose the factor discovery problem in a systematic way, this article proposes an optimal transport based framework such that potential factors are extracted from the data and the resulting conditionals are consolidated into their barycenter, from which the procedure can continue inductively. Extending ideas proposed in <cit.>, we design a novel objective function that minimizes the unexplained variance of the data, featured as the barycenter's variance.

This article introduces several clustering and principle surface algorithms for solving the factor discovery problem in the Gaussian setting, which benefit from second moment information. Specifically,


  * The barycentric clustering algorithms (Algorithms <ref> and <ref>) leverage the clusters' covariance matrices and are capable of recognizing non-isotropic clusters with varying radii and sizes.


  * Barycentric k-means (Algorithms <ref>, designed for isotropic Gaussians) runs as efficiently as k-means, yet with little loss in robustness compared to the general barycentric clustering algorithms.


  * Affine factor discovery (Algorithm <ref>), an efficient principle curve/surface algorithm that uncovers continuous latent variables, utilizes the variance of each conditional distribution.


  * All algorithms proposed reduce to the classical k-means and principle surface algorithm in the simplified setting when all conditionals are equal up to translations, indicating that the optimal transport framework includes the classical intra-class variance-based framework as a special case.




The methodology developed in this article, in addition to its value as a practical tool for clustering and continuous factor discovery, opens the way to further inquiry into various directions:


  * Efficient non-isotropic algorithm: as discussed in Sections <ref> and <ref>, the main difficulty in solving the general non-isotropic problem is due to the non-commutativity of the covariances Σ_k. The minimum requirement for commutativity is that all Σ_k are simultaneously diagonalizable, that is, there exists some orthonormal basis with which all Σ_k become diagonal matrices. Then, the variance formula of Corollary <ref> can be applied to each coordinate, greatly simplifying the objective (<ref>) and its gradients. This commutative setting is non-isotropic and more general than the isotropic setting, and the corresponding clustering algorithm might inherit both the efficiency of Algorithms <ref> and the generality of Algorithms <ref> and <ref>.


  * Task-specific metrics: the Wasserstein barycenter minimizes the L^2 deformation of the data, so our formulations are all based on the Euclidean distance. Yet for many tasks, the most suitable underlying geometry is not Euclidean or homogeneous. For instance, in the space of images, certain deformations, such as translation and rotation, are more acceptable than others, such as tearing and smearing, even if they have similar Euclidean norm. A promising approach is to extend the algorithms to use the (squared) geodesic distance in the manifold that underlies the data, such as the Fermat distance introduced in <cit.>.

 

  * Iterative factor discovery: factor discovery can be performed iteratively to uncover a hierarchy of latent factors, which could correspond to diverse latent spaces . It could be of practical interests to develop a principled way of conducting multi-stage factor discovery. For gene expression data, for instance, one could first apply Affine factor discovery to reduce the dimensionality from the order of millions to several hundreds, then apply Barycentric k-means to identify the clusters, which might correspond to meaningful demographic groups, and finally apply Affine factor discovery on the resulting barycenter to obtain continuous factors that may offer biomedical insights.


  * Curse of dimensionality: for high dimensional distributions ρ, training can become more difficult as the landscape contains more local minima. Moreover, as the pairwise distances between the data points become similar, the gap between the global minimum and many of the bad local minima become less prominent. One effective way to constrain and regularize the solution is to implement the assignment distribution ν(·|x) or the hard assignment z(x) by neural networks. This modification is straightforward for principal curve/surface, while for clustering, one can place the points ={1,… K} on a grid in a low-dimensional ^k and apply a neural net that maps from ^k to ^d. This construction resembles the self-organizing maps of <cit.>, while the use of neural networks allow us to scale to high dimensions.








§ ACKNOWLEDGMENTS
 The work of E. G. Tabak was partially supported by NSF grant DMS-1715753 and ONR grant N00014-15-1-2355.




acm







Appendices:






§ EXISTENCE OF DERIVATIVE


Here we establish that ∂Σ_y / ∂ P_k^i, the partial derivatives of the barycenter's covariance implicitly defined by (<ref>) with respect to the assignment probabilities P_k^i = ν(k|x_i), always exist. Theorem 1.1 from <cit.> is a useful result on matrix derivatives which we restate below. Denote by 𝒮_d, 𝒮^+_d ⊆ℳ_d the linear subspace of symmetric matrices and the cone of positive-definite matrices.

The principal matrix square root function S∈𝒮^+_d → S^1/2∈𝒮^+_d is Fréchet differentiable to any order, and the first order derivative is given by the operator 
    (∇ S^1/2)(H) := ∇ S^1/2|_H = ∫_0^∞ e^-S^1/2t· H · e^-S^1/2tdt

such that for any H∈𝒮_d and S+hH∈𝒮^+_d 
    lim_h→ 01/h[(S+hH)^1/2-S^1/2 - h(∇ S^1/2)(H)] = 0.




Now we prove the existence of the partial derivatives through the Implicit Function Theorem.

For Σ_1,…Σ_K ∈𝒮_d^+, the solution Σ_y to
the covariance formula (<ref>) depends differentiably on Σ_1,…Σ_K.



For convenience, define the function 
    F(Σ_y,Σ_1,…Σ_K) = ∑_k=1^K P_k (Σ_y^1/2Σ_kΣ_y^1/2)^1/2 - Σ_y

It is a composition of C^1
functions on ∏_i=1^K+1 S^+_d and thus is C^1. To confirm that the gradient ∇_Σ_yF is non-singular, perturb Σ_y along an arbitrary direction S∈𝒮_d,

    (∇_Σ_yF)(S) =   ∑_k=1^K P_k (∇(Σ_y^1/2Σ_kΣ_y^1/2)^1/2)((∇Σ_y^1/2)(S)Σ_kΣ_y^1/2 + Σ_y^1/2Σ_k(∇Σ_y^1/2)(S)) - S
    
    =   ∑_k=1^K ∫_0^∞ e^-(Σ_y^1/2Σ_kΣ_y^1/2)^1/2t[(∫_0^∞ e^-Σ_y^1/2uSe^-Σ_y^1/2udu)Σ_kΣ_y^1/2
       +Σ_y^1/2Σ_k( ∫_0^∞ e^-Σ_y^1/2uSe^-Σ_y^1/2udu )] e^-(Σ_y^1/2Σ_kΣ_y^1/2)^1/2tdt - S


To evaluate integrals of the form ∫_0^∞ e^-Σ^1/2tS e^-Σ^1/2tdt, we can apply the eigendecomposition Σ = UDU^T

    ∫_0^∞ e^-(UDU^T)^1/2tS e^-(UDU^T)^1/2tdt = U (∫_0^∞ e^-D^1/2t U^T S U e^-D^1/2t dt )U^T 
    = U(T∘ U^T S U)U^T

where ∘ is Hadamard product and T_ij = 1/√(λ_i)+√(λ_j).

Thus, using   
Σ_y^1/2Σ_kΣ_y^1/2 = U_kD_kU_k^T,  Σ_y = U_yD_yU_y^T  
and the corresponding T_k,T_y, we obtain

    (∇_Σ_yF)(S) = ∑_k=1^K P_k U_k[T_k∘ U_k^T [ U_y(T_y∘ U_y^T S U_y)U_y^T Σ_kΣ_y^1/2 + Σ_y^1/2Σ_k U_y(T_y∘ U_y^T S U_y)U_y^T] U_k ] U_k^T - S


To check non-singularity, we set (∇_Σ_yF)(S) = 0 and vectorize the equation to disentangle S. We apply the identity that vec(AXB) = (B^T⊗ A)vec(X) where ⊗ is the Kronecker product <cit.>. Meanwhile, vectorizing the Hadamard product yields

    vec(T∘ X)  = diag(vec(T))vec(X) = (D^1/2⊗ I + I ⊗ D^1/2)^-1 vec(X).

Then, the equation (∇_Σ_yF)(S) = 0 becomes,

    vec(S)    = vec{∑_k=1^K P_k U_k[T_k∘ U_k^T [ U_y(T_y∘ U_y^T S U_y)U_y^T Σ_kΣ_y^1/2 + Σ_y^1/2Σ_k U_y(T_y∘ U_y^T S U_y)U_y^T] U_k ] U_k^T }
       = ∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1 (U_k^T ⊗ U_k^T) ( Σ_y^1/2Σ_k⊗ I + I ⊗Σ_kΣ_y^1/2)
            (U_y ⊗ U_y) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T) vec(S)

Splitting the term Σ_y^1/2Σ_k⊗ I = (Σ_y^1/2Σ_kΣ_y^1/2⊗ I)· (Σ_y^-1/2⊗ I), we get

    vec(S)    = {[∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k ⊗ I)  (U_k^T ⊗ U_k^T)] ·
         [(U_y ⊗ U_y) (D_y^-1/2⊗ I) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)] + 
         [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(I ⊗ D_k)  (U_k^T ⊗ U_k^T)]·
         [(U_y ⊗ U_y) (I ⊗ D_y^-1/2) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]} vec(S).

Applying the three identities

    D_k⊗ I = (D_k^1/2⊗ I + I⊗ D_k^1/2)(D_k^1/2⊗ I - I⊗ D_k^1/2) + I⊗ D_k,


    I⊗ D_k = (D_k^1/2⊗ I + I⊗ D_k^1/2)^2 - D_k^1/2⊗ I - 2 D_k^1/2⊗ D_k^1/2,


    D_y^-1/2⊗ D_y^-1/2 = (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1(D_y^-1/2⊗ I + I ⊗ D_y^-1/2),


we obtain


    vec(S)    = {[∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I - I ⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·
         [(U_y ⊗ U_y) (D_y^-1/2⊗ I) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]
          + [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(I ⊗ D_k)  (U_k^T ⊗ U_k^T)]·
         [(U_y ⊗ U_y) (D_y^-1/2⊗ D_y^-1/2) (U_y^T ⊗ U_y^T)]} vec(S)
    
       = {[∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·[(U_y ⊗ U_y) (D_y^-1/2⊗ D_y^-1/2) (U_y^T ⊗ U_y^T)]
          - [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I) (U_k^T ⊗ U_k^T)] ·[(U_y ⊗ U_y) (I ⊗ D_y^-1/2) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]
          - [∑_k=1^K P_k (U_k ⊗ U_k) (I ⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·[(U_y ⊗ U_y) (D_y^-1/2⊗ I) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]
          - [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k^1/2⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·
           [(U_y ⊗ U_y) (D_y^-1/2⊗ D_y^-1/2) (U_y^T ⊗ U_y^T)]} vec(S)
    
       = {(Σ_y ⊗ I) ·[(U_y ⊗ U_y) (D_y^-1/2⊗ I) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]
          + (I ⊗Σ_y) ·[(U_y ⊗ U_y) (I ⊗ D_y^-1/2) (D_y^1/2⊗ I + I ⊗ D_y^1/2)^-1 (U_y^T ⊗ U_y^T)]
          - [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k^1/2⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·
         [(U_y ⊗ U_y) (D_y^-1/2⊗ D_y^-1/2) (U_y^T ⊗ U_y^T)]} vec(S)
       = { I - [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k^1/2⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·(Σ_y^-1/2⊗Σ_y^-1/2) } vec(S)

Hence, it follows that

    [∑_k=1^K P_k (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k^1/2⊗ D_k^1/2) (U_k^T ⊗ U_k^T)] ·(Σ_y^-1/2⊗Σ_y^-1/2) vec(S) = O

Denote the lengthy matrix by [∑_k=1^K P_k Y_k ] Y. Then, Y is positive-definite, and since

    (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1(D_k^1/2⊗ D_k^1/2)

are diagonal with positive entries, Y_k are also positive-definite, so the equation holds if and only if S=O. We conclude that the gradient (∇_Σ_y F) is always non-singular, and the implicit function theorem implies that Σ_y depends differentiably on {Σ_1,…Σ_K}⊆∏_i=1^K𝒮_d^+.


It follows that since Σ_k depends differentiably on P_k^i, the derivatives ∂Σ_y / ∂ P_k^i exist.





§ COMPUTATION OF DERIVATIVE


To solve for the gradient ∇_P^iTr[Σ_y], set Λ_k^i = ∂Σ_y / ∂ P_k^i ∈𝒮_d as an unknown variable. Rather artificially, define the term

    Ω_k^i    := 1/P_k∂ (P_k)^2Σ_k/∂ P_k^i = 1/P_k∂ P_k ∑_i=1^N P^i_k (x^i-x_k) · (x^i-x_k)^T /∂ P_k^i
       = Σ_k + (x_i-x_k)· (x_i-x_k)^T + 2∑_i=1^N P^i_k (x^i-x_k) ·(-∂x_k/∂ P^i_k)^T
       = Σ_k + (x_i-x_k)· (x_i-x_k)^T

Taking partial derivative ∂ P_k^i on both sides of the covariance formula (<ref>), we obtain

    Λ_k^i    = ∑_h≠ k(∇(Σ_y^1/2((P_h)^2Σ_h)Σ_y^1/2)^1/2)((∇Σ_y^1/2)(Λ_k^i)((P_h)^2Σ_h)Σ_y^1/2 + Σ_y^1/2((P_h)^2Σ_h)(∇Σ_y^1/2)(Λ_k^i))
          + (∇(Σ_y^1/2((P_k)^2Σ_k)Σ_y^1/2)^1/2)((∇Σ_y^1/2)(Λ_k^i)((P_k)^2Σ_k)Σ_y^1/2 + Σ_y^1/2P_kΩ_k^iΣ_y^1/2 + Σ_y^1/2((P_k)^2Σ_k)(∇Σ_y^1/2)(Λ_k^i))
       = ∑_h=1^K U_h[(P_h)^-1T_h∘ U_h^T [ U_y(T_y∘ U_y^T Λ_k^i U_y)U_y^T (P_h)^2Σ_hΣ_y^1/2 + Σ_y^1/2(P_h)^2Σ_h U_y(T_y∘ U_y^T Λ_k^i U_y)U_y^T] U_h ] U_h^T
          + U_k[(P_k)^-1 T_k∘ U_k^T (Σ_y^1/2P_kΩ_k^iΣ_y^1/2 + Σ_y^1/2) U_k ] U_k^T
       = ∑_h=1^K P_h U_h[T_h∘ U_h^T [ U_y(T_y∘ U_y^T Λ_k^i U_y)U_y^T Σ_hΣ_y^1/2 + Σ_y^1/2Σ_h U_y(T_y∘ U_y^T Λ_k^i U_y)U_y^T] U_h ] U_h^T
          + U_k[T_k∘ U_k^T (Σ_y^1/2Ω_k^iΣ_y^1/2 + Σ_y^1/2) U_k ] U_k^T

Vectorize and simplify it by the previous computations,

    vec(Λ_k^i)    = { I - [∑_h=1^K P_h (U_h ⊗ U_h) (D_h^1/2⊗ I + I ⊗ D_h^1/2)^-1(D_h^1/2⊗ D_h^1/2) (U_h^T ⊗ U_h^T)] ·(Σ_y^-1/2⊗Σ_y^-1/2) }
         · vec(Λ_k^i) + (U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1 (U_k^T ⊗ U_k^T)(Σ_y^1/2⊗Σ_y^1/2) vec(Ω_k^i)

Therefore,

    vec(Λ_k^i)    = (Σ_y^1/2⊗Σ_y^1/2)
         [∑_h=1^K P_h (U_h ⊗ U_h) (D_h^1/2⊗ I + I ⊗ D_h^1/2)^-1(D_h^1/2⊗ D_h^1/2) (U_h^T ⊗ U_h^T)]^-1
         [(U_k ⊗ U_k) (D_k^1/2⊗ I + I ⊗ D_k^1/2)^-1 (U_k^T ⊗ U_k^T)](Σ_y^1/2⊗Σ_y^1/2) · vec(Ω_k^i)

Denote the solution by Λ_k^i = vec^-1(W_k vec(Ω_k^i)), we obtain an expression for the gradient of the objective function

    ∇_P^i Tr[Σ_y] = ∑_k=1^K Tr[Λ_k^i]_k= ∑_k=1^K vec(I)^T · W_k · vec(Ω_k^i)_k






§ PROOF OF THEOREM <REF>



We closely follow the proof of the variance decomposition theorem in <cit.>. Let us assume that = {1,… K} with clusters ρ_k and weights P_k, and that all clusters ρ_k ∈_2,ac(^d). It follows from Theorem 6.1 of <cit.> that the barycenter μ is unique and μ∈_2,ac(^d). The general formula (<ref>) will follow from the standard approximation method in <cit.>.

By Theorem 2.1 of <cit.>, as ρ_k,μ are all Gaussian, there exists a unique transport map T_k that transports ρ_k to μ and has the form (<ref>). Similarly, there are unique optimal transport maps T_kh from any ρ_k to any ρ_h. We denote them by

    T_k(x) = A_kx+b_k,  T_kh = A_khx + b_kh

where the A's are positive-definite matrices. Consider T^-1_kh∘ T_h^-1: It is an affine map that transports ρ_k back to itself and has a positive-definite matrix A_kh^-1A_h^-1A_k. Since ρ_k is a non-degenerate Gaussian, there is only one such affine map, namely the identity (One could examine each eigensubspace of A_kh^-1A_h^-1A_k and consider how it operates there). It follows that T_h^-1∘ T_k is exactly the optimal transport map T_kh.

Define the random variable Y∼μ and the random variables X_k = T_k^-1(Y) ∼ρ_k. Then, by the optimality of T_k,

    ∫ W_2^2(ρ(·|z),μ) dν(z) = ∑_k=1^K P_k [Y-X_k^2]

Meanwhile, the identity T_h^-1∘ T_k = T_kh implies that X_h = T_kh(X_k) and that W_2^2(ρ_k,ρ_h) = [X_k-X_h^2]. Hence,

    ∬ W_2^2(ρ(·|z_1),ρ(·|z_2)) dν(z_1)dν(z_2) = ∑_k,h=1^K P_k P_h [X_k-X_h^2]

Yet, Remark 3.9 of <cit.> implies that Y is exactly the average of X_k:

    Y = ∑_k=1^K P_k X_k

which leads to

    ∑_k=1^K P_k [Y-X_k^2] = ∑_k,h=1^K P_k P_h [X_k-X_h^2]

Hence, we have established equality (<ref>).







§ DERIVATION OF GRADIENTS IN SECTION <REF>



For each z, the partial derivative of the conditional standard deviation σ(z) times the marginal latent distribution ν with respect to the conditional latent distribution ν_i = ν(z|x_i) is given by

    ∂(σ(z)ν(z))/∂ν_i   = ∑_j=1^N ||x_j - x(z)||^2 ν_j + ||x_i - x(z)||^2 ·∑_j=1^N ν_j - 2∑_j=1^N ν_j (x_j-x)^T·∂x(z)/∂ν_i/ 2N (∑_j=1^N ν_j ·∑_j=1^N ||x_j - x(z)||^2 ν_j)^1/2
       = 1/2N[σ(z)+||x_i-x(z)||^2/σ(z)]


Meanwhile, for each ν_i, the partial derivatives of (<ref>) are given by

    ∂ν_i(z)/∂z_j = [ -z_j/||z||^3 + N/α^2z_j ||z-z_i||^2/||z||^4 + 1_i=jN/α^2z-z_i/||z||^2] ν_i(z)

So the Jacobian matrix J_zp of the vector p:=[ν_i] with respect to the parameter vector z:=[z_j] is

    J^T_zp(z)    = z·[ ( -1/||z||^3 + N/α^2||z-z_i||^2/||z||^4) ν_i(z) ]^T_i
    + diag( N/α^2z-z_i/||z||^2ν_i(z) )
       = z·[ ( -1/||z|| + ||z-z_i||^2/ϵ^2) ν_i(z)/||z||^2]^T_i
    + diag(  z-z_i/ϵ^2ν_i(z) )


To show that we can differentiate under the integral sign of σ in (<ref>), we rewrite the standard deviation σ(z) into

    σ(z) = [1/2∬ ||x-y||^2 dρ(x|z)dρ(y|z)]^1/2 = (1/2∑_i,j=1^N ||x_i-x_j||^2 ν_i ν_j)^1/2
    = (∑_i,j=1^N ||x_i-x_j||^2 ν_i ν_j )^1/2/√(2)∑_i=1^N ν_i

Define the matrix D = [D_ij] = [||x_i-x_j||^2]. Then, the distribution term dν(z) in σ can be cancelled, and we obtain a simpler formula for σ:

    σ   = 1/√(2)N∫(∑_i,j=1^N D_ijν_iν_j )^1/2 dz = 1/√(2)N∫ ||√(D)p|| dz

It is straightforward to show that for each z_0 ∈ℝ^N (z≠0), there exists some compact neighborhood U (z_0 ∈ U^o ⊆U⊆ℝ^N-{0}) such that the integrand ||√(D)p|| and its z gradient are uniformly bounded by some integrable function. Thereby, Theorem 3.2 of <cit.> shows that at z_0, we are allowed to take derivatives under the integral sign:

    ∇_zσ
    =   ∫∇_z(σ(z) ν(z)) dz = ∫∇_z(σ(z)ν(z))/ν(z) dν(z)
    
    =   ∫1/ν(z) J^T_zp(z) ·∇_p(σ(z)ν(z)) dν(z)
    ∂σ/∂z_i
    =   ∫1/2N·ν(z){z_i/||z||^2∑_j=1^N [σ(z)+||x_j-x(z)||^2/σ(z)] ·[ ||z-z_j||^2/ϵ^2 - 1/||z||] ν_j(z)
       + z-z_i/ϵ^2ν_i(z) [σ(z)+||x_i-x(z)||^2/σ(z)] } dν(z)

In particular, since the Jacobian J_zp consists of a rank-one matrix and a diagonal matrix, computing the above integrand for any z takes only linear time, O(N).



