Cooperative Distribution Alignment via JSD Upper Bound
Abstract
Unsupervised distribution alignment estimates a transformation that maps two or more source distributions to a shared aligned distribution given only samples from each distribution.
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 and are limited in efficiently aligning multiple distributions.
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
Jensen-Shannon Divergence (JSD).
Importantly, our problem reduces to a minmin, i.e., cooperative, problem and can provide a natural evaluation metric for unsupervised distribution alignment.
We present empirical results of our framework on both simulated and real-world datasets to demonstrate the benefits of our approach. Code is available at https://github.com/inouye-lab/AUB .
Introduction
In many cases a practitioner has access to multiple related but distinct distributions such as agricultural measurements from two farms, experimental data collected in different months or sales data before and after a major event.
Unsupervised distribution alignment (UDA) is the ML task aimed at aligning these related but distinct distributions in a shared space, without any pairing information between the samples from these distrbibutions (i e., unsupervised).
This task has many applications such as generative modeling (e.g., [1]), unsupervised domain adaptation (e.g., [2, 3]), batch effect mitigation in biology (e.g., [4]) and fairness-aware learning (e.g., [5]).
The most common approach for obtaining such alignment transformations stems from Generative
Adversarial Networks (GAN) [6], which can be viewed as minimizing a lower bound on the JensenShannon Divergence (JSD) between real and generated distributions.
The lower bound is tight if and only if the inner maximization is solved perfectly.
CycleGAN [1] maps between two datasets via two GAN objectives between the 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. [7–11]) in part because of the competitive nature of the min-max optimization problem.
Perhaps more importantly, the research community has resorted to surrogate evaluation metrics for GAN because likelihood computation is intractable.
Specifically, the commonly accepted Frechet Inception Distance (FID) [12] is only applicable to image or auditory data for which there exist publicly available classifiers trained on large-scale data.
Moreover, the concrete implementation of FID can have issues due to seemingly trivial changes in image resizing algorithms [13].
Recently, flow-based methods with a tractable likelihood have been proposed for the UDA task [2, 14– 16].
Specifically, iterative flow methods [14, 15] proposed alternative approaches to distribution alignment via iteratively building up a deep model via simpler maps.
AlignFlow [2] leverages invertible models to make the model cycle-consistent (i.e., invertible) by construction and introduces exact log-likelihood loss terms derived from standard flow-based generative models that complement the adversarial loss terms.
On the other hand, log-likelihood ratio minimizing flows (LRMF) [16] 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, iterative flow models [14, 15] do not explicitly reduce a global divergence measure and actually solve non-standard adversarial problems via alternating optimization (see [15]).
AlignFlow [2] assumes that the shared density model for all distributions is a fixed Gaussian and lacks an explicit alignment metric.
LRMF [16] may only partially align distributions if the target distribution is not in the model class. Also, the LRMF metric is limited because it is only defined for two distributions and depends on the shared density model class.
To address these issues, we unify existing flow-based methods (both AlignFlow and LRMF) under a common cooperative (i.e., non-adversarial) framework by proving that a minimization over a shared density model is a variational upper bound of the JSD.
The unifying theory also suggests a natural domain-agnostic metric for UDA that can be applied to any domain including tabular data (where FID is inapplicable).
This metric is analogous to the Evidence Lower Bound (ELBO), i.e., it is a variational bound that is useful for both training models and comparing models via heldout test evaluation.
Furthermore, this unification enables straightforward and parameter-efficient multi-distribution alignment because the distributions share a latent space density model.
We summarize our contributions as follows:
• We prove that a minimization over a shared variational density model is a variational upper bound on a generalized version of JSD that allows for more than two distributions.
Importantly, we theoretically quantify the bound gap and show that it can be made tight if the density model class is flexible enough.
• Based on this JSD upper bound, we derive a novel unified framework for cooperative (i.e., non-adversarial) flow-based UDA that includes a novel domain-agnostic AUB metric and explain its relationship to prior flow-based alignment methods.
• Throughout experiments, we demonstrate that our framework consistently shows superior performance in both our proposed and existing measures on simulated and real-world datasets.
We also empirically show that our model is more parameter-efficient than the baseline models.
Notation We will denote distributions as P X ( x ) where X is the corresponding random variable.
Invertible functions will be denoted by T ( · ) .
We will use X j ∼ P X j to denote the observed random variable from the j -th distribution.
We will use Z j ≜ T j ( X j ) ∼ P
Z j ≡ P T j ( X j ) to denote the latent random variable of the j -th distribution after applying T j to X j (and note that X j = T − 1 j ( Z j ) ).
We will denote the mixtures of these observed or latent distributions as P X mix ≜ (cid:80) j w j P X j and
P Z mix ≜ (cid:80) j w j P Z j , where w is a probability vector.
We denote KL divergence, entropy, and cross entropy as KL( · , · ) , H( · ) , and H c ( · , · ) , respectively, where KL( P, Q ) =
H c ( P, Q ) − H
( P ) .
Alignment Upper Bound Loss
In this section, we will present our main theoretical result by proving an upper bound on the generalized JSD divergence, deriving our loss function based on this upper bound, and then showing that minimizing this upper bound results in aligned distributions assuming the model components have a large enough capacity.
Background: Normalizing Flows and Invertible Models.
Normalizing flows are generative models that have tractable distributions where exact likelihood evaluation and exact sampling are [Figure] [Equation] [Figure]
𝑄 𝑃
[Figure]
𝑃
[Figure] possible [17].
Flow models leverage the change of variables formula to create an invertible mapping T such that P X ( x ) =
P Z ( T ( x ))
| J T ( x ) | where P Z is a known latent distribution and | J T ( x ) | is the absolute value of determinant of the Jacobian of the invertible map T .
To sample from P X , one needs to first sample from the latent distribution P Z
and then apply the inverse transform T − 1 .
Therefore, the key challenge in designing invertible models is to have computationally efficient inverse evaluation for sampling and Jacobian determinant calculation for training.
Many approaches have been proposed by parameterizing mapping function T as deep neural networks including autoregressive structures [18, 19], coupling layers [20, 21], ordinary differential equations[22], and invertible residual networks[23, 24].
Flow models can be then learned efficiently by maximizing the likelihood of 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 w 1 = w 2 = 0 .
Definition 2.1 (Generalized Jensen-Shannon Divergence (GJSD) [25]) .
Given k distributions { P Z j } kj =1 and a corresponding probability weight vector w , the generalized Jensen-Shannon divergence is defined as (proof of equivalence in Appendix A.):
[Equation]
2.1 GJSD Variational Upper Bound
The goal of distribution alignment is to find a set of transformations { T j ( · ) } kj =1 (which will be invertible in our case) such that the latent distributions align, i.e., P T j ( X j ) =
P T j ′
X j ′ ) or equivalently
P Z j = P
Z j ′ for all j ̸ = j ′ .
Given the properties of divergences, this alignment will happen if and only if GJSD( P Z 1 , · · · , P Z k ) = 0 .
Thus, ideally, we would minimize GJSD directly, i.e.,
[Equation]
(1) where T is a class of invertible functions.
However, we cannot evaluate the entropy terms in Eqn.because we do not know the density of P X j ; we only have samples from P X j .
Therefore, we will upper bound the first entropy term in Eqn.
1 ( H( (cid:80) j w j P T j ( X j ) ) ) using a variational density model and decompose the other entropy terms via the change of variables formula for invertible functions.
Theorem 2.2 (GJSD Variational Upper Bound) .
Given an variational density model class Q , we form a GJSD variational upper bound:
[Equation]
where P Z mix ≜ (cid:80) j w j P Z j and the bound gap is exactly min Q ∈Q KL( P Z mix , Q ) .
Proof of Theorem 2.2.
For any Q ∈ Q , we have the following upper bound:
[Equation]
where the first equals is merely inflating by H c ( P Z mix , Q ) , the inequality is by the fact that KL divergence is non-negative, and the bound gap is equal to KL( P Z mix , Q ) .
The Q that achieves the minimum in the upper bound is equivalent to the Q that minimizes the bound gap, i.e.,
[Equation]
mix , Q ) , where the second equality is because the entropy terms are constant with respect to Q and the last is by the definition of KL divergence.
The tightness of the bound depends on how well the class of density models Q (e.g., mixture models, normalizing flows, or autoregressive densities) can approximate P Z mix ; notably, the bound can be made tight if P Z mix ∈ Q .
Also, one key feature of this upper bound is that the cross entropy term can be evaluated using only samples from P X j and the transformations T j ,
i.e., H c ( P Z mix , Q ) = (cid:80) j w j E P Xj [ − log Q
( T j ( x j ))] .
However, we still cannot evaluate the other entropy terms H( P Z j ) since we do not know the density functions of P Z j (or P X j ).
Thus, we leverage the fact that the T j functions are invertible to define an entropy change of variables.
Lemma 2.3 (Entropy Change of Variables) .
Let X ∼ P X and Z ≜ T ( X ) ∼ P Z , where T is an invertible transformation.
The entropy of Z can be decomposed as follows:
[Equation]
where | J T ( x ) | is the absolute value of the determinant of the Jacobian of T .
The key insight from this lemma is that H( P X ) is a constant with respect to T and can thus be ignored when optimizing T , while E P X
[log | J T ( x ) | ] can be approximated using only samples from
P X (formal proof in Appendix A).
2.2 Alignment Upper Bound (AUB)
Combining Theorem 2.2 and Lemma 2.3, we can arrive at our final objective function which is equivalent to minimizing the variational upper bound on the GJSD:
[Equation]
where the cross entropy term is replaced by its definition provided above, and the last term − (cid:80) j w j H( P X j ) is constant with respect to T j functions so they can be ignored during optimization.
We formally define this loss function as follows.
Definition 2.4 (Alignment Upper Bound Loss) .
Given k continuous distributions { P X j } k j =1 , a class of continuous distributions Q , and a probability weight vector w , the alignment upper bound loss is defined as follows:
[Equation]
where T j are invertible and | J T j ( 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 T j .
We now consider the theoretical optimum if we optimize over all invertible functions.
Theorem 2.5 (Alignment at Global Minimum of L AUB ) .
If L AUB is minimized over the class of all invertible functions, a global minimum of L AUB implies that the latent distributions are aligned, i.e.,
P T j ( X j ) =
P T j ′
X j ′ ) 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 P T j ( X j ) =
Q for all T j (formal proof in Appendix A).
This alignment of the latent distributions also implies the translation between any of the observed component distributions.
The proof follows directly from Theorem 2.
and the change of variables formula.
Corollary 2.6 (Translation at Global Minimum of L AUB ) .
Similar to Theorem 2.5, a global minimum of L AUB implies translation between any component distributions using the inverses of T j , i.e.,
P T − 1 j ′ ( T j ( X j ))
=
P X j ′ for all j ̸ = j ′ .
As seen in Alg. 1, we use a simple alternating optimization scheme for training our translation models and variational distribution with cooperative (i.e., min-min) AUB objective, i.e., we aim to optimize:
[Equation]
We emphasize that our framework allows any invertible function for T j (e.g., coupling-based flows [20], neural ODE flows [22], or residual flows [23]) and any density model class for Q z (e.g., kernel densities in low dimensions, mixture models, autoregressive densities [26], or normalizing flows [21]).
Even VAEs [27] could be used where the log likelihood term is upper bounded by the negative
ELBO, which will ensure the objective is still an upper bound of GJSD.
Input: Datasets { X j } kj =1 for k domains; n as the batch size; x j as a minibatch for the j -th domain; normalizing flow models { T j ( x j ; θ j ) } kj =1 ; density model Q ( z ; ϕ ) ; learning rate η ; maximum epoch E max
Output: { ˆ θ } kj =1 ; for epoch = 1 , E max do for each batch { x j } kj =1 do ϕ ← ϕ + η ∇ ϕ 1 k (cid:80) kj =1 1 n (cid:80) ni
=1 log Q ( T j ( x i,j ; θ j ); ϕ ) end for for each batch { x j } kj =1 do ∀ j, θ j ← θ j + η ∇
AUB for UDA is like ELBO for density estimation.
Although AUB and ELBO are for fundamentally different tasks, we would like to point out the similarities between AUB and ELBO.
First, both are variational bounds of the quantity of interest where the tightness of the bounds depend on the optimization of the variational distributions.
Second, both can be made tight if the class of variational distributions is powerful enough.
Third, both can be used to train a model by minimizing the objective on training data.
Fourth, while neither can be used as an absolute performance metric, they can both be used to evaluate the relative performance of models on held-out test data.
Thus, AUB can be used for UDA as ELBO has been used for density estimation with the same strengths and weaknesses such as being a relative metric and requiring an auxiliary model for evaluation.
Relationship to Prior Works
AlignFlow without adversarial terms is a special case.
As illustrated in Fig.
2, AlignFlow [2] 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).
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.
By using the same architecture of T and Q , AlignFlow can be also viewed as sharing the last few layers of the T ’s; however, we note that our approach allows for Q that are not flows, e.g., autoregressive densities or mixture models as in our toy experiments that even have alternative non-SGD learning algorithms.
LRMF is special case with only one transformation.
As illustrated in Fig.
2, Log-likelihood ratio minimizing flows (LRMF) [16] is also a special case of our method for only two distributions, where one transformation is fixed at the identity (i.e., T 2 = Id).
While the final LRMF objective is a special case of ours, the theory is developed from a different but complementary perspective.
The LRMF metric depends on the shared density model class, which enables a zero point (or absolute value) of the metric to be estimated but requires fitting extra density models.
Usman et al.
[16] do not uncover
AlignFlow (MLE only)
LRMF
Ours
Fixed dist.
Invertible function
Changed dist.
Q dist.
Getting closed
Gettingclosed
Gettingclosed 𝑋 !
𝑋 " 𝑍 " 𝑍 ! 𝑇 !
𝑇 " 𝑋 !
𝑋 " 𝑍 ! 𝑋 !
𝑋 " 𝑍 ! 𝑍 " learnable 𝑄 learnable 𝑄 fixed 𝑄 𝑇 !
𝑇 !
𝑇 "
[Figure]
the connection of the objective as an upper bound on JSD regardless of the density model class 3 .
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.
In other words, our bound holds regardless of the model class Q , whereas the LRMF-JSD discussion further assumes that the target distribution needs to be in the Q model class (see Theorem 2.5 and our simulated experiments).
We analyze the performance of our proposed framework comparing to the relevant flow-based baseline models. In subsection 4.1, we experiment on a toy dataset to clearly show the benefit of our method over the baseline flow-based models in a controlled environment.
In subsection 4.2 and subsection 4.3, we demonstrate AUB’s superiority over baseline flow-based models on real-world datasets including tabular data and high-dimensional MNIST data. In subsection 4.4, we additionally conduct Domain
Adaptation [2, 28] experiments for validating AUB alignment in the context of a downstream task. Implementation details are provided in Appendix B.
4.1 Toy dataset comparison with related works
Single T vs. Double T (LRMF vs. Ours).
We first compare our method with LRMF [16] method.
The task is to translate between the two half-circle distributions (i.e., “moons”).
We compared our AUB setup with two maps T 1 and T 2 to the LRMF setup with only T 1 where Q is the set of independent Gaussian distributions.
As illustrated in Fig.
3, the LRMF method fails to transform between X 1 and X 2 .
Even though Q can model the transformed data T 1 ( X 1 ) , a Gaussian-based Q cannot fully model the half-circle distribution of X 2 .
Therefore, LRMF fails to transform between two distributions.
However, in the AUB setup, both latent distributions of T 1 ( X 1 ) and T 2 ( X 2 ) can be modeled by the same Gaussian distribution Q because of the flexibility in both transformations, which leads to better translation results in this illustrative example.
In conclusion, the performance of LRMF is limited by the choice of density models Q ; if Q fails to model the target distribution, distribution alignment may not be achieved.
Simple Fixed Q vs. Learnable Q (AlignFlow vs. Ours).
Next we compare our AUB setup with the AlignFlow [2, 3] setup.
The task is to translate between two random Gaussian mixture datasets (i.e., "blobs"). We compared our AUB setup where Q is a mixture of Gaussians with the AlignFlow setup where Q is a fixed standard normal distribution (the models for T 1 and T 2 are the same).
As illustrated in Fig.
3, the AlignFlow method fails to transform between X 1 and X 2 , because the transformed dataset T 1 ( X 1 ) and T 2 ( X 2 ) failed to reach the fixed standard normal distribution Q .
However, in the AUB setup, the shared density model Q adapted to the distributions of T 1 ( X 1 ) and T 2 ( X 2 ) to enable a tighter alignment bound and thus the translation results are better.
In conclusion, the performance of the AlignFlow model is limited by the power of the invertible functions.
4.2 Unsupervised Domain Alignment on Tabular Datasets
To showcase the application-agnostic AUB metric and the parameter efficiency of our AUB framework, we conduct two experiments on real-world tabular datasets.
In both experiments, we used four UCI tabular datasets [29] (MINIBOONE, GAS, HEPMASS, and POWER), following the same preprocessing as the MAF paper [30].
Train, validation, and test sets are 80%, 10%, and 10% of the data respectively.
Also, the experiments are measured by test AUB defined in Definition 2.4, where a lower AUB score indicates better performance (see end of section 2.2 for discussion on using test AUB for evaluation).
We emphasize that there is no natural metric for evaluating GAN-based alignment methods on tabular datasets.
Thus, these experiments demonstrate one of the key benefits of our proposed framework over a GAN-based approach.
Our first experiment was designed to compare alignment performance between our proposed method and the baseline methods.
To separate each dataset into two distributions, we choose the last input feature from 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.
Given the divided dataset, two transformation functions ( T ) are trained to align the distributions.
For baselines, we use AlignFlow MLE, Adv.
only, and hybrid versions and LRMF on top of the original implementation.
Because AlignFlow Adv.
and hybrid setups optimize over a mixed objective of AUB (special case) and adversarial losses, to be fair, we additionally fit an identity-initialized flow model Q to the final T ’s. We use the same T and Q models wherever possible across all methods (e.g., the same Q is used for LRMF and AUB
but AlignFlow has a fixed Q ).
As shown in Table 1, our method shows better performance compared to other methods across all datasets.
In particular, because ours and LRMF (trained with learnable Q ) outperform AlignFlow
MLE (trained with fixed Q ), we can see that learnable Q plays an important role in distribution alignment.
We also observe that ours shows better performance than LRMF, where the gap may come from aligning in the shared latent space (ours) rather than aligning in the original data space (LRMF).
AlignFlow Adv.
only and hybrid versions show worse performance than ours with a large gap, which implies our proposed cooperative training is competitive to adversarial methods in aligning distributions of tabular data.
The second experiment demonstrates alignment between 8 distributions and is designed to show that our proposed method can be more efficient in terms of model parameter compared to the baseline methods.
For this multi-distribution experiment, we separate each dataset into 8 domains by choosing the last three features from a given dataset and dividing the dataset by the three medians, e.g., (+ + +) , (+ + − ) , ..., ( − − − ) .
The results are shown in Table 2 where RealNVP (5) and RealNVP (10) indicates the number of coupling layers used in the RealNVP architecture respectively.
LRMF is excluded because LRMF is only designed for aligning two distributions.
Comparing AlignFlow RealNVP (5) and RealNVP (10), we can see that AlignFlow can achieve better alignment with more than double amount of parameters. However, as can be observed in a [Table]
[Table] comparison between AlignFlow with RealNVP (10) and ours, our proposed method can achieve similar performance with the baseline model but with less than half the parameters.
These two results suggest that our model can be more parameter efficient than AlignFlow for multi-distribution alignment.
We hypothesize that our approach may scale better with respect to the number of distributions because our Q shares parameters across the distributions and can capture the similarities between distributions.
4.3 Unsupervised Distribution Alignment on MNIST Dataset
We perform an image translation task on MNIST dataset 4 [31] to demonstrate that our distribution alignment method can be applied to high-dimensional datasets and to validate our AUB metric in a case where FID is also available.
We train ours and baseline models with the digit images of 0, 1, and 2, and compare the translated results both quantitatively and qualitatively.
Specifically, we use RealNVP invertible models for all translation maps T j , as well as the density model Q .
Note that all methods are flow-based models and thus images translated back to the original domain are exactly the same, which implies exact cycle consistency.
As represented in Table 3, both of our approaches outperform the baseline models in terms of FID and AUB.
The lower AUB score indicates our method seems to align the distributions better than baselines.
We believe this result comes from our model setup, i.e., a learnable shared density model and transformation to a shared latent 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 the relatively simple models that we are using across all methods, i.e.,
RealNVP T and RealNVP Q .
We expect this is caused by the restrictions of only using one T for translation without a shared latent space and the fact that the Q distribution must be able to model the target distribution to ensure alignment. better performance FID score for by averaging scores from is shown [Table]
The better quantitative performance of our methods can be corroborated by the qualitative results as seen in Fig.
AlignFlow shows less stable translation results than our method especially for translating to digit ‘2’ from digits ‘0’ and ‘1’.
The second column of LRMF is set to be black because it does not have a latent representation, and LRMF fails to translate in this situation, which is why the translated results are nearly the same.
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.
Note that our method has a shared space, so different transformation functions for each domain are trained together.
This shared space may provide advantage in terms of sample complexity and computational complexity compared to multiple independent flows as in AlignFlow.
Similarly, we conducted experiments on a ten-domain translation in Appendix C to illustrate that our method with a shared space can be easily scaled to more domain distributions with less number of parameters than the baseline model. Additionally, we visualize an interpolation over the shared space from the ten-domain experiment. Detailed descriptions and the results are provided in Appendix D 4.4 Domain Adaptation on USPS-MNIST dataset
We additionally conducted a domain adaptation (DA) experiment between MNIST to USPS for externally validating our alignment performance compared to baseline methods.
We first reduce the dimensionality of both datasets to 32 using a pretrained variational autoencoder (VAE), which is trained jointly on both MNIST and USPS images without any label information (i.e., an unsupervised
VAE).
We then learn a translation maps in this latent space via our method and baselines.
Finally, following the typical DA evaluation protocol, we train the classifier with source domain data and evaluate the performance by applying the classifier on the target domain data translated to the source domain.
Image translation results can be obtained by forwarding the latent translated results to the pretrained decoder.
As seen in Table 4, our method (right column) performs better than baseline methods.
This result is further empirical evidence that our novel cooperative training can be comparable to adversarial training in certain cases.
As shown in Fig.
5, adversarial loss shows mode collapse, e.g., AF(Adv. only), AF(1e-2).
without careful hyperparameter tuning.
The other thing to note is AF (MLE) works for domain alignment itself (i.e., the translated result is in MNIST domain), but it does not maintain the class information (i.e., the digit changes from USPS to MNIST).
We conjecture that this issue stems from AlignFlow’s simple and fixed Q distribution where arbitrary rotations of the latent space have equivalent likelihoods throughout the whole training process and thus class information is lost while transforming to and from such a Q distribution.
On the other hand, our learnable Q guides the alignment training process such that class information is partially preserved during transformation.
LRMF failed to align since the Q distribution is not complex enough to model the density of the target distribution.
This is the consistent result with our toy dataset experiment (Fig. 3, (a)-(c)).
LRMF does not work well with the simple Q while ours can align two different distributions with a simple
Q because of the shared space.
Please note that all models in the table use approximately the same number of model parameters.
Further comparisons are provided in Appendix E.
[Table]
Pros and cons of our method compared to adversarial method.
Flow-based methods have different benefits and limitations compared to adversarial methods for distribution alignment.
As one clear difference, our framework provides an application-agnostic yet theoretically grounded evaluation metric for comparing alignment methods.
Additionally, our min-min problem is fundamentally different than a min-max problem and avoids issues unique to min-max problems (as seen in Fig. 6).
However, a key limitation of our approach compared to GANs is that it is restricted to invertible models. Also, our method requires training a density model, whereas GANs train a discriminator.
Overall, we suggest our alignment approach is a feasible and fundamentally different alternative to adversarial, which is currently the only dominant approach.
With our foundation, future work could focus on the performance aspects just as adversarial learning has been improving over the past seven years but is still an active area of research.
[Figure]
Alternating minimization against vanishing gradient problem.
As revealed in previous works [16, 32], the vanishing gradient problem can occur while minimizing a JSD approximation.
On the same line, AUB can also suffer from vanishing gradient if the distributions are disjoint, and Q is fit very well.
However, any Q that does not match P Z mix provides an upper bound that we conjecture is smoother than the true JSD.
Thus, in experiments, we have found that alternating minimization does not suffer from vanishing gradients because Q is not fully fitted at each step but provides a smooth upper bound.
Yet, deeper theoretic and empirical analysis is needed in future work to fully understand this case. Extensibility.
As mentioned in section 1 and section 2, our proposed idea is a general framework that can be used with any invertible flow models (e.g., Residual Flows [23], Flow++ [33]) and any density model (e.g., PixelCNN++ [26], FFJORD [22]).
Hence, we argue our proposed idea is not limited to specific flow-based model or density model.
We expect our framework could show better performance in distribution alignment if it is combined with the aforementioned state-of-the-art components.
Conclusion
In this paper, we propose a novel variational upper bound on the generalized JSD that leads to a theoretically grounded alignment loss.
We then show that this framework unifies previous flow-based distribution alignment approaches and demonstrate the benefits of our approach compared to these prior flow-based methods.
In particular, our framework allows a straightforward extension to multidistribution alignment that could be more parameter efficient than naïvely extending prior approaches.
More broadly, we suggest that our AUB metric can be useful as an application-agnostic metric for comparing distribution alignment methods (analogously to how ELBO is used to evaluate density estimation methods).
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 distribution alignment.
We hope this paper provides one step in that direction.
Acknowledgement.
The authors thank the anonymous reviewers for their insightful comments and helpful discussion that improved the final paper.
The authors acknowledge support from ARL (W911NF-2020-221) and NSF (2212097).
A Proofs
Proof of Equivalence in Definition 2.1 in the main paper.
While the proof of the equivalence is wellknown
[Equation]
we reproduce here for completeness.
As a reminder
[Equation] the KL divergence is defined as: [Equation] where H c ( ·
[Equation]
· ) denotes the cross entropy and H( · ) denotes entropy. Given this [Equation] we can now easily derive the equivalence:
[Equation]
Proof of Lemma 2.3 in the main paper.
First, we note the following fact from the standard change of variables formula:
[Equation]
We can now derive our result using the change of variables for expectations (i [Equation] e., LOTUS) and the probability change of variables from above: [Equation]
Proof of Theorem 2.5 in the main paper.
Given any fixed Q
[Equation]
minimizing L AUB decouples into minimizing separate normalizing flow losses where Q is the base distribution.
For each normalizing flow
[Equation]
there exists an invertible T j such that T j ( X j ) ∼ Q
[Equation]
and this achieves the minimum value of
L AUB .
More formally, [Equation] Given that KL( P, Q ) ≥ 0 and equal to 0 if and only if P = Q , the global minimum is achieved only if P T j ( X j ) =
Q, ∀ j and there exist such invertible functions (e.g., the optimal Monge map between P X j and Q for squared Euclidean transportation cost [34]).
Additionally, the optimal value is (cid:80) j w j H( P X j ) , which is constant with respect to the T j transformations.
B Additional Experiment Details B.1 Toy Dataset Experiment LRMF vs. Ours Experiment
• T for LRMF setup: T 1 : 8 channel-wise mask for Real-NVP model with s and t derived from 64 hidden channels of fully connected networks.
T 2 : Identity function.
• T for RAUB setup: T 1 and T 1 : 8 channel-wise mask for Real-NVP model with s and t derived from 64 hidden channels of fully connected networks.
• Q for both: A single Gaussian distribution with trainable mean and trainable variances.
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.
B.2 Tabular Dataset Experiment
Our invertible transformation function T adapts general purpose RealNVP (5) and RealNVP (10) layers with detailed parameters provided in Table 5.
Successive coupling layers are concatenated by alternating the format between keeping odd-parity index of the data samples and transforms the even-parity index of the data and vice versa.
Note that our density model Q for AUB experiments shares the same architecture of RealNVP (5) throughout the two-domain tabular dataset experiment and RealNVP (10) for multi-domain experiment.
For AlignFlow hybrid and adversarial only models, in order to adapt their model for a tabular dataset, we change the invertible model T _ src and T _ tgt to
RealNVP(5) and change the discriminator to a fully connected network which has a hidden dimension of 256.
[Table]
C Multi-domain translation
To illustrate that our method can be easily scaled to more domain distributions even for high dimensional data
[Table] we present qualitative examples of translating between every digit and every other digit for MNIST in Fig. 7 with quantitative performances in Table 6. 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
[Table]
our approach 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
[Table]
[Figure]
model in AlignFlow. In other words, the shared space of AlignFlow is limited because of the fixed simple density model. The superiority of our method compared to AlignFlow in multiple-domain translation can also be verified through qualitative comparisons in Fig. The leftmost column is an input image and the second and third macro columns are the results from AlignFlow and ours.
By forwarding a given k -th latent T k ( x k ) 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 fifth, eighth 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 [35] where maintaining the original identity is important. D Implementation on Latent Space 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 − 1 j to get the image sampled in j th domain.
Examples of generated images for each domain in MNIST data are shown in Fig. 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 order to visualize the benefit of our shared latent space, we further perform interpolation in the latent space. We first randomly select two distinct real images in one domain (in this case two ‘0’s), 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 remaining domains to generate ‘translated-interpolated’ images, i.e., the corresponding interpolations in each of the remaining domains. As shown in Fig. 9, 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 shared space aligns the domains so that some latent space directions have similar semantic meaning for all domains.
Here are the translated results that we used in Table 4. The results consistently show that ours show the best performance. Please note that these samples are randomly chosen.