Year: 2021

Abstract:
Unsupervised dataset alignment estimates a transformation that maps two or more source domains to a shared aligned domain given only the domain datasets. This task has many applications including generative modeling, unsupervised domain adaptation, and socially aware learning. Most prior works use adversarial learning (i.e., min-max optimization), which can be challenging to optimize and evaluate. A few recent works explore non-adversarial flow-based (i.e., invertible) approaches, but they lack a unified perspective. Therefore, we propose to unify and generalize previous flow-based approaches under a single non-adversarial framework, which we prove is equivalent to minimizing an upper bound on the JensenShannon Divergence (JSD). Importantly, our problem reduces to a min-min, i.e., cooperative, problem and can provide a natural evaluation metric for unsupervised dataset alignment. We present empirical results of our framework on both simulated and real-world datasets to demonstrate the benefits of our approach.

heading: 1 INTRODUCTION
text: In many cases, a practitioner has access to multiple related but distinct datasets such as agricultural measurements from two farms, experimental data collected in different months, or sales data before and after a major event. Unsupervised dataset alignment (UDA) is the ML task aimed at aligning these related but distinct datasets in a shared space, which may be a latent space, without any pairing information between the two domains (i.e., unsupervised). This task has many applications such as generative modeling (e.g., (Zhu et al., 2017)), unsupervised domain adaptation (e.g., (Grover et al., 2020; Hu et al., 2018)), batch effect mitigation in biology (e.g., (Haghverdi et al., 2018)), and fairness-aware learning (e.g., (Zemel et al., 2013)).
The most common approach for obtaining such alignment transformations stems from Generative Adversarial Networks (GAN)(Goodfellow et al., 2014), which can be viewed as minimizing a lower bound on the Jensen-Shannon Divergence (JSD) between real and generated distributions. The lower bound is tight if and only if the inner maximization is solved perfectly. CycleGAN (Zhu et al., 2017) maps between two datasets via two GAN objectives between the two domains and a cycle consistency loss, which encourages approximate invertibility of the transformations. However, adversarial learning can be challenging to optimize in practice (see e.g. (Lucic et al., 2018; Kurach et al., 2019)) in part because of the competitive nature of the min-max optimization problem. Perhaps more importantly, the research community only has reasonable GAN evaluation metrics for certain data types. Specifically, the commonly accepted Frechet Inception Distance (FID) (Heusel et al., 2017) is only applicable to image or auditory data, which have standard powerful pretrained classifiers, and even the implementation of FID can have evaluation issues (Parmar et al., 2021). No clear metrics exist for tabular data or non-perceptual data.
Recently, flow-based methods that leverage invertible models have been proposed for the UDA task (Grover et al., 2020; Usman et al., 2020). AlignFlow (Grover et al., 2020) leverages invertible models to make the model cycle-consistent (i.e., invertible) by construction and introduce exact log-likelihood loss terms derived from standard flow-based generative models as a complementary loss terms to the adversarial loss terms. Yet, AlignFlow still leverages adversarial learning. Loglikelihood ratio minimizing flows (LRMF) (Usman et al., 2020) use invertible flow models and
density estimation for distribution alignment without adversarial learning and define a new metric based on the log-likelihood ratio. However, LRMF depends heavily on the density model class and can only partially align datasets if the target distribution is not in the chosen density model class. Additionally, the LRMF metric depends on this density model class and is only defined for two datasets.
Therefore, to provide an alternative to adversarial learning and generalize previous flow-based approaches, we propose a unified non-adversarial UDA framework, which we prove is equivalent to minimizing an upper bound on the JSD. Importantly, our problem reduces to a min-min, i.e., cooperative, problem, and the JSD upper bound can provide a natural evaluation metric for UDA that can be applied in any domain. Our framework requires two parts, the outer minimization requires an invertible model and the inner minimization requires a density model (e.g., Gaussian mixture models or normalizing flows (Dinh et al., 2017)). We summarize our contributions as follows:
• We prove that a minimization problem over density models is an upper bound on a generalized version of JSD that allows for more than two distributions. Importantly, we also theoretically quantify the bound gap and show that it can be made tight if the density model class is flexible enough.
• We use this JSD upper bound to derive a novel regularized loss function for UDA and explain its relationship to prior methods.
• We empirically demonstrate the benefits of our method compared to prior flow-based models on both simulated and real-world datasets.
Notation We will denote distributions as PX(x) where X is the corresponding random variable. Invertible functions will be denoted by T (·). We will use Xj ∼ PXj to denote the observed random variable from the j-th distribution. We will use Zj ≜ Tj(Xj) ∼ PZj ≡ PTj(Xj) to denote the latent random variable of the j-th distribution after applying Tj to Xj (and note that Xj = T−1j (Zj)). We will denote the mixtures of these observed or latent distributions as PXmix ≜ ∑ jwjPXj and
PZmix ≜ ∑
jwjPZj , where w is a probability vector. We denote KL divergence, entropy, and cross entropy as KL(·, ·), H(·), and Hc(·, ·), respectively, where KL(P,Q) = Hc(P,Q)−H(P ).

heading: 2 REGULARIZED ALIGNMENT UPPER BOUND LOSS
text: In this section, we will introduce our main theoretical result proving an upper bound on the generalized JSD divergence, deriving our loss function based on this, and then showing that minimizing this upper bound results in aligned distributions assuming large enough capacity of the model components.
Background: Normalizing Flows and Invertible Models Normalizing flows are generative models that have tractable distributions where exact density evaluations and efficient samplings are ensured (Kobyzev et al., 2021). Such models leverage the change of variables formula to create an invertible mapping T such that PX(x) = PZ(T (x))|JT (x)| where PZ is a known latent distribution and |JT (x)| is the absolute value of determinant of the Jacobian of the invertible map T . For sampling in distribution PX , one need to first sample from the latent distribution PZ and then apply the inverse transform T−1. Therefore the key challenge in designing invertible models is to have computationally efficient inverse evaluation and Jacobian determinant calculation. Many approaches have been proposed by parameterizing mapping function T as deep neural networks including autoregressive structures (Kingma et al., 2016; Papamakarios et al., 2021), coupling layers (Dinh et al., 2017)(Kingma and Dhariwal, 2018), oridinary differential equations(Grathwohl et al., 2018), and invertible residual networks(Chen et al., 2019; Behrmann et al., 2019). Flow models can be then learned efficiently by maximizing the likelihood for the given data.
Background: Generalized JSD We remind the reader of the generalized Jensen-Shannon divergence for more than two distributions, where the standard JSD is recovered if w1 = w2 = 0.5.
Definition 1 (Generalized Jensen-Shannon Divergence (GJSD) (Lin, 1991)). Given k distributions {PXj}kj=1 and a corresponding probability weight vector w, the generalized Jensen-Shannon divergence is defined as (proof of equivalence in appendix):
GJSDw(PX1 , · · · , PXk) ≜ ∑ jwj KL(PXj , ∑ jwjPXj )
≡ H (∑ jwjPXj ) − ∑ jwj H(PXj ) . (1)

heading: 2.1 GJSD UPPER BOUND
text: The goal of distribution alignment is to find a set of transformations {Tj(·)}kj=1 (which will be invertible in our case) such that the latent distributions align, i.e., PTj(Xj) = PTj′ (Xj′ ) or equivalently PZj = PZj′ for all j ̸= j
′. Given the properties of divergences, this alignment will happen if and only if GJSD(PZ1 , · · · , PZk) = 0. Thus, ideally, we would minimize GJSD directly with respect to Tj , i.e.,
min T1,··· ,Tk∈T GJSD(PT1(X1), · · · , PTk(Xk)) ≡ min T1,··· ,Tk∈T
H (∑ jwjPTj(Xj) ) − ∑ jwj H(PTj(Xj)) ,
where T is a class of invertible functions. However, we cannot evaluate the entropy terms in Eqn. 2 because we do not know the density of PXj ; we only have samples from PXj . Therefore, we will upper bound the first entropy term in Eqn. 2 (H (∑ jwjPXj ) ) using an auxiliary density model and decompose the other entropy terms by leveraging the change of variables formula for invertible functions.
Theorem 1 (GJSD Upper Bound). Given an auxiliary density model class Q, we form a GJSD upper bound:
GJSDw(PZ1 , · · · , PZk) ≤ min Q∈Q
Hc(PZmix , Q)− ∑ jwj H(PZj ) ,
where the bound gap is exactly minQ∈Q KL(PZmix , Q).
Proof of Theorem 1. For any Q ∈ Q, we have the following upper bound:
GJSDw(PZ1 , · · · , PZk) = Hc(PZmix , Q)−Hc(PZmix , Q)︸ ︷︷ ︸ =0
+H(PZmix)− ∑ jwj H(PZj )
= Hc(PZmix , Q)−KL(PZmix , Q)− ∑ jwj H(PZj )
≤ Hc(PZmix , Q)− ∑ jwj H(PZj ) ,
where the inequality is by the fact that KL divergence is non-negative and the bound gap is equal to KL(PZmix , Q). The Q that achieves the minimum in the upper bound is equivalent to the Q that
minimizes the bound gap, i.e., Q∗ = argmin
Q∈Q Hc(PZmix , Q)− ∑ jwj H(PZj )︸ ︷︷ ︸
Constant w.r.t. Q
(2)
= argmin Q∈Q Hc(PZmix , Q) −H(PZmix)︸ ︷︷ ︸ Constant w.r.t. Q
(3)
= argmin Q∈Q KL(PZmix , Q) . (4)
The tightness of the bound depends on how well the class of density modelsQ (e.g., mixture models, normalizing flows, or autoregressive densities) can approximate PZmix ; notably, the bound can be made tight if PZmix ∈ Q. Also, one key feature of this upper bound is that the cross entropy term can be evaluated using only samples from PXj and the transformations Tj , i.e., Hc(PZmix , Q) =∑
jwjEPXj [− logQ(Tj(xj))]. However, we still cannot evaluate the other entropy terms H(PZj ) since we do not know the densities of PZj (or PXj ). Thus, we leverage the fact that the Tj functions are invertible to define an entropy change of variables. Lemma 2 (Entropy Change of Variables). Let X ∼ PX and Z ≜ T (X) ∼ PZ , where T is an invertible transformation. The entropy of Z can be decomposed as follows:
H(PZ) = H(PX) + EPX [log |JT (x)|] , (5) where |JT (x)| is the determinant of the Jacobian of T .
The key insight from this lemma is that H(PX) is a constant with respect to T and can thus be ignored when optimizing T , while EPX [log |JT (x)|] can be approximated using only samples from PX .

heading: 2.2 ALIGNMENT UPPER BOUND (AUB)
text: Combining Theorem 1 and Lemma 2, we can arrive at our final objective function which is equivalent to minimizing an upper bound on the GJSD:
GJSDw(PZ1 , · · · , PZk) ≤ min Q∈Q
Hc(PZmix , Q)− ∑ jwj H(PZj ) (6)
= min Q∈Q
∑ jwjEPXj [− logQ(Tj(x))|JTj (x)|]− ∑ jwj H(PXj ) , (7)
where the last term− ∑
jwj H(PXj ) is constant with respect to Tj functions so they can be ignored. We formally define this loss function as follows. Definition 2 (Alignment Upper Bound Loss). Given k continuous distributions {PXj}kj=1, a class of continuous distributionsQ, and a probability weight vector w, the alignment upper bound loss is defined as follows:
LAUB(T1, · · · , Tk; {PXj}kj=1,Q,w) ≜ min Q∈Q ∑ jwjEPXj [− log |JTj (x)|Q(Tj(x))] , (8)
where Tj are invertible and |JTj (x)| is the absolute value of the Jacobian determinant.
Notice that this alignment loss can be seen as learning the best base distribution given fixed flow models Tj . We now consider the theoretical optimum if we optimize over all invertible functions. Theorem 3 (Alignment at Global Minimum of LAUB). If LAUB is minimized over the class of all invertible functions, a global minimum of LAUB implies that the latent distributions are aligned, i.e., PTj(Xj) = PTj′ (Xj′ ) for all j ̸= j ′. Notably, this result holds regardless of Q.
Informally, this can be proved by showing that the problem decouples into separate normalizing flow losses where Q is the base distribution and the optimum is achieved only if PTj(Xj) = Q for all Tj (formal proof in the appendix). This alignment of the latent distributions also implies the translation between any of the observed component distributions. The proof follows directly from Theorem 3 and the change of variables formula. Corollary 4 (Translation at Global Minimum of LAUB). Similar to Theorem 3, a global minimum of LAUB implies translation between any component distributions using the inverses of Tj , i.e., PT−1
j′ (Tj(Xj)) = PXj′ for all j ̸= j ′.
Regularization via Transportation Cost While the alignment objective is the most challenging part of UDA, we suggest that regularization may also be useful for practical and stable alignment (or translation) between datasets because there are many optimal alignment solutions—even infinitely many in most cases (see appendix for two examples). We alleviate this issue by adding expected transportation cost (usually squared Euclidean distance) as a regularization to our objective inspired by optimal transport (OT) concepts.
Definition 3 (Regularized Alignment Upper Bound Loss, RAUB). Given similar setup as in Def. 2 and a transportation cost function c(a, b) ≥ 0 for transporting a point from a to b, the regularized alignment upper bound loss is defined as:
LRAUB(T1, · · · , Tk; {PXj}kj=1,Q,w, λ, c) ≜ min
Q∈Q
∑ jwjEPXj [− log |JTj (x)|Q(Tj(x)) + λc(x, Tj(x))] .
(9)
Pseudo-code for training our model with RAUB objective can be found in the appendix Section C.

heading: 3 RELATIONSHIP TO PRIOR WORKS
text: AlignFlow without adversarial terms is a special case AlignFlow (Grover et al., 2020) without adversarial loss terms is a special case of our method for two distributions where the density model class Q only contains the standard normal distribution (i.e., a singleton class) and no regularization is used (i.e., λ = 0). Thus, AlignFlow can be viewed as initially optimizing a poor upper bound on JSD; however, the JSD bound becomes tighter as training progresses because the latent distributions independently move towards the same normal distribution.
LRMF is special case with only one transformation Log-likelihood ratio minimizing flows (LRMF) (Usman et al., 2020) is also a special case of our method for only two distributions, where one transformation is fixed at the identity (i.e., T2 = Id) and no regularization is applied (i.e., λ = 0). While the final practical LRMF objective is a special case of ours, the theory is developed from a different but complementary perspective. The LRMF metric developed requires an assumption about a given density model class, which enables a zero point (or absolute value) of the metric to be estimated but requires fitting extra domain density models. Usman et al. (2020) also do not uncover the connection of the objective as an upper bound on JSD regardless of the density model class. Additionally, to ensure alignment, LRMF requires that the density model class includes the true target distribution because only one invertible transform is used, while our approach can theoretically align even if the shared density model class is weak (see Theorem 3 and our simulated experiments). See Fig. 2 for a comparison of our approach to prior works.
Cooperative networks versus adversarial networks Analogous to the generator G and the discriminator D in adversarial learning, our framework has two main networks, Tj and Qz . We can use any invertible function for Tj (e.g., coupling-based flows (Dinh et al., 2017), neural ODE flows (Grathwohl et al., 2018), or residual flows (Chen et al., 2019)) and any (approximate) density models for Qz (e.g., kernel densities (in low dimensions), mixture models, autoregressive densities (Salimans et al., 2017), normalizing flows (Kingma and Dhariwal, 2018), or even VAEs (Kingma and Welling, 2019)). Thus, our framework has similar modularity compared to adversarial approaches. In contrast, we have a min-min, i.e., cooperative, optimization problem, but our transformations must be invertible. As another difference, the inner optimization problem may be more challenging (i.e., fitting a density model Qz may be more difficult than fitting an auxiliary discriminator D), but the overall min-min is likely to be more stable than the min-max problem—specifically, note that our problem can align even if the density model Qz is not optimal (see Theorem 3). We expect some of these limitations to be alleviated as new invertible models and density models are continually being developed. Therefore, the proposed approach provides a fundamental alternative to adversarial learning with different strengths and weaknesses.

heading: 4 EXPERIMENT

heading: 4.1 2D DATASET COMPARISON WITH RELATED WORKS
text: We compare our method with related works in order to illustrate the limitations of prior methods and how our method overcomes these limitations. All implementation details for the 2D dataset are available in the appendix. Single T vs. Double T (LRMF vs. Ours) We first compare our method with LRMF (Usman et al., 2020) method. We construct the experiment to have the task: Translation between the two half-circled distributions: X1 and X2 in the moons dataset. In this example, we made two models, one with LRMF setup and one with our AUB setup. As illustrated in Fig. 3, the LRMF method fails to transform between X1 and X2. Even though Q can model well enough for T1(X1), Q can only model the mean and variance of X2 which is obviously not informative enough. Therefore, the LRMF fails to transform between two datasets. While in the AUB setup, both T1(X1) and T2(X2) are modeled to the same distribution which Q can be learned to fit with high likelihood which leads to better translation results. In conclusion, the performance of the LRMF model is limited by the power of the density model Q which means if Q fails to model one of the domain distribution in high likelihood, data alignment cannot be achieved with good performances. Simple Fixed Q vs. Learnable Q (AlignFlow vs. Ours) Next we compare our method with AlignFlow(Grover et al., 2020; Hu et al., 2018) setup. We construct the experiment to have the task: Transform between the two random patterns X1 and X2 from the randomly generated datasets. Again, we made two models with AlignFlow and our AUB setups respectively. As illustrated in Fig. 3, the AlignFlow method fails to transform between X1 and X2, because the transformed dataset T1(X1) and T2(X2) failed to reach the normal distribution Q. While in the AUB setup, the density model Q is learned to help fit the transformed distributions T1(X1) and T2(X2), which allows them to be aligned with each other easier. In conclusion, the performance of the AlignFlow model is limited by the performances of the invertible functions. Regularized vs. Un-regularized (Some prior works vs. Ours) We finally show the importance of the regularization term. We construct the experiment to have the task: Transform between two concentric circles with the same mean but slighted different radius. In this example, we make two models with our AUB approach, but one with regularization and one without. As illustrated in the Figure 4, both models are able to transform between two distributions. However, the transformation pattern is not natural in terms of the movement of each point. Each pair created by the unregularized model has bigger transportation cost compared to the pairs created by the regularized model. Therefore, we suggest that by adding transportation cost, the resulting transformation between samples will be closer to the identity function and therefore more stable.

heading: 4.2 REAL-WORLD DATASETS
text: Metrics In this subsection, we use two metrics to measure the overall model performances: Frechet Inception Distance (FID) (Heusel et al., 2017) implemented by (Seitzer, 2020) and our metric Alignment Upper Bound (AUB).
Unsupervised domain translation We first perform three image translation tasks on MNIST dataset; more specifically, we are flipping digit images in ”0/1”, ”0/2” , and ”1/2” pairs shown in Fig. 5. We use RealNVP invertible models for all translation maps Tj , as well as the density model Q. We evaluate our model performances in Table 1 along with baseline models: LRMF and MLE version of AlignFlow. Note that all methods are flow-based models and thus images translated back to original domain are exactly the same, which implies exact cycle consistency.
As represented in Table. 1, both of our approaches outperform the baseline models in terms of FID and AUB. The lower score in AUB indicates our model has the tighter bound than baseline model. This implies the statistical distance between the shared distribution {Zj}kj=1 and our density model Q is smaller than baselines, meaning our model has the better dataset alignment performance. We believe this result comes from our model setup, i.e., a learnable density model and domain specific invertible transformation functions with shared space. Specifically, AlignFlow with a fixed standard normal distribution as their Q obtains worse AUB because the Q is not powerful enough to model the complex shared space trained from the real world dataset. On the other hand, LRMF shows the lack of stability when trained with relatively simple models that we are using, i.e., RealNVP T and Real NVP Q. This is because the transformation function and the density model should be able to directly model the complex real distribution in their setup. By comparing our approach with and without transportation cost, we observe introducing transportation cost yields insignificant decrease in AUB and an increase in FID. This indicates the transportation cost regularizes our method well without degeneration, so that our method can find better solution (which is theoretically unique).
The impressive performance of our method in AUB shows a consistent pattern in FID score as well. Across entire translation cases, our method shows the better performance in terms of FID. This result can be intuitively understood by Fig. 5. AlignFlow shows less stable translation results than our method especially for translating to digit ‘2’ from digits ‘0’ and ‘1’. We believe this phenomenon comes from the lack of expressivity of their Q model. On the other hand, our model shows stable results across all translation cases, which are also quantitatively verified via the lower FID score. Also, from the comparison on the latent representations (second column of each macro column) of
ours with and without transportation cost, we can observe that the latent representation of the λ = 1 setup contains numerous excessive values. This implies adding transportation cost (lambda = 1) prevents the excessive movement in the latent space, thus, our model can find a simpler and smoother solution during the training.
Multi-domain translation To illustrate that our method can be easily scaled to more domain distributions, we present qualitative examples of translating between every digit and every other digit for MNIST in Fig. 6 with performances in Table 2. To show the effect of transportation cost proposed in our paper, we also compare our model with and without transportation cost. Note that we omit LRMF in this experiment because the multi-domain situation is hard to deal with for LRMF setup due to LRMF’s two distribution setup and assymetric model structure.
As shown in Table 2, our approach with and without transportation cost shows better performance in terms of FID and AUB than AlignFlow since our learnable density estimator can model more complex distribution than fixed simple density model in AlignFlow. In other words, the shared space of AlignFlow is limited because of the fixed simple density model. Note that the gap in the AUB score between AlignFow and our model is larger in this multi-domain setting (Table 2) than in the two domain case (Table 1). We hypothesize that this is because the shared density model across many domains needs to be more complex than for only two domains and thus, our learnable Q can create a better upper bound to optimize than AlignFlow. In this experiment, while ours with transportation cost performs slightly worse, the results are comparable and it does provide a natural regularization that ensures a theoretically identifiable solution.
The superiority of our method compared to AlignFlow in multiple-domain translation can also be verifed through qualitative comparisons in Fig. 6. The leftmost column is an input image and the second and third macro columns are the results from ours and AlignFlow. by forwarding a given k-th latent Tk(xk) into 10 inverse transformation functions, respectively. It is easy to observe that ours has more clear results in most cases than baseline. Moreover, our model shows better performance in maintaining the original identity (e.g., width and type of a stroke) than the baseline, as seen in third, fourth and ninth rows. This is because we jointly train our transformation functions with a learnable density model, while AlignFlow independently train their transformation functions. This benefit of our approach may be crucial for other datasets such as human faces (Choi et al., 2018) where maintaining the original identity is important.

heading: 5 DISCUSSION AND CONCLUSION
text: Transfer Learning Capability The performance of our model can be improved by adapting pretrained density models. Fig. 7(a) shows preliminary results. Specifically, a MLE version of AlignFlow model may still struggle to translate between USPS and MNIST dataset after learning for 200 epochs while ours can show qualitative results after only 1 epoch if we are using pretrained powerful density model for Q Salimans et al. (2017).
Generative Tasks Our model is also capable of generating samples in each domain. One needs to sample from the density model Q to have a latent image z first, and then forward to the inverse function T−1j to get the image sampled in j
th domain. Examples of generated images for each domain in MNIST data are shown in Fig. 7(b). The quality of the generative result also reflects the tightness of the bound between the latent space and the latent density model as illustrated in Eqn. 4.
In this paper, we proposed a novel upper bound on the generalized JSD that leads to a theoretically grounded alignment loss function. We then show that this framework unifies previous flow-based approaches to dataset alignment and demonstrate the benefits of our approach compared to prior flow-based methods. Additionally, we show that our framework could enable the use of pre-trained density models which would enable a type of transfer learning for distribution alignment. More broadly, we expect that our AUB metric can be useful as a domain-agnostic metric for comparing dataset alignment methods beyond images. An alignment metric that is not tied to a particular pretrained model (as for FID) or to a particular data type will be critical for systematic progress in unsupervised dataset alignment. We hope this paper provides one step in that direction.

heading: A PROOFS
text: Proof of Equivalence in Def. 1. While the proof of the equivalence is well-known, we reproduce here for completeness. As a reminder, the KL divergence is defined as:
KL(P,Q) = EP [log P (x)Q(x) ] = EP [− logQ(x)]− EP [− logP (x)] = Hc(P,Q)−H(P ) , (10)
where Hc(·, ·) denotes the cross entropy and H(·) denotes entropy. Given this, we can now easily derive the equivalence:
GJSDw(PX1 , · · · , PXk) = ∑ jwj KL(PXj , PXmix) (11)
= ∑
jwj(Hc(PXj , PXmix)−H(PXj )) (12) = ∑ jwj Hc(PXj , PXmix)− ∑ jwj H(PXj ) (13)
= ∑ jwjEPXj [− logPXmix ]− ∑ jwj H(PXj ) (14)
= ∑ jwj ∫ X −PXj (x) logPXmix(x)dx− ∑ jwj H(PXj ) (15)
= ∫ X − ∑ jwjPXj (x) logPXmix(x)dx− ∑ jwj H(PXj ) (16)
= ∫ X −PXmix(x) logPXmix(x)dx− ∑ jwj H(PXj ) (17)
= H(PXmix)− ∑ jwj H(PXj ) . (18)
Proof of Lemma 2. First, we note the following fact from the standard change of variables formula: PX(x) = PZ(T (x))|JT (x)|
⇒ PX(x)|JT (x)|−1 = PZ(T (x)) . (19)
We can now derive our result using the change of variables for expectations (i.e., LOTUS) and the probability change of variables from above:
H(PZ) = EPZ [− logPZ(z)] = EPX [− logPZ(T (x))] = EPX [− logPX(x)|JT (x)|−1] = EPX [− logPX(x)] + EPX [− log |JT (x)|−1] = H(PX) + EPX [log |JT (x)|] .
Proof of Theorem 3. Given any fixed Q, minimizing LAUB decouples into minimizing separate normalizing flow losses where Q is the base distribution. For each normalizing flow, there exists an invertible Tj such that Tj(Xj) ∼ Q, and this achieves the minimum value of LAUB. More formally,
min T1,··· ,Tk
LAUB(T1, · · · , Tk) (20)
= min T1,··· ,Tk
∑ jwjEPXj [− log |JTj (x)|Q(Tj(x))] (21)
= ∑
jwj minTj EPXj [− log |JTj (x)|Q(Tj(x))] + H(PXj )−H(PXj ) (22) = ∑
jwj minTj EPXj [− log |JTj (x)|Q(Tj(x))] + H(PXj )− EPXj [− logPXj (x)]) (23)
= ∑ jwj H(PXj ) + ∑ jwj minTj EPXj [log PXj (x)|JTj (x)| −1 Q(Tj(x)) ] (24)
= ∑ jwj H(PXj ) + ∑ jwj minTj EPXj [log PTj(Xj)(Tj(x)) Q(Tj(x)) ] (25)
= ∑ jwj H(PXj ) + ∑ jwj minTj EPTj(Xj) [log PTj(Xj)(z) Q(z) ] (26)
= ∑ jwj H(PXj ) + ∑ jwj minTj KL(PTj(Xj), Q) . (27)
Given that KL(P,Q) ≥ 0 and equal to 0 if and only if P = Q, the global minimum is achieved only if PTj(Xj) = Q,∀j and there exist such invertible functions (e.g., the optimal Monge map between PXj and Q for squared Euclidean transportation cost (Peyré and Cuturi, 2019)). Additionally, the optimal value is ∑ jwj H(PXj ), which is constant with respect to the Tj transformations.

heading: B EXAMPLES OF NON-UNIQUE ALIGNMENT SOLUTIONS
text: B.1 GAUSSIAN EXAMPLE
Suppose the component distributions are normal distributions, i.e., X1 ∼ N (µ1, I) and X2 ∼ N (µ2, I), and for even greater simplicity, we assume T2 is the identity, i.e., T2(x) = x. Then, a global optimal solution could be T1(x) = U(x− µ1 + µ2) for any orthogonal matrix U , i.e., there are infinitely many invertible functions that align the distributions. Note that this lack of unique solutions is not restricted to orthogonal rotations (see appendix for a more complex example).
B.2 COMPLEX EXAMPLE
Consider the 1D case where Q only contains the uniform distribution. Thus, T1 and T2 must map their distributions to the uniform distribution for alignment. One solution would be that T1 = F1 and T2 = F2 where F1 and F2 are the CDFs of PX1 and PX2 . Yet, there are infinitely many other possible solutions. Consider an invertible function that subdivides the unit interval into an arbitrarily large number of equal length intervals and then shuffles these intervals with a fixed arbitrary permutation. More formally, we could define this as:
Sm,π(x) =  x− 1m + π(1) m if x ∈ [0, 1 m ) x− 2m + π(2) m if x ∈ [ 1 m , 2 m ) ... ...
x− mm + π(m) m if x ∈ [ m−1 m , 1]
, (28)
where π(·) is a permutation of the integers 1 to m. Given this, then other optimal solutions could be T1 = Sm,π ◦F1 and T2 = F2 for any m > 1 and any permutation π. This idea could be generalized to higher dimensions as well by mapping to the multivariate uniform distribution and subdividing the unit hypercube similarly.

heading: C ALGORITHM
text: We summarize our computation of regularized alignment upper bound in Alg. 1.
Algorithm 1 Training algorithm for our model
Input: datasets {Xj}kj=1 for k domains; normalizing flows {Tj(xj ; θj)}kj=1; density model Q(z;ϕ); learning rate η; transportation cost factor λ maximum epoch Emax; initial parameters value {T (0)j (xj ; θj)}kj=1, Q(0)(z;ϕ); Output: {θ∗}kj ; for epoch = 1, Emax do
for each batch {xj}kj=1 do ϕ← ϕ+ η · λθ logQ({Tj(xj ;ϕ)}kj=1) ▷ Update Q using data from all domains end for for each batch {xj}kj=1 do
for domainj = 1, k do θ ← θ + η · log |JTj (x)|Q(Tj(x))− λc(x, Tj(x)) ▷ Update T independently across
all domains end for
end for end for

heading: D DETAILED PARAMETERS USED IN 2D DATASET EXPERIMENT
text: D.1 LRMF VS. OURS EXPERIMENT
• T for LRMF setup: T1: 8 channel-wise mask for Real-NVP model with s and t derived from 64 hidden channels of fully connected networks. T2: Identity function.
• T for RAUB setup: T1 and T1: 8 channel-wise mask for Real-NVP model with s and t derived from 64 hidden channels of fully connected networks. Regularization coefficient λ = 0
• Q for both: A single Gaussian distribution with trainable mean and trainable variances.
D.2 ALIGNFLOW VS. OURS EXPERIMENT
• T for both: 2 channel-wise mask for RealNVP model with s and t derived from 8 hidden channels of fully connected networks.
• Q for Alignflow setup: A single fixed normal distribution.
• Q for RAUB setup: A learnable mixture of Gaussian with 3 components. Regularization coefficient λ = 0
D.3 REGULARIZED VS. UNREGULARIZED EXPERIMENT
1. T for both: 8 channel-wise mask for RealNVP model with s and t derived from 64 hidden channels of fully connected networks.
2. Q for both: A learnable mixture of Gaussian with 2 components.
3. λ for unregularized Experiment : λ = 0
4. λ for regularized Experiment : λ = 1
E INTERPOLATION OVER LATENT SPACE
We use the same model used in the ”Multi-domain translation” part in section 4.2 and perform interpolation in the latent space. We first randomly select two distinct real images in one domain (in this case two 0s), and do a linear interpolation of the selected two images in the latent space. Then we translate all the interpolated images (including the two selected images) to all of the domains to generate ”translated-interpolated” images—i.e., the corresponding interpolations in each of the domains.
As shown in Fig. 8, all ”translated-interpolated” results can preserve the trend of the stroke width of the digits from the original interpolated domain. These results suggest that our approach aligns the domains so that some latent space directions have similar semantic meaning for all domains.

heading: F TABULAR/STRUCTURED DATA EXPERIMENT
text: We follow the same preprocessing for four UCI tabular datasets (Dua and Graff, 2017) from the MAF paper (Papamakarios et al., 2017), specifically MINIBOONE, GAS, HEPMASS, and POWER. In order to obtain a binary label for creating domains, we choose the last input feature for each dataset and discretize it based on whether it is higher or lower than the median value, which ensures the datasets are of equal size. We use the MLE version of AlignFlow as a baseline to compare with our model and use the AUB score as the metric for comparison.
As shown in the Table. 3, our method shows better performance with a large gap across all comparisons. This provides strong evidence that our proposed method (trained with learnable Q) better aligns the two domains than AlignFlow (trained with fixed Q).
It is worth emphasizing that there is no natural metric for evaluating GAN (Goodfellow et al., 2014) in tabular dataset. We believe this demonstrates one of the key benefits of our proposed method over GAN-based alignment methods.

heading: G MULTI-TRANSLATION TASKS IN CELEBA DATASET
text: In order to verify our proposed method performs well in a more complex setup than MNIST, we conducted experiments on the CelebA (Liu et al., 2015) dataset. We first create three domains (Black Hair, Blond Hair and Brown Hair) by using the hair color attribute. We further center crop
the images to have the size 148 × 148 then resize them to 64 × 64. For the baseline model, we again uses MLE version of AlignFlow. We use both FID and AUB as a metric to compare the performances.
As shown in the Table 4, our proposed method demonstrates better performance in almost all experiments in terms of both FID and AUB. This consistently validates that the flexibility of Q in our framework has clear advantages over the fixed Q in AlignFlow for the domain alignment task. Intuitively, we think the tighter upper bound from our proposed method plays an important role in aligning distributions in the latent space.
We also provide qualitative results in Fig. 9. Compared to the AlignFlow (MLE) baseline, our model performs significantly better in translating between selected attributes. In particular, while our approach maintains a valid face structure during translation, AlignFlow is unable to maintain a valid face structure. Note that absolute image quality is not the focus of our paper as we use relatively simple flow models, but rather the comparison to prior methods and to show the feasibility of our approach in higher dimensions. We expect better qualitative results could be achieved using more advanced flow models such as GLOW, Residual Flows, Flow++, etc. but given that the scope of our work which focuses on the theoretical and foundational approach, we leave this to future work.

Review summaries:
The paper presents a unified framework for domain alignment through non-adversarial approaches. Instead of solving a min-max adversarial problem, this paper aims to solve a min-min optimization equivalent to minimizing the upper bound of JSD. It proves the equivalence and offers a computational and feasible solution that can be inserted into flow-based non-adversarial approaches in a plug-and-play fashion. Preliminary results on generating digits from different classes or domains have demonstrated the correctness of the model.

This paper proposes a flow-based method for the unsupervised data set alignment problem. It firstly reveals that the minimization problem over density models can be addressed by optimizing the upper bound of generalized JSD. Based on this theoretical result, the authors propose a new regularizer for the multi-distribution alignment problem. Compared to prior work, e.g., AlignFlow and LRMF, the authors derive a more general framework for the unsupervised data set alignment problem. The authors also provide extensive experiments to verify its effectiveness and superiority.