Hierarchical Optimal Transport for Multimodal Distribution Alignment
John Lee Max Dabagia Eva L. Dyer Christopher J. Rozell


Signatures of the orbital angular momentum of an infrared light beam in the two-photon transition matrix element: A step toward attosecond chronoscopy of photoionization
    Gopal Dixit
    Received: date / Revised version: date
=========================================================================================================================================================================








In many machine learning applications, it is necessary to meaningfully aggregate, through alignment, different but related datasets. Optimal transport (OT)-based approaches pose alignment as a divergence minimization problem: the aim is to transform a source dataset to match a target dataset using the Wasserstein distance as a divergence measure under alignment constraints. We introduce a hierarchical formulation of OT which leverages clustered structure in data to improve alignment in noisy, ambiguous, or multimodal settings. To solve this numerically, we propose a distributed ADMM algorithm that  exploits the Sinkhorn distance, thus it has an efficient computational complexity that scales quadratically with the size of the largest cluster. When the transformation between two datasets is unitary, we provide performance guarantees that describe when and how well  cluster correspondences can be recovered with our formulation, and then describe the worst-case dataset geometry for such a strategy. We apply this method to synthetic datasets that model data as mixtures of low-rank Gaussians and study the impact that different geometric properties of the data have on alignment. Next, we applied our approach to a neural decoding application where the goal is to predict movement directions and instantaneous velocities from populations of neurons in the macaque primary motor cortex. Our results demonstrate that when clustered structure exists in datasets, and is consistent across trials or time points, a hierarchical alignment strategy that leverages such structure can provide significant improvements in cross-domain alignment.






§ INTRODUCTION



In many machine learning applications, it is necessary to meaningfully aggregate, through alignment, different but related datasets (e.g., data across time points or under different conditions or contexts).
Alignment is an important problem at the heart of transfer learning <cit.>, point set registration <cit.>, and shape analysis <cit.>, but is generally NP hard. In recent years, distribution alignment methods that use optimal transport (OT) to quantify similarity between two distributions have increased in popularity due to their attractive mathematical properties and impressive performance in a variety of tasks <cit.>. However, using OT to solve unsupervised distribution alignment problems that must simultaneously match two datasets' distributions (using OT) while also learning a transformation between their latent spaces, is extremely challenging, especially when the data has complicated multi-modal structure. Leveraging additional structure in the problem is thus necessary to regularize OT and constrain the solution space.

Here, we leverage the fact that heterogeneous datasets often admit  clustered or multi-subspace structure to improve OT-based distribution alignment. Our solution to this problem is to simultaneously estimate the cluster alignment across two datasets using their local geometry, while also solving a global alignment problem to meld these local estimates. While it is advantageous to regularize the OT problem with known cluster pairings <cit.>, we are instead concerned with the substantially harder unsupervised setting where such information is missing. We introduce a hierarchical formulation of OT for clustered and multi-subspace datasets called Hierarchical Wasserstein Alignment (HiWA)[MATLAB code can be found at https://github.com/siplab-gt/hiwa-matlabhttps://github.com/siplab-gt/hiwa-matlab. Neural datasets and Python code are provided at http://nerdslab.github.io/neuralignhttp://nerdslab.github.io/neuralign].

We empirically show that when data are well approximated with Gaussian mixture models (GMMs) or lie on a union of subspaces, we may leverage existing clustering pipelines (e.g., sparse subspace clustering <cit.> <cit.>) to improve alignment.
When the transformation between datasets is unitary, we provide analyses that reveal key geometric and sampling insights, as well as perturbation and failure mode analyses. To solve the problem numerically, we propose an efficient distributed ADMM algorithm that also exploits the Sinkhorn distance, thus benefiting from efficient computational complexity that scales quadratically with the size of the largest cluster.

To test and benchmark our approach, we applied it to synthetic data generated from mixtures of low-rank Gaussians and studied the impact of different geometric properties of the data on alignment to confirm the predictions of our theoretical analysis. 
Next, we applied our approach to a neural decoding application where the goal is to predict movement directions from populations of neurons in the macaque primary motor cortex. Our results demonstrate that when clustered structure exists in neural datasets and is consistent across trials or time points, a hierarchical alignment strategy that leverages such structure can provide significant improvements in unsupervised decoding from ambiguous (symmetric) movement patterns. This suggests OT can be applied to a wider range of neural datasets, and shows that a hierarchical strategy avoids local minima encountered by a global alignment strategy that ignores clustered structure.






§ BACKGROUND AND RELATED WORK




Transfer learning and distribution alignment.
A fundamental goal in transfer learning is to aggregate related datasets by learning a mapping between them.
We wish to learn a transformation T ∈, where  refers to some class of transformations that aligns distributions under a notion of probability divergence (·|·)  between a target distribution μ and a reference (source) distribution ν:

    min_T ∈(T(μ)|ν) .


Various probability divergences have been proposed in the literature, such as Euclidean least-squares (when data ordering is known) <cit.>, Kullback-Leibler (KL) <cit.>, maximum mean discrepancy (MMD) <cit.>, and the Wasserstein distance <cit.>, where trade-offs are often statistical (e.g., consistency, sample complexity) versus computational.
Alignment problems are ill-posed since the space of 𝒯 is large, so a priori structure is often necessary to constrain 𝒯 based on geometric assumptions. Compact manifolds like the Grassmann or Stiefel <cit.> are primary choices when little information is present, as they preserve isometry. Non-isometric transformations, though richer, demand much more structure (e.g., manifold or graph structure) <cit.>. 

Low-rank and union of subspaces models.
Principal components analysis (PCA), one of the most popular methods in data science, assumes a low-rank model where the top-k principal components of a dataset provide the optimal rank-k approximation under an Euclidean loss.
This has been extended to robust (sparse errors) settings <cit.>, and multi- (union of) subspaces settings where data can be partitioned into disjoint subsets where each subset of data is locally low-rank <cit.>.
Transfer learning methods based on subspace alignment <cit.> work well with zero-mean unimodal datasets, but struggle on more complicated modalities (e.g., Gaussian mixtures or union of subspaces) due to a mixing of covariances.
Related to our work, <cit.> performs multi-subspace alignment by greedily assigning correspondences between subspaces using chordal distances; this however discards valuable information about a distribution's shape.

Optimal transport.
Optimal transport (OT) <cit.> is a natural type of divergence for registration problems because it accounts for the underlying geometry of the space.
In Euclidean settings, OT gives rise to a metric known as the Wasserstein distance (μ,ν) which measures the minimum effort required to “displace” points across measures μ and ν (understood here as empirical point clouds).
Therefore, OT relieves the need for kernel estimation to create an overlapping support of the measures μ,ν.
Despite this attractive property, it has both a poor numerical complexity of O(n^3 log n) (where n is the sample size) and a dimension-dependent sample complexity of O(n^-1/d), where the data dimension is d <cit.>.
Recently, an entropically regularized version of OT known as the Sinkhorn distance <cit.> has emerged as a compelling divergence measure; it not only inherits OT's geometric properties but also has superior computational and sample complexities of O(n^2) and O(n^-1/2)[Dependent on a regularization parameter <cit.>.], respectively.
It has also become a versatile building block in domain adaptation <cit.>.
Prior art <cit.> has largely exploited the OT's push-forward as the alignment map since this map minimizes the OT cost between the source and target distributions while allowing a priori structure to be easily incorporated (e.g., to preserve label/graphical integrity). Such an approach, however, is fundamentally expensive when d ≪ n since the primary optimization variable is a large transport coupling (i.e., ^n × n), while in reality the alignment mapping is merely ^d ↦^d.
Moreover, it assumes that the source and target distributions are close in terms of their squared Euclidean distance (i.e., an identity transformation), but this does not generally hold between arbitrary latent spaces.

Hierarchical OT and related work. The idea of learning an affine or unitary transformation to align datasets with an OT-based divergence has previously been studied in <cit.>, a problem known as OT Procrustes. However, these methods don't use problem-specific or clustered structure in data. Hierarchical OT is a recent generalization of OT <cit.> that is an effective and efficient way of injecting structure into OT but it has never been used to jointly solve alignment problems – our work represents a first attempt at doing so. Thus, a key contribution of this paper is putting both of these two ingredients together to develop a scalable strategy that leverages multimodal structure in data solve the OT Procrustes problem.









§ HIERARCHICAL WASSERSTEIN ALIGNMENT




Preliminaries and notation.
Consider clustered datasets {X̌_i ∈^D × n_x,i}_i=1^S and {Y̌_j ∈^D × n_y,j}_j=1^S whose clusters are denoted with the indices i,j and whose columns are treated as ^D embedding coordinates.
The number of samples in the i-th (j-th) cluster of dataset X̌ (dataset Y̌) is given by  n_x,i (n_y,j). We express the empirical measures of clusters X̌_i and Y̌_j as 
μ_i := 1/n_x,i∑_k=1^n_x,iδ_X̌_i(k) and 
ν_j := 1/n_y,j∑_l=1^n_y,jδ_Y̌_j(l), respectively,
where δ_x̌ refers to a point mass located at coordinate x̌∈^D.
The squared 2-Wasserstein distance between μ_i and ν_j is defined as

    _2^2(μ_i,ν_j)
        :=
        min_Q̌∈(n_x,i,n_y,j)∑_k=1^n_x,i∑_l=1^n_y,j
        Q(k,l) X̌_i(k) - Y̌_j(l)_2^2

where Q̌ is a doubly stochastic matrix that encodes point-wise correspondences (i.e., the (k,l)-th entry describes the flow of mass between δ_X̌_i(k) and δ_Y̌_j(l)), 
X̌_i(k) is the k-th column of matrix X̌_i, and the constraint (m,n) := {Q̌∈^m × n_+ : Q̌_n = _m/m , Q̌^⊤_m = _n/n } refers to the uniform transport polytope (with _m a length m vector containing ones). We will use · to denote the operator norm, X̌^† to denote the pseudo-inverse of X̌, and Ǐ_d to denote the d × d identity matrix. 


Overview.
Although unsupervised alignment is challenging due to the presence of local minima, the imposition of additional structure will help to prune them away.
Our key insight is that hierarchical structure decomposes a complicated optimization surface into simpler ones that are less prone to local minima.
We formulate a hierarchical Wasserstein approach to align datasets with known (or estimated) clusters {μ_i}_i=1^S,{ν_j}_j=1^S but whose correspondences are unknown.
The task therefore is to jointly learn the alignment T and the cluster-correspondences:

    min_P̌∈_S , T ∈∑_i=1^S ∑_j=1^S P_ij_2^2 ( T (μ_i) , ν_j )  ,

where the matrix P̌ encodes the strength of correspondences between clusters, with a large P_ij value indicating a correspondence between clusters i,j, and a small value indicating a lack thereof.
We note that _S := (S,S) is a special type of transport polytope known as the S-th Birkhoff polytope.
Interestingly, this becomes a nested (or block) OT formulation, where correspondences are resolved at two levels: the outer level resolves cluster-correspondences (via P̌) while the inner level resolves point-wise correspondences between cluster points (via the Wasserstein distance).


Alignment over the Stiefel manifold.
Assuming clusters lie on subspaces and principal angles between subspaces are “well preserved” across X̌ and Y̌ (we make this precise in Theorem <ref>), an isometric transformation suffices.
Hence, we solve (<ref>) with _D,D, the Stiefel manifold which is defined as _k,d := {Ř∈^k × d : Ř^⊤Ř = Ǐ_d }. Explicitly, we can re-formulate equation (2) as:

    min_P̌ , Ř , {Q̌_ij}∑_i,j P_ij C_ij (Ř,Q̌_ij)
          s.t.  P̌∈_S ,   Ř∈_D,D ,   Q̌_ij∈(n_x,i,n_y,j),




    where       C_ij (Ř,Q̌_ij)
           :=
        1/D∑_k,lQ̌_ij(k,l) ŘX̌_i(k) - Y̌_j(l)_2^2

measures pairwise cluster divergences using the squared 2-Wasserstein distance under a Stiefel transformation Ř acting on the i^ th cluster.


Finally, we include entropic regularization over transportation couplings P̌ and all Q̌_ij's to modify the Wasserstein distances to Sinkhorn distances, so as to take advantage of its superior computational and sample complexities. Omitting constraints for brevity, our final problem is given as

    min_P̌,Ř,{Q̌_ij}∑_i,j( P_ij C_ij (Ř,Q̌_ij) + H_γ_2(Q̌_ij) )
        + H_γ_1(P̌) ,

where γ_1,γ_2>0 are the entropic regularization parameters and the negative entropy function is defined as H_γ(P̌) := γ∑_i,j P_ijlog P_ij.
Parameters γ_1,γ_2 control the correspondence entropy, therefore (<ref>) approximates (<ref>) when γ_1,γ_2 > 0, but reverts to the original problem (<ref>) as γ_1,γ_2 → 0.


Distributed ADMM approach.
Problem (<ref>) is non-convex due to multilinearity in the objective and its Stiefel manifold domain.
Although alternating directions method of multipliers (ADMM) is a convergent convex solver framework <cit.>, it is being applied in increasingly many non-convex settings <cit.>.
Since (<ref>) readily admits a splitting structure that separates the individual C_ij blocks, we develop a distributed ADMM approach.
We proceed to split (<ref>) as follows:

    min_P̌,Ř,{Ř_ij,Q̌_ij}∑_i,j( P_ij C_ij(Ř_ij,Q̌_ij) + H_γ_2(Q̌_ij) )
        + H_γ_1(P̌) 
          s.t.  Ř_ij = Ř,
        ∀ i,j ,

noting that the set constraints are omitted for brevity.
The augmented Lagrangian is given by

    ℒ_μ
        = 
           ∑_i,j(
        P_ij C_ij(Ř_ij,Q̌_ij)
        + ⟨μ/DΛ̌_ij , Ř_ij - Ř⟩
        + μ/2DŘ_ij - Ř_F^2 
        + H_γ_2(Q̌_ij)
        ) 
        + H_γ_1(P̌) ,

where μ>0 is the ADMM parameter and {Λ_ij} are Lagrange multipliers.
Full details of the update steps are included in the Supplementary Material.
The algorithm may be summarized in two steps (Alg. <ref>): (i) a distributed step that asks all cluster pairs to individually find their optimal transformations Ř_ij in parallel, and (ii) a consensus step that aggregates all the locally estimated transformations according to a weighting that is proportional to correspondence strengths P_ij.




Parameters.
Entropic parameters γ_1,γ_2 relax the one-to-one cluster correspondence assumption, balancing a trade off between alignment precision (small γ) and sample complexity (large γ).
Numerically, negative entropy adds strong convexity to the program, reducing sensitivity towards perturbations at the cost of a slower convergence rate.
The ADMM parameter μ controls the `strength' of the consensus, or from an algorithmic viewpoint, the gradient step size.

Distributed consensus.
Update steps for Q̌_ij, Ř_ij,Ľ_ij can be performed in parallel over all cluster pairs (S^2 in total), making it amenable for a distributed implementation. The runtime complexity of this algorithm is presented in the supplementary Materials.

Robustness against initial conditions.
We intentionally build robustness against initial conditions by ordering updates for Ř_ij and Q̌_ij before P̌ such that when μ is sufficiently small, the ADMM sequence is influenced more by the data than by initial conditions.









§ THEORETICAL GUARANTEES FOR CLUSTER-BASED ALIGNMENT





While the previous section explains how to align clustered datasets, in this section, we aim to answer the question of when and how well they can be aligned.
We provide necessary conditions for cluster-based alignability as well as alignment perturbation bounds according to equation (<ref>)'s formulation.
To simplify our analysis, we make the following assumptions:
(i) each of the clusters contain the same number of datapoints n,
(ii) the ground truth cluster correspondences are P̌^⋆ = Ǐ_S / S (i.e., diagonal containing 1/S). However, this analysis can be extended to the case where the number of points is unequal without loss of generality.
Detailed proofs are given in the Supp. Material.

The following result is a criterion that, if met, ensures the existence of a global minimizer of the cluster-correspondence P̌^⋆. This criterion requires that matched clusters must be closer in Wasserstein distance than mismatched clusters, according to a threshold determined by Wasserstein's sample complexity (i.e., an asymptotic rate dependent on the clusters' sample sizes and intrinsic dimensions).
Since these sample complexity results are based on the Wasserstein distance, we expect a less stringent criterion when using the Sinkhorn distance in (<ref>) (due to superior sample complexity <cit.>).




Let all clusters be strictly low-rank where the dimension of the i-th cluster in the x-th dataset is d_x,i.
Let d_x,i,d_y,j > 4, ∀ i,j ∈ S.
Define Ĉ^⋆_ij := min_Ř∈_D,D, Q̌_ij∈_n C_ij (Ř,Q̌_ij).
Problem (<ref>) yields the solution P̌^⋆ = Ǐ_S/S with probability at least 1-δ if, ∀ i,j : i ≠ j, the following criterion is satisfied:

    Ĉ_ij^⋆
        + Ĉ_ji^⋆
        - Ĉ_ii^⋆
        - Ĉ_jj^⋆  >  
               B_x,i(δ) + B_y,i(δ)
             + B_x,j(δ) + B_y,j(δ)




    where   B_z,k(δ) :=   c_z,k n^ -2/d_z,k
                          + √(log (1/δ) / 2 n ) ,
             
        c_z,k = 1458 ( 2 + 1/3^d_z,k/2 - 2 - 1) .



The proof contains two parts.
In the first part, we consider perturbation conditions of the cost matrix Č in a (non-variational) optimal transport program over the Birkhoff polytope.
To be unperturbed from P̌^⋆ = Ǐ_S/S, we require that C_ij + C_ji - C_ii - C_jj > 0, ∀ i,j : i ≠ j.
In the second part, we extend this condition to the the finite-sample regime by utilizing recently developed concentration bounds <cit.> for the p-Wasserstein distance, which essentially raises the disambiguity lower bound due to finite-sample uncertainty. (Supp. Material, Section 2)




Now, even if we know the global correspondence  P̌^⋆, we still do not have the full picture about the alignment's quality.
For example, all matching clusters may have very similar covariances, but principal angles between the clusters are “distorted” across the datasets.
Our next theorem gives us an upper bound on the alignment error (for unitary transformations), and makes precise the notion of global structure distortion.




Consider data matrices {X̌_i,Y̌_i ∈^D × n}_i=1^c with known point-wise correspondence matrices {Q̌_ii∈_n}_i=1^c.
Define matrices

    X̌ := [X̌_1 Q̌_11 , X̌_2 Q̌_22 , … , X̌_c Q̌_cc] , 
            Y̌ := [Y̌_1 , Y̌_2 , … , Y̌_c] .

Set ε^2 := Y̌^⊤Y̌ - X̌^⊤X̌_F.
If the criterion stated in theorem <ref> is satisfied,
X̌ is full row rank, and εX̌^†≤1/√(2) (X̌X̌^†)^-1/2, then

    min_P̌∈_c , Ř∈_D,D∑_i,j P_ij C_ij (Ř)
          ≤  
        ( X̌X̌^† + 2 )^2 X̌^†^2 ε^4 + D,

where D = ∑_i=1^c  (   X̌_i ( Ǐ / n - Q̌_iiQ̌_ii^⊤ ) X̌_i^⊤ + ( 1/n - 1 ) Y̌_i Y̌_i^⊤ ) is a data-dependent constant.


We utilize a recent perturbation result on the Procrustes problem (on a Frobenius norm objective) by Arias-Castro et al. <cit.> and adapt it to our squared 2-Wasserstein objective. (Supp. Material, Section 3)




Note that ε plays a major role in the alignment error bound and quantifies the notion of global structure distortion, which allows us to understand on how phenomena like covariate shift or misclustering impacts alignment.
To shed some light in this regard, we consider a simple analysis on a cluster-pair's error contribution to ε, denoted as ε_ij.
Consider the decomposition of the (i,j)-th block of the Gramians related to clusters i and j, where their respective singular value decompositions are
X̌_i Q̌_ii = Ǎ_i Σ̌_x,iV̌^⊤ and
Y̌_j = B̌_j Σ̌_y,jV̌^⊤.
Defining the blockwise error between clusters i,j as

    ε_ij
        :=
        Y̌_i^⊤Y̌_j - Q̌_ij^⊤X̌_i^⊤X̌_j Q̌_jj_F
        =
        Σ̌_y,iB̌_i^⊤B̌_j Σ̌_y,j 
              - Σ̌_x,iǍ_i^⊤Ǎ_j Σ̌_x,j_F ,

two components stand out: (i) angular shift, which is characterized by differences in principal angles between B̌_i^⊤B̌_j and Ǎ_i^⊤Ǎ_j, and (ii) spectral shift, which is characterized by differences in spectra.

Finally, we show that the subspace configuration of a dataset's clusters can also affect alignment.
Pretend for a moment that external alignment information were present to aid in the disambiguation between two clusters.
The following lemma tells us when such information is useless (Proof in Supp. Material, Section 4).




Consider clusters X̌_i, Y̌_j ∈^D × n and known point-wise correspondences Q̌_ij∈(n,n).
Denote the left and right singular vectors of Y̌_j Q̌_ij^⊤X̌_i^⊤ associated with the non-zero singular values as Ǔ̃ , Ṽ̌̃∈^D × r with r ≤ D.
Define the set of orthogonal transformations that are constrained to agree with known angular directions as

    (Ǔ',V̌') :=
        {Ř∈^D × D_+ :
           Ř^⊤Ř = Ǐ , ŘV̌' = Ǔ' } ,

where Ǔ',V̌' ∈_D,r with r ≤ D.
Given Ǔ',V̌' ∈^D × r' with r' ≤ D, we have

    min_Ř∈(Ǔ',V̌') C_ij(Ř)
        ≥min_Ř∈_D,D C_ij(Ř) ,

with equality holding when ⟨Ǔ̃ , Ǔ' ⟩ = ⟨Ṽ̌̃ , V̌' ⟩. 
 


Direct consequences of this lemma are the following:
When a dataset has equally-spaced subspaces, it has a maximally uninformative geometric configuration since angular information from other clusters (i.e., Ǔ',V̌') can never increase the inter-cluster distance C_ij (i.e., equality in (<ref>) always holds); it is hence a worst-case scenario for alignment.
This also explains why alignment in very high-dimensional space is harder: All subspaces may be orthogonal to each other, and hence offer no “geometric” advantage.









§ NUMERICAL EXPERIMENTS







 §.§ Synthetic low-rank Gaussian mixture dataset







In this section, we validate our method as well as demonstrate its limiting characteristics under symmetric-subspace and finite-sample regimes.
To generate our synthetic data, we repeat the following procedure for each of the S clusters. We first randomly generate Gaussian distribution parameters μ_i ∈^d,Σ_i ∈^d : Σ_i ≽ 0 (positive semi-definite), then randomly sample n data-points from these parameters, and finally project them into a random subspace V̌_i ∈^D × d in a D>d dimensional embedding.
In these experiments, we assume that the  clusters are known, but the cluster-correspondence across datasets is unknown. We measure performance with respect to two metrics:
(i) alignment error, defined as the relative difference between the recovered versus true rotation acting on the data Ř̂X̌ - Ř^⋆X̌_F^2 / Ř^⋆X̌_F^2, and
(ii) correpondence error, defined as the sum of absolute differences between the recovered and the true correspondences ∑_ij |P̂̌̂-P̌^⋆|_ij.

To understand how global geometry impacts alignment,
we applied HiWA in two different settings (Figure <ref>a-b): (i) a worst-case setting where subspaces are equally spaced with a subspace similarity of V̌_i^⊤V̌_j = 1, ∀ i ≠ j, and (ii) the random setting where subspaces are randomly selected from the Grassmann manifold. 
We observe that equally-spaced subspaces have significantly inferior performance when compared to randomly-spaced subspaces, providing some evidence that equally spaced subspaces are indeed the worst-case scenario in alignment, as suggested by Lemma <ref>.

Next, we studied the effect of dimensions d and sample size n on the accuracy of alignment (Figure <ref> (c-d)). We tested HiWA across various dataset conditions by varying parameters d={2,3,4,5} and  n={12,25,50,100,200} while approximately maintaining the average subspace correlations (i.e., 𝔼V̌_i^⊤V̌_j) by fixing the cluster size S=5 and tuning D to control the subspace spacing. In both cases, sample complexities are better than the theoretical rate of O(n^-1/d), which is likely due to the Sinkhorn distance's superior sample complexity.
In Figure <ref>e-f, we conduct an ablation study and evaluate our algorithm against benchmark methods in transfer learning and point set registration in two settings: a simple one in low-d (e) and a harder one in higher-d (f). Specifically, we compare HiWA when clusters are known (but pairwise correspondences are unknown), HiWA with clustering via sparse subspace clustering <cit.> (HiWA-SSC) to represent completely unsupervised alignment, a Wasserstein alignment variant with no cluster-structure (WA) which is akin to OT Procrustes <cit.>, subspace alignment <cit.>, correlation alignment <cit.>, and iterative closest point (ICP) <cit.>.
HiWA exhibits strongest performance, with HiWA-SSC trailing closely behind (since clusters are independently resolved), followed by WA, then other algorithms.
Subspace alignment methods have remarkably poor performance in higher dimensions due to their inability to resolve subspace sign ambiguities, while ICP demonstrates its notorious dependence on good initial conditions.
These results indicates HiWA's strong robustness against initial conditions and good scaling properties.

  




 §.§ Neural population decoding example



Decoding intent (e.g., where you want to move your arm) or evoked responses (e.g., what you are looking at or listening to) directly from neural activity is a widely studied problem in neuroscience, and the first step in the design of a brain machine interface (BMI). A critical challenge with BMIs is that neural decoders need to be recalibrated (or re-trained) due to drift in neural responses or electrophysiology measurements/readouts <cit.>. A recent method for semi-supervised brain decoding finds a transformation between projected neural responses and movements by solving a KL-divergence minimization problem <cit.>. Using this approach, one could build robust decoders that work across days and shifts in neural responses through alignment.





We test the utility of hierarchical alignment for neural decoding on datasets collected from the arm region of primary motor cortex of a non-human primate (NHP) during a center out reaching task <cit.>. After spike sorting and binning the data, we applied factor analysis to reduce the data dimensionality to 3D (source distribution) and applied HiWA to align the neural data to a 3D movement distribution (target distribution) (Figure <ref>). We compared its performance to (procrustes) Wasserstein alignment (WA) without hierarchical structure, and a baseline brute force search method called distribution alignment decoding (DAD) <cit.>. We examined the prediction accuracy of the target reach direction for the motor decoding task (i.e., the cluster classification accuracy).

Next, we examined the impact of the sampling density (Figure <ref>b) on alignment performance. Our results demonstrate that HiWA continues to produce consistent cluster correspondences (> 70% accuracy), even as the number of samples per cluster drops to 8.
In comparison, DAD is competitive at larger sample sizes but its performance rapidly drops off as sampling density decreases because it requires estimating a distribution from samples. WA suffers from the presence of many local minima and fails to find the correct cluster correspondences. Our results suggest that HiWA consistently provides stable solutions, outperforming competitor methods for this application.

Finally, to study the impact of local and global geometry on whether an unlabeled source and target can be aligned, we applied HiWA to permutations of eight subsets of reach directions (movement patterns). When just two reach directions are considered (Figure <ref>c, Columns 1-4), global geometry becomes useless in determining the correct rotation. In this case, we observe that HiWA is only capable of consistent alignment when cluster asymmetries are sufficiently extreme in both the source and target. When three reach directions are considered (Figure <ref>c, Columns 5-8), the global geometry can be used, yet there still exist symmetrical cases where recovering the correct rotation is difficult without adequate local asymmetries or some supervised (labeled) data to match clusters. These results suggest that hierarchical structure can be critical in resolving ambiguities in alignment of globally symmetric movement distributions.








§ CONCLUSION




This paper introduces a new method for hierarchical alignment with Wasserstein distances, provided an efficient numerical solution with analytical guarantees. We tested our method and compared its performance against other methods on a synthetic mixture model dataset and on a real neural decoding dataset. Future directions include extensions to non-rigid transformations, and applications to higher dimensional neural datasets that do not rely on external measured behavioral covariates.




  §.§.§ Acknowledgments



JL was supported by DSO National Laboratories of Singapore, ED and MD were supported by NSF grant IIS-1755871, and CR was supported by NSF grant CCF-1409422 and CAREER award CCF-1350954.



unsrt











§ SUPPLEMENTARY MATERIALS




§ ADMM ALGORITHMIC DETAILS




 §.§ Derivation of ADMM update steps


ADMM admits the following sequence of updates:

    (Ř_ij , Q̌_ij )
           ←Ř_ij∈ S(D,D) 
    Q̌_ij∈(n_x,i,n_y_j)
        P_ij C_ij(Ř_ij,Q̌_ij) 
        + μ/2DŘ_ij - Ř + Λ̌_ij_F^2
        + H_γ_2(Q̌_ij) ,
        
    P̌   ←P̌∈_c∑_i,j P_ij
          C_ij(Ř_ij,Q̌_ij)
        + H_γ_1(P̌) ,
        
    Ř   ←Ř∈ S(D,D)∑_i,jŘ_ij - Ř + Λ̌_ij_F^2 ,
        
    Λ̌_ij   ←Λ̌_ij + Ř_ij - Ř .


Update (<ref>) involves an alternating minimization over Q̌_ij and Ř_ij whereby the first variable is fixed while the second is minimized, followed by the second fixed and the first minimized, and the procedure is repeated until convergence is achieved.
When solving for Ř_ij we have the following Stiefel manifold optimization:

    Ř_ij   ←Ř_ij∈ S(D,D)
        P_ij C_ij(Ř_ij,Q̌_ij) +
        μ/2DŘ_ij - Ř + Λ̌_ij_F^2 
       =
        Ř_ij∈ S(D,D)(
        ( 2 P_ijY̌_j Q̌_ij^⊤X̌_i^⊤ 
              + μ ( Ř - Λ̌_ij ) )
        Ř_ij^⊤)
        = ǓV̌^⊤ ,

where 2 P_ijY̌_j Q̌_ij^⊤X̌_i^⊤ + μ ( Ř - Λ̌_ij ) = ǓĎV̌^⊤ is its SVD.
We employ the Sinkhorn algorithm (Algorithm 1 of <cit.>) to solve for Q̌_ij, using an entropic parameter of γ_2/P_ij and uniform marginals.
Ř_ij and Q̌_ij are retrieved once the alternating minimization converges.

Update (<ref>) also employs the Sinkhorn algorithm over the cost matrix generated by C_ij(Ř_ij,Q̌_ij) using variables found in update (<ref>), along with an entropic parameter of γ_1 and uniform marginals.

Update (<ref>) is a consensus update over a Stiefel manifold optimization:

    Ř←Ř∈ S(D,D)∑_i,jŘ_ij - Ř + Λ̌_ij_F^2
        = ǓV̌^⊤ ,

where ∑_i,jŘ_ij + Λ̌_ij = ǓĎV̌^⊤ is its SVD.




 §.§ Computational complexity of distributed ADMM algorithm


The main computational complexity of the ADMM algorithm comes from line 8, which demands the solving of S^2 Sinkhorn problems per ADMM iteration.
Fortunately, these Sinkhorn problems (i.e., update steps for Q̌_ij, Ř_ij,Ľ_ij) may be conducted in parallel, making it amenable for a distributed implementation. When fully parallelized, the algorithm has a per-iteration computational complexity of O(n_i n_j), where n_i,n_j refers to the number of points in the largest clusters of X̌,Y̌ respectively (compared to vanilla Sinkhorn's O(n_x n_y) complexity where n_x,n_y refers to the total number of points in respective datasets, assuming D ≪max(n_i,n_j)). In Figure <ref>, we record run times of two versions of the same algorithm, with and without parallelism, operating on a similar test dataset, at varying cluster sizes S (50 data points per cluster). 10 random trials were conducted and the means and standard deviations were reported (error lines). As expected, the parallel implementation scales better. These results were computed using an i7 intel processor with 6 cores clocked at 2.6 GHz.






 §.§ Stopping criteria.


In lines <ref> and <ref> of Algorithm <ref>, possible stopping criteria are (i) Ř - Ř_F ≤τ where the difference is between the current and previous iteration's transformation and τ is the tolerance, and (ii) t ≤ T where T is the maximum number of iterations.





§ PROOF OF THEOREM <REF>




 §.§ Part 1: Geometric perturbations conditions over the Birkhoff polytope


First, we provide the following lemma that illuminates some basic geometrical insights of the general OT cost matrix, whose coupling is in the Birkhoff polytope _c := {P̌∈^c × c : P̌_c = P̌^⊤_c = _c }. We define the OT program with respect to a cost matrix Č as

    L(Č) := P̌∈_cP̌Č .

The following lemma describes the conditions on Č for P̌ to remain unperturbed at P̌^⋆.


 
Define the set of vertices on the c-th Birkhoff polytope _c that are within a δ-ball from P̌^⋆ as 

    _δ (P̌^⋆) =
        {P̌∈_c ∖{P̌^⋆} :
           P̌ - P̌^⋆_F ≤δ} .

Define the set of matrices that denote directions from P̌^⋆ to each neighboring vertex P̌_i as

    _δ (P̌^⋆) =
        {V̌∈^c × c :
           V̌ = P̌ - P̌^⋆ , P̌∈_δ(P̌^⋆) } .

For the linear program's solution P̌^⋆ = L(Č) to remain unchanged, ∀V̌∈_δ (P̌^⋆) for δ = 2, the cost matrix Č should satisfy

    ČV̌ > 0 .



Birkhoff-von Neuman's theorem <cit.> states that the optimal transport solution must lie on the convex hull of the c-th Birkhoff polytope _c, and that its vertices are in fact permutation matrices.
We therefore say that an LP solution P̌^⋆ is a vertex on _c.
The outline of this proof is straightforward: so that P̌^⋆ remains unchanged, Č should not cause P̌^⋆ to move to an adjacent edge of the Birkhoff polytope, nor should it cause it to extend beyond its adjacent edge because then, the neighboring vertex would assume the new solution.
For the rest of this proof, we shall let P̌^⋆ = Ǐ without any loss of generality.
We define the set of nearest neighbors to P̌^⋆, which are simply permutation matrices that can be described as taking P̌^⋆ and exchanging any two columns.
For a P̌ with any two columns of P̌^⋆ exchanged, notice that the difference matrix V̌ = P̌ - P̌^⋆ is a symmetric matrix of mostly zeros except for two off-diagonal +1 entries and two diagonal -1 entries, hence P̌-P̌^⋆_F = 2.
Formally, we describe the set of nearest neighbors with _δ (P̌^⋆) with δ = 2, defined by (<ref>).
Next, we define edges that are adjacent to P̌^⋆ using the set _δ (P̌^⋆) also with δ = 2, defined by (<ref>).
Note that there are K := c  2 - 1 = c(c-1)/2 - 1 neighboring vertices and adjacent edges.

We will now show how perturbing Č in just one direction V̌ changes P̌^⋆.
First, consider a cost matrix that produce P̌^⋆, which is defined as Č^⋆ := ∑_V̌_i ∈_δ (P̌^⋆)V̌_i, meaning that it is equi-angle from all V̌∈_δ (P̌^⋆), for δ = 2.
By enumerating over all V̌, we may derive Č^⋆ = ^⊤ - c Ǐ, along with the fact that Č^⋆V̌/Č^⋆_F V̌_F = 1/√(c-1).
Define P̌_1 = L(Č_1) with Č_1 such that V̌_1Č_1≤ 0 and V̌_iČ_1/V̌_i_F Č_1_F = 1/√(c-1), ∀ i ∈ 2, …, K.
Define P̌_2 = L(Č_2) with Č_2 such that V̌_iČ_2/V̌_i_F Č_2_F = 1/√(c-1), ∀ i ∈ K.
We make the claim that P̌_1 = P̌_2 and proceed with a proof by contradiction.
As mentioned before, ∀V̌∈_δ (P̌^⋆ ) for P̌^⋆ = Ǐ, δ = 2 has exactly four non-zero entries, i.e., V_ii = V_jj = -1 and V_ij = V_ji = +1, where i,j ∈ c , i ≠ j.
Writing out V̌Č explicitly, we have V̌Č = -(C_ii + C_jj) + (C_ij + C_ji).
To ensure that Č_1 and Č_2 does not interact with other edges of the polytope, we fix V̌_iČ = V̌_iČ^⋆ = c for i=2,…,K.
Since we constructed Č_1 and Č_2 to differ only by the condition V̌_1Č_1≤ 0 or V̌_1Č_2 > 0, and any V̌_i affects only four entries of the cost matrix, we may greatly simplify our analysis of Č and P̌ to only these four entries.
As such, we extract these four entries of P̌, represent it using P̌̂̌∈_2, and parameterize using t ∈ [0,1] it as

    P̌̂̌(t) =
        t [ 1 0; 0 1 ] +
        (1-t) [ 0 1; 1 0 ] .

With this parameterized form, we may reexpress the optimization for L(Č) as

    min_t ∈ [0,1]
        t(C_ii + C_jj) + (1-t) (C_ij + C_ji) 
        =
        min_t ∈ [0,1](C_ii + C_jj/C_ij + C_ji - 1 ) t

The above minimization has three cases.
If C_ii + C_jj = C_ij + C_ji then there exists no unique solution t.
If C_ii + C_jj > C_ij + C_ji then t = 0.
If C_ii + C_jj < C_ij + C_ji then t = 1.
The first two cases directly corresponds to V̌_1Č_1≤ 0, while the third case corresponds to V̌_1Č_2 > 0.
The fact that t is not consistent between all cases demonstrates a contradiction.
Moreover, t=1 produces the solution P̌̂̌ = Ǐ, and if this holds for all V̌∈_δ (P̌^⋆), then all off-diagonal entries must be zero and therefore P̌^⋆ = Ǐ must be the minimizer for L(Č).



A direct consequence of lemma <ref> is the following.




The solution to the linear program defined by (<ref>) is P̌^⋆ = Ǐ if the linear cost matrix Č satisfies the following property

    C_ij + C_ji - C_ii - C_jj > 0, ∀ i,j ∈ c , i ≠ j .



Analyzing (<ref>), we observe that any V̌∈_δ (P̌^⋆) for δ = 2 has only four symmetric non-negative entries, which we condense V̌ and Č into ^2 × 2 matrices at these four support locations respectively as
V̌̂̌ = [ -1 +1; +1 -1 ] and Ĉ̌ = [ C_ii C_ji; C_ij C_jj ].
It thus follows that an explicit computation produces:

    P̌Č≥ 0
    ⇒P̌̂̌Ĉ̌ = (C_ij + C_ji) - (C_ii + C_jj) ≥ 0 .

Since the set _δ (P̌^⋆) spans all permutations between i,j ∈ c , i ≠ j, we conclude with (<ref>).



In the variational setting, C_ij(Ř,Q̌_ij)'s are themselves linearly coupled with Ř and Q̌_ij.
The following proposition introduces a trivial criterion on C_ij(Ř,Q̌_ij)'s to guarantee that P̌ remains unperturbed from P̌^⋆.
For pedagogical reasons, we shall assume that n →∞ for this proposition but subsequently relax this.




Problem (<ref>) yields the solution P̌^⋆ if, ∀ i,j : i ≠ j, the following criterion is satisfied:

    min_Ř,Q̌_ij C_ij(Ř,Q̌_ij) 
        + min_Ř,Q̌_ji C_ji(Ř,Q̌_ji) 
        - min_Ř,Q̌_ii C_ii(Ř,Q̌_ii) 
        - min_Ř,Q̌_jj C_jj(Ř,Q̌_jj)  > 0 .



Consider a set of c clusters where clusters i,j satisfy

    C_ij(Ř'_ij,Q̌'_ij) 
        + C_ji(Ř'_ji,Q̌'_ji) 
        - C_ii(Ř',Q̌'_ii) 
        - C_jj(Ř',Q̌'_jj) ≤ 0 ,

where

    (Ř'_ij,Q̌'_ij)
           := Ř∈ S(d,d), Q̌_ij∈_n C_ij(Ř) , 
    
        (Ř'_ji,Q̌'_ji) 
           := Ř∈ S(d,d), Q̌_ji∈_n C_ji(Ř) , 
    
        (Ř',Q̌'_ii,Q̌'_jj)
           := Ř∈ S(d,d), Q̌_ii,Q̌_jj∈_n C_ii(Ř,Q̌_ii) + C_jj(Ř,Q̌_jj) .

Since the 2-Wasserstein is a valid metric, its distance between any two clusters must satisfy C_ij(Ř) ≥ 0 with equality holding if and only if the clusters are exactly similar.
If P̌^⋆ = Ǐ, similar clusters are denoted with matching indices, and it must follow that C_ii(Ř') + C_jj(Ř') = 0.
This implies that (<ref>) must be false since C_ij(Ř'_ij), C_ji(Ř'_ji) > 0 for mismatched clusters.
Due to this contradiction, the disambiguity criterion (<ref>) must hold for all cluster pairs i,j : i ≠ j.



This proposition provides a disambiguity criterion, requiring that matched clusters (i.e., C_ii,C_jj) should be more similar than mismatched clusters (i.e., C_ij,C_ji) up to some disambiguity threshold (in the case of n →∞, the threshold is 0).
To extend this proposition to the finite-sample regime, we require a higher disambiguity-threshold to offset uncertainty due to finite samples.







 §.§ Part 2: Disambiguity criterion in the finite-sample regime


We utilize a recent p-Wasserstein concentration bound by Weed and Bach <cit.> that describes finite sample behavior on the Wasserstein distance for data embedded in high-dimensional space, but whose clusters are themselves approximately low-dimensional.
We will proceed our analysis with the language of probability measures μ to make our analysis consistent with <cit.>. We thus define the equivalent measure analogs as follows.

 
Let clusters X̌_i ∈^D × n_x,i and Y̌_j ∈^D × n_y,j be respectively denoted by empirical measures as

    μ̂_x,i := 1/n_x,i∑_k = 1^n_x,iδ_X̌_i(k) ,
            μ̂_y,i := 1/n_y,j∑_k = 1^n_y,jδ_Y̌_j(k) ,

where δ_X̌_i(k) refers to a discrete point located at X̌_i(k).
At the limit, we denote the measures as

    μ_x,i := lim_n_x,i→∞1/n_x,i∑_k = 1^n_x,iδ_X̌_i(k) , 
            μ_y,j := lim_n_y,j→∞1/n_y,j∑_k = 1^n_y,jδ_Y̌_j(k) .





Denote a linear transformation Ř applied on the measure as

    Ř∘μ̂_x,i
        := 1/n_x,i∑_k = 1^n_x,iδ_ŘX̌_i(k) ,
            Ř∘μ_x,i
        := lim_n_x,i→∞1/n_x,i∑_k = 1^n_x,iδ_ŘX̌_i(k) .

The transformed inter-cluster distance between clusters may thus be denoted as

    Ĉ_ij (Ř) := _2^2( Ř∘μ̂_x,i , μ̂_y,i ) ,
            C̃_ij (Ř) := _2^2( Ř∘μ_x,i , μ_y,i ) .




Now, we may proceed to state results from <cit.>.
The following result pertains to the sample complexity of measures μ in ^D that are supported on an approximately low-dimensional set in ^d, where d ≪ D.
First we require some definitions.



Given a set S ⊆^d, let _ε(S) denote the ε-covering number of set S, which is defined as the minimum number m of closed balls B_1,…,B_m of diameter ε such that S ⊆⋃_1 ≤ i ≤ m B_i.



For any set S ⊆^d, the ε-fattening of S is S_ε := { y : D(y,S) ≤ε}.




Let S be a set that satisfies _ε'(S) ≤ (3 ε')^-d for all ε' ≤ 1/27 and for some d > 2p.
Suppose there exists a positive constant σ such that μ satisfies μ(S_ε) ≥ 1 - e^-ε^2/2σ^2 for all ε > 0.
If p log1/σ≥ 1/18, then for all n ≤ (18pσ^2 log1/σ)^-d/2,

    [_p^p (μ,μ̂_n)] ≤ c n^-p/d ,

where c = 27^p (2 + 1/3^d/2 - p - 1).



This proposition states that the degree that μ is concentrated (as parameterized by σ) around set S (approximately supported in low-dimensions) affects how “long” (in terms of n) we can enjoy the fast convergence rate of n^-p/d.
We will leverage this result to obtain the following theorem on cluster correspondence disambiguity with respect to sample complexity.




Let sets S_x,i and S_y,j satisfy the conditions for S in proposition <ref> for some σ_x,i,σ_y,j > 0 and d_x,i,d_y,j > 4.
If log1/σ_x,i≥1/36 and log1/σ_y,j≥1/36, then for all n_x,i≤ (36σ_x,i^2 log1/σ_x,i)^-d_x,i/2 and n_y,j≤ (36σ_y,j^2 log1/σ_y,j)^-d_y,j/2,

    [| C̃_ij(Ř) - Ĉ_ij(Ř) |] ≤ 
        c_x,i n_x,i^-2/d_x,i + c_y,j n_y,j^-2/d_y,j ,

where

    c_z,k = 729 (2 + 1/3^d_z,k/2 - 2 - 1) .



Denote μ_i,μ_j as measures and μ̂_i,μ̂_j as their empirical estimates.
By the triangle inequality,

    [_2^2(μ̂_i,μ̂_j)]
       ≤[_2^2(μ̂_i,μ_i) + _2^2(μ_j,μ̂_j)] 
       ≤[_2^2(μ̂_i,μ_i) + _2^2(μ_i,μ_j) + _2^2(μ̂_j,μ_j)]
    ⇒[|_2^2(μ_i,μ_j) - _2^2(μ̂_i,μ̂_j)|]
       ≤[_2^2(μ̂_i,μ_i)] + [_2^2(μ̂_j,μ_j)] 
    ⇒ [| C̃_ij(Ř) - Ĉ_ij(Ř) |]
       ≤ c_x,i n_x_i^-2/d_x,i + c_y,j n_y,j^-2/d_y,j ,

where the last line is a direct application of definition <ref> and proposition <ref>.





Let S_x,i,S_y,i,S_x,j,S_y,j be sets that satisfy the conditions for S in proposition <ref> for some σ_x,i,σ_y,i,σ_x,j,σ_y,j > 0 and d_x,i,d_y,i,d_x,j,d_y,j > 4.
If log1/σ_x,i , log1/σ_y,i , log1/σ_x,j , log1/σ_y,j≥1/36, then for all n_x,i≤ (36σ_x,i^2 log1/σ_x,i)^-d_x,i/2,
n_y,i≤ (36σ_y,i^2 log1/σ_y,i)^-d_y,i/2,
n_x,j≤ (36σ_x,j^2 log1/σ_x,j)^-d_x,j/2, and
n_y,j≤ (36σ_y,j^2 log1/σ_y,j)^-d_y,j/2,
cluster correspondences in problem (<ref>) may be disambiguated to achieve P̌^⋆ = Ǐ when the following criterion is fulfilled for all i,j ∈ c  : i ≠ j:

    [   min_ŘĈ_ij(Ř) 
            + min_ŘĈ_ji(Ř)
            - min_ŘĈ_ii(Ř)
            - min_ŘĈ_jj(Ř) ] 
        > 2 (B_x,i + B_y,i + B_x,j + B_y,j ) .

where the constants are defined as

    B_z,k := c_z,k n_z,k^-2/d_z,k ,
          
        c_z,k := 729 ( 2 + 1/3^d_z,k/2 - 2 - 1) .



This follows a direct application of criterion (<ref>) and corollary <ref>.





For all n ≥ 0 and 0 ≤ p < ∞,

    [ _p^p(μ,μ̂_n) ≥_p^p(μ,μ̂_n) + t ]
        ≤exp(-2nt^2) .






Define Ĉ^⋆_ij := min_Ř∈ S(D,D)Ĉ_ij (Ř).
If the conditions in Lemma <ref> are satisfied then problem (<ref>) yields the solution P̌^⋆ = Ǐ with probability at least 1-δ if, ∀ i,j : i ≠ j, the following criterion is satisfied:

    Ĉ_ij^⋆
        + Ĉ_ji^⋆
        - Ĉ_ii^⋆
        - Ĉ_jj^⋆  >  
        2 (   B_x,i(δ) + B_y,i(δ)
            + B_x,j(δ) + B_y,j(δ) )

where

    B_z,k(δ) :=   c_z,k n_z,k^ -2/d_z,k
                          + √(log (1/δ) / 2 n_z,k)  ,
          
        c_z,k := 729 ( 2 + 1/3^d_z,k/2 - 2 - 1) ,

where d_z,k refers to the intrinsic dimension of the k-th cluster from the z-th dataset.


For some measure μ and its empirical finite-sample estimate μ̂_n, proposition <ref> may be equivalently stated with the choice of t = √(log1/δ / 2n ) as:

    | _p^p(μ,μ̂_n) - _p^p(μ,μ̂_n) |
        ≤√(log (1/δ) / 2n ) ,

holds with at least probability 1 - δ.
Under the conditions stated in proposition (<ref>), and combining its result with the above relation, we have

    _p^p (μ,μ̂_n) ≤ c n^-p/d + √(log (1/δ) / 2n ) ,

where c = 27^p (2 + 1/3^d/2 - p - 1).
Combining this for all terms in the left-hand side of (<ref>) yields the stated result.






 §.§ Putting everything together


The final proof of Theorem <ref> is a simplified version of Theorem <ref>'s since we assert a stronger (but cleaner) exact low-rank assumption to streamline communication.
When the data is exactly supported in low-dimensions (as opposed to approximately), the ε-fattening disappears (i.e., ε→ 0) thus any positive σ < ε will send n →∞, implying that the rapid convergence in dimensions d ≪ D holds for n →∞.
Hence an identical result holds, with the sole condition that d > 4. □






§ PROOF OF THEOREM <REF>


We apply a very recent perturbation bound for the Procrustes problem developed by Arias-Castro et al. <cit.> to subsequently state a cluster-based alignment bound.
First, we outline the perturbation bound for the classical Procrustes problem below.




Consider short matrices X̌,Y̌∈^d × n with d < n and X̌ having full rank.
Set ε^2 = Y̌^⊤Y̌ - X̌^⊤X̌_p, where ·_p denotes the Schatten p-norm.
Denote the singular value decomposition of X̌ = ǓΣ̌V̌^⊤, where Σ̌ contains diagonal elements σ_1 ≥σ_2 ≥…σ_d > 0 = … = 0, and let X̌^† be the pseudo-inverse of X̌, i.e., X̌^† = ǓΣ̌^†V̌^⊤, where Σ̌^† = (σ_1^-1,σ_2^-1,…,σ_d^-1,0,…,0).
If X̌^†ε≤1/√(2) (X̌X̌^†)^-1/2 then

    min_Ř∈ S(d,d)ŘX̌ - Y̌_p
        ≤
        (X̌X̌^†+2) X̌^†ε^2.




Directly applying Theorem <ref> using the Schatten 2-norm (i.e., the Frobenius norm) yields

    min_Ř∈ S(d,d)ŘX̌ - Y̌_F^2
           =  min_Ř∈ S(d,d)∑_i=1^c
             (   Q̌_ii^⊤X̌_i^⊤X̌_i Q̌_ii
                  + Y̌_i^⊤Y̌_i
                  - 2 Y̌_iQ̌_ii^⊤X̌_i^⊤Ř^⊤ ) 
       ≤ 
           ( ( X̌X̌^† + 2 )
                  X̌^†ε^2 )^2 := B^2 ,


    ⇒max_Ř∈ S(d,d)∑_i=1^c
         ( 2Y̌_iQ̌_ii^⊤X̌_i^⊤Ř^⊤ )
        ≥ ( X̌X̌^⊤ + Y̌Y̌^⊤ )
        - B^2 ,

where ε^2 := Y̌^⊤Y̌ - X̌^⊤X̌_F
and under the conditions that X̌ is full rank and X̌^†ε≤1/√(2) (X̌X̌)^†)^-1/2.
When the criterion given by corollary <ref> is satisfied, we are guaranteed cluster-correspondences P̌^⋆ = Ǐ.
Therefore ∑_ij P_ij C_ij(Ř) = ∑_i C_ii(Ř).
We utilize this to lower bound ∑_i C_ii(Ř) as follows:

    min_Ř∈ S(d,d)∑_i=1^c C_ii(Ř) 
    
        =   min_Ř∈ S(d,d)∑_i=1^c
             (   X̌_i (Q̌_ii) X̌_i^⊤
                  + Y̌_i (Q̌_ii^⊤) Y̌_j^⊤
                  - 2 Y̌_i Q̌_ii^⊤X̌_i^⊤Ř^⊤ ) 
    
        =   ∑_i=1^c
             (   X̌_i (Q̌_ii) X̌_i^⊤
                  + Y̌_i (Q̌_ii^⊤) Y̌_j^⊤ )
            - max_Ř∈ S(d,d)∑_i=1^c
               ( 2 Y̌_j Q̌_ii^⊤X̌_i^⊤Ř^⊤ ) 
    ≤   
            B^2 +
            ∑_i=1^c
             (   X̌_i ( (Q̌_ii) - Q̌_iiQ̌_ii^⊤ ) X̌_i^⊤ 
                  + Y̌_j ( (Q̌_ii) - Ǐ ) Y̌_j^⊤  ) 
    
        =   
            B^2 +
            ∑_i=1^c
             (   X̌_i ( 1/n - Q̌_iiQ̌_ii^⊤ ) X̌_i^⊤ 
                  + Y̌_j ( 1/n - Ǐ ) Y̌_j^⊤  ) .

□






§ PROOF OF LEMMA <REF>



To ease notation, let Ǎ = Y̌_j Q̌_ij^⊤X̌_i^⊤, and Č = 1/nX̌_i X̌_i^⊤ + 1/nY̌_j Y̌_j^⊤.
Let Ǎ be decomposed by the singular-value decomposition as Ǎ = Ǔ̃Σ̃̌̃Ṽ̌̃^⊤.
Lastly, let Ǔ = [Ǔ',Ǔ”] and V̌ = [V̌',V̌”], where ⟨Ǔ”,Ǔ' ⟩ = ⟨V̌”,V̌' ⟩ = 0.
Then it follows that

    ( Ǔ^⊤Ǔ̃Σ̃̌̃Ṽ̌̃^⊤V̌ )
           ≤( Σ̃̌̃ ) 
        = ( Ǔ̃^⊤Ǔ̃Σ̃̌̃Ṽ̌̃^⊤Ṽ̌̃ )
        
    ( Ǔ^⊤ǍV̌ )
           ≤( Ǔ̃^⊤ǍṼ̌̃ )
        
    ( Č - 2 Ǔ'^⊤ǍV̌' - 2 Ǔ”^⊤ǍV̌”)
           ≥( Č - 2 Ǔ̃^⊤ǍṼ̌̃)
        
    min_Ǔ”, V̌”∈ S(d,d-r) : 
    ⟨Ǔ',Ǔ”⟩ = ⟨V̌',V̌”⟩ = 0( Č - 2 ǍV̌' Ǔ'^⊤ - 2 ǍV̌”Ǔ”^⊤)
           ≥min_Ǔ̃, Ṽ̌̃∈ S(d,d)( Č - 2 ǍṼ̌̃Ǔ̃^⊤)
        
    min_Ř∈(Ǔ',V̌') C_ij (Ř)
           ≥min_Ř∈ S(d,d) C_ij (Ř) .


What remains is for us to show the condition for equality.
From (<ref>), we have that

    ( Ǔ^⊤Ǔ̃Σ̃̌̃Ṽ̌̃^⊤V̌ )
        =
        ( Ṽ̌̃^⊤V̌Ǔ^⊤Ǔ̃Σ̃̌̃ )
        ≤( Σ̃̌̃ ) ,

with equality holding if Ṽ̌̃^⊤V̌Ǔ^⊤Ǔ̃ = Ǐ, implying that Ǔ̃^⊤Ǔ = Ṽ̌̃^⊤V̌, which imply that

    ⟨Ǔ̃ , Ǔ' ⟩
        =
        ⟨Ṽ̌̃ , V̌' ⟩,
            ⟨Ǔ̃ , Ǔ”⟩
        =
        ⟨Ṽ̌̃ , V̌”⟩ ,

which are obtained via the substitutions Ǔ = [Ǔ',Ǔ”] and V̌ = [V̌',V̌”]. □








§ ADDITIONAL FIGURES


In this section, we provide an additional figure to compare HiWA with other competitor methods on the brain decoding example.





