Cooperative Distribution Alignment via JSD Upper Bound
Abstract
Unsupervised distribution alignment estimates a transformation that maps two ormore source distributions to a shared aligned distribution given only samples fromeach distribution.
This task has many applications including generative modeling,unsupervised domain adaptation, and socially aware learning.
Most prior worksuse adversarial learning (i.e., min-max optimization), which can be challengingto 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 inefficiently 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 min-min, i.e., cooperative,problem and can provide a natural evaluation metric for unsupervised distributionalignment.
We present empirical results of our framework on both simulated andreal-world datasets to demonstrate the benefits of our approach.
Introduction
In many cases, a practitioner has access to multiple related but distinct distributions such as agriculturalmeasurements from two farms, experimental data collected in different months or sales data beforeand after a major event.
Unsupervised distribution alignment (UDA) is the ML task aimed at aligningthese related but distinct distributions in a shared space, without any pairing information betweenthe samples from these distrbibutions (i e., unsupervised).
This task has many applications such asgenerative modeling (e.g., [1]), unsupervised domain adaptation (e.g., [2, 3]), batch effect mitigationin 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 ifand only if the inner maximization is solved perfectly.
CycleGAN [1] maps between two datasetsvia two GAN objectives between the two domains and a cycle consistency loss, which encouragesapproximate invertibility of the transformations.
However, adversarial learning can be challenging to optimize in practice (see e.g. [7, 8, 9, 10, 11])in part because of the competitive nature of the min-max optimization problem.
Perhaps moreimportantly, the research community has resorted to surrogate evaluation metrics for GAN becauselikelihood computation is intractable.
Specifically, the commonly accepted Frechet Inception Distance(FID)
[12] is only applicable to image or auditory data, where publicly available classifiers trainedon large-scale data exist.
Moreover, the concrete implementation of FID can have issues due toseemingly trivial changes in image resizing algorithms [13].
Recently, flow-based methods with a tractable likelihood have been proposed for the UDA task[14, 15, 2, 16].
Specifically, iterative flow methods [14, 15] proposed alternative approaches todistribution 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 andintroduces exact log-likelihood loss terms derived from standard flow-based generative models ascomplementary loss terms to the adversarial loss terms.
On the other hand, log-likelihood ratiominimizing flows (LRMF) [16] use invertible flow models and density estimation for distributionalignment 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 areactually solving non-standard adversarial problems via alternating optimization (see [15]).
AlignFlowlacks an explicit alignment metric and assumes that the shared density model for all distributions isa fixed Gaussian.
Also, the LRMF metric is limited because it is only defined for two distributionsand strongly depends on the shared density model class; if the target distribution is not in this modelclass, LRMF can only partially align distributions.
To address these issues, we unify existing flow-based methods (both AlignFlow and LRMF) undera common cooperative (i.e., non-adversarial) framework by proving that a minimization over ashared density model is a variational upper bound of the JSD.
The unifying theory also suggests anatural 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-efficientmulti-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 ifthe 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 (including a novel domain-agnostic AUB metric) andexplain its relationship to prior flow-based alignment methods.•
Throughout experiments, we demonstrate that our framework consistently shows superiorperformance in both our proposed and existing measures on simulated and real-worlddatasets.
We also quantitatively verify that our model is more parameter-efficient than thebaseline 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 randomvariable 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 latentrandom 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 crossentropy 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 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 upperbound results in aligned distributions assuming large enough capacity of the model components.
Background: Normalizing Flows and Invertible Models. Normalizing flows are generativemodels that have tractable distributions where exact density evaluations and efficient samplingsare ensured [17].
Such models leverage the change of variables formula to create an invertiblemapping 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 .
For sampling in
𝑃 !
[Equation]
𝑇 & 𝑄 𝑄 𝑃 ’ !
𝑃 ’ "
[Figure]
distribution P X , one need to first sample from the latent distribution P Z
and then apply the inversetransform T − 1 .
Therefore the key challenge in designing invertible models is to have computationallyefficient inverse evaluation and Jacobian determinant calculation.
Many approaches have beenproposed by parameterizing mapping function T as deep neural networks including autoregressivestructures [18, 19], coupling layers [20][21], oridinary differential equations[22], and invertibleresidual networks[23, 24].
Flow models can be then learned efficiently by maximizing the likelihoodfor 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 .
5 .95 Definition 2.1 (Generalized Jensen-Shannon Divergence (GJSD) [25]) .
Given k distributions{ P X j } kj =1 and a corresponding probability weight vector w , the generalized Jensen-Shannon divergence is defined as (proof of equivalence in supplementary material.):
[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 beinvertible 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 andonly if GJSD( P Z 1 , · · · , P Z k ) = 0 .
Thus, ideally, we would minimize GJSD directly, i.e.,min [Equation] [Equation] [Equation]
P
[Equation]
where T is a class of invertible functions.
However, we cannot evaluate the entropy terms in Eqn.
1because we do not know the density of P X j ; we only have samples from P X j .
Therefore, we willupper bound the first entropy term in Eqn.
1 ( H( (cid:80) j w j P T j ( X j ) ) ) using a variational density modeland decompose the other entropy terms by leveraging the change of variables formula for invertiblefunctions.
Theorem 2.2 (GJSD Variational Upper Bound) .
Given an variational density model class Q , weform a GJSD variational upper bound:
[Equation]
where 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 KLdivergence is non-negative and the bound gap is equal to KL( P Z mix , Q ) .
The Q that achievesthe 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 withrespect 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 bemade tight if P Z mix ∈ Q .
Also, one key feature of this upper bound is that the cross entropy termcan be evaluated using only samples from P X j and the transformations T j ,
i.e., H c ( P Z mix , Q ) = 122 (cid:80)
[Equation]
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 densities of P Z j (or P X j ).
Thus, we leverage the fact that the T j functionsare 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 aninvertible 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 beignored when optimizing T , while E P X
[log | J T ( x ) | ] can be approximated using only samples from
P X (formal proof in supplementary material).
2.2 Alignment Upper Bound (AUB)
Combining Theorem 2.2 and Lemma 2.3, we can arrive at our final objective function which isequivalent 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.
We formallydefine this loss function as follows.
Definition 2.4 (Alignment Upper Bound Loss) .
Given k continuous distributions { P X j } k j =1 , a classof continuous distributions Q , and a probability weight vector w , the alignment upper bound loss isdefined as follows: [Equation]
[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 flowmodels 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 allinvertible 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 flowlosses 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 supplementary material).
This alignment of the latent distributions also implies thetranslation between any of the observed component distributions.
The proof follows directly from
Theorem 2.5 and the change of variables formula.
Corollary 2.6 (Translation at Global Minimum of L AUB ) .
Similar to Theorem 2.5, a global minimumof 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 in Supplementary material, we use a simple alternating optimization scheme fortraining our translation models and variational distribution with cooperative (i.e., min-min) AUBobjective, i.e., we aim to optimize:
[Equation]
[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., kerneldensities (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.
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 the ELBO.
First, both are variational bounds of the quantity of interest where the tightness of the bounds dependon the optimization of the variational distributions.
Second, both can be made tight if the class ofvariational distributions is powerful enough.
Third, both can be used to train a model by minimizingthe 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, webelieve AUB can be used for UDA like ELBO has been used for density estimation with the samestrengths (and weaknesses such as it being a relative metric).
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 densitymodel class Q only contains the standard normal distribution (i.e., a singleton class).
Thus, AlignFlowcan be viewed as initially optimizing a poor upper bound on JSD; however, the JSD bound becomestighter as training progresses because the latent distributions independently move towards the samenormal distribution.
By using the same architecture of T and Q , AlignFlow can be also viewed assharing the last few layers of 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 alternativelearning algorithms than SGD.
LRMF is special case with only one transformation. As illustrated in Fig.
2, Log-likelihoodratio 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 aspecial 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 absolutevalue) of the metric to be estimated but requires fitting extra density models.
[16] do not uncoverthe connection of the objective as an upper bound on JSD regardless of the density model class 1 .
Additionally, to ensure alignment, LRMF requires that the density model class includes the true targetdistribution because only one invertible transform is used, while our approach can theoretically aligneven if the shared density model class is weak (see Theorem 2.5 and our simulated experiments). Tobe specific, our bound holds regardless of the model class Q , whereas the LRMF-JSD discussionfurther assumes that the model class needs to be in the Q model class. Again, our theory does notmake any assumptions on the model class Q and yet still proves that the objective is a true upperbound on JSD.  Experiment
We analyze the performance of our proposed framework comparing to the relevant flow-based baselinemodels.
In subsection 4.1, we conduct our experiment in toy dataset to clearly show the benefit of
AlignFlow (MLE only)
LRMF
Ours
Fixed dist.
Invertible function
Changed dist.
Q dist.
Getting closed
Gettingclosed
Gettingclosed 𝑋 !
𝑋 " 𝑍 " 𝑍 ! 𝑇 !
𝑇 " 𝑋 !
𝑋 " 𝑍 ! 𝑋 !
𝑋 " 𝑍 ! 𝑍 " learnable 𝑄 learnable 𝑄 fixed 𝑄 𝑇 !
𝑇 !
𝑇 "
[Figure]
our method over the baseline flow-based models in a controlled environment (implementation detailsare in supplemantary material).
In subsection 4.2 and subsection 4.3, we demonstrate our superiorityover baseline flow-based models on real-world datasets including tabular data and high-dimensional
MNIST data.
We additionally conduct Domain Adaptation [28, 2] experiments for verifying ouralignment performance in the Appendix.
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.
We construct the experiment to have the task:
Translation between the two half-circled distributions:
X 1 and X 2 in the moons dataset.
In this example, we made two models, one with LRMF setupand one with our AUB setup.
As illustrated in Fig.
3, the LRMF method fails to transform between
X 1 and X 2 .
Even though Q can model well enough for T 1 ( X 1 ) , Q can only model the mean andvariance of X 2 which is obviously not informative enough.
Therefore, the LRMF fails to transformbetween two datasets.
While in the AUB setup, both T 1 ( X 1 ) and T 2 ( X 2 ) are modeled to the samedistribution 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 Qwhich means if Q fails to model one of the domain distribution in high likelihood, data alignmentcannot be achieved with good performances.
Simple Fixed Q vs. Learnable Q (AlignFlow vs. Ours).
Next we compare our method with
AlignFlow[2, 3] setup.
We construct the experiment to have the task: Transform between the tworandom patterns X 1 and X 2 from the randomly generated datasets.
Again, we made two modelswith AlignFlow and our AUB setups respectively.
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 toreach the normal distribution Q .
While in the AUB setup, the density model Q is learned to help fitthe transformed distributions T 1 ( X 1 ) and T 2 ( X 2 ) , which allows them to be aligned with each othereasier.
In conclusion, the performance of the AlignFlow model is limited by the performances of theinvertible functions.
4.2 Unsupervised Domain Alignment on Tabular Datasets
To demonstrate the benefits of our framework, we conduct two experiments on real-world tabulardatasets.
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 measuredby test AUB defined in Definition 2.4, where a lower AUB score indicates better performance (seeend of section 2.2 for discussion on using test AUB for evaluation).
We emphasize that there is nonatural metric for evaluating GAN-based alignment methods on tabular datasets.
We believe thisdemonstrates 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 methodand the baseline models.
To separate each dataset into two distributions, we choose the last inputfeature 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 accomplish distribution alignment.
Regarding baseline models, we use AlignFlowwith 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) andadversarial losses, to be fair, we additionally fit an identity-initialized Q to the final T ’s. We also usethe same T and Q functions (if exists in setup) with ours across all the baselines.
Note that the Q isfrom the same class of distributions Q with our approach.
As shown in Table [Table] our method shows better performance compared to other methods across alldatasets.
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 distributionalignment.
We also observe that ours shows better performance than LRMF, where the gap may comefrom 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 aligningdistributions of tabular data.
[Table]
[Table]
The second experiment is designed to show that our proposed idea is more efficient in terms of modelparameter compared to the baseline model.
For implementing experiments with multiple distributionsto align rather than just two, we separate each dataset into 8 domains by choosing last three featuresfrom a given dataset and divide the dataset by the three medians, e.g., (+ + +) , (+ + − ) , ..., ( − − − ) .
The results are shown in Table 2 RealNVP (5) and RealNVP (10) indicates the number of couplinglayers used in the RealNVP architecture respectively.
LRMF comparison is excluded because it isdesigned for aligning only two distributions.
From a comparison between two AlignFlow setups with RealNVP (5) and RealNVP (10), we cansee that AlignFlow can achieve better alignment with more than double amount of parameters.
However, as can be observed in a comparison between AlignFlow with RealNVP (10) and ours,our proposed method can achieve similar performance with the baseline model but with less thanhalf parameters.
These two results suggest that our model can be more parameter efficient than
AlignFlow for multi-distribution alignment.
We hypothesize that our approach can scale better to any number of distributions because our Q shares parameters across the 8 distributions and can capturethe similarities between distributions.
4.3 Unsupervised Distribution Alignment on MNIST Dataset
We perform an image translation task on MNIST dataset 2 [31] to demonstrate that our distributionalignment method could be applied to high-dimensional datasets.
We train ours and baseline modelswith the digit images of 0, 1, and 2, and compare the translated results both quantitatively andqualitatively.
Specifically, We use RealNVP invertible models for all translation maps T j , as well asthe density model Q .
Note that all methods are flow-based models and thus images translated back tooriginal domain are exactly the same, which implies exact cycle consistency.
Details for the setup ofthe experiments are available in supplementary material.
As represented in Fig.
4, both of our approaches outperform the baseline models in terms of FIDand 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 { Z j } kj =1 and our density model
Q is smaller than baselines, meaning our model has the better distribution alignment performance.
We believe this result comes from our model setup, i.e., a learnable shared density model andtransformation to a shared latent space.
Specifically, AlignFlow with a fixed standard normaldistribution as their Q obtains worse AUB because the Q is not powerful enough to model thecomplex shared space trained from the real world dataset.
On the other hand, LRMF shows the lackof 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 fortranslation without a shared latent space and the fact that the Q distribution must be able to model thetarget distribution to ensure alignment.
The better quantitative performance of our methods can be corroborated by the qualitative resultsas seen in Fig.
AlignFlow shows less stable translation results than our method especially fortranslating to digit ‘2’ from digits ‘0’ and ‘1’.
We believe this phenomenon comes from the lack ofexpressivity of their Q model.
On the other hand, our model shows stable results across all translationcases, which are also quantitatively verified via the lower FID score.
Note that we have a shared space,so different transformation functions for each domain are trained together.
This takes advantagesover the baseline in terms of sample complexity and computational complexity compared to multipleindependent flows as in AlignFlow.
Discussion
Pros and cons of our method compared to adversarial method.
Flow-based methods havedifferent benefits and limitations compared to adversarial methods for distribution alignment.
As oneclear difference, our approach enables an application-agnostic yet theoretically grounded evaluationmetric for comparing between models.
Additionally, our min-min problem is fundamentally differentthan a min-max problem and avoids issues unique to min-max problems.
However, our approach [Figure] certainly has other limitations compared to GANs (e.g., it is restricted to invertible models andrequires a density estimator).
Overall, we ultimately want to convey our alignment approach is afeasible and fundamentally different alternative to adversarial, which is currently the only dominantapproach.
With our foundation, future work could focus on the performance aspects just as adversariallearning has been improving over the past seven years but is still an active area of research.
Benefit of AUB metric.
Theoretically grounded dataset-agnostic metrics do not exist for GANbased models [32, 33].
This is indeed an unsolved problem for distribution alignment via adversarialmethods.
Our AUB score provides a theoretically grounded metric that is dataset and domainagnostic and does not require any pretrained classifier.
Specifically, it is an upper bound of GJSDbetween latent distributions up to a constant, which only depends on the training dataset—this isdirectly analogous to the fact that the negative log-likelihood is only equal to KL divergence upto a constant.
Thus, while the optimal AUB is not zero, the relative performance of competingmethods can be assessed, and held-out test AUB can be compared to train AUB to determine themodel’s generalization ability (even during training as seen in Fig. 5).
We emphasize that our AUBscore enables relative evaluation on tabular datasets, which is not usually considered because FID isinapplicable.
Extensibility.
As mentioned in section 1 and section 2, our proposed idea is a general frameworkthat can be harmonized with any invertible flow models , (e.g., Residual Flows [23], Flow++ [34]) andany density models (e.g., PixelCNN++ [26], FFJORD [22]).
Hence, we argue our proposed idea isnot limited to specific flow-based model or density model.
We expect our framework would showbetter performance in domain alignment if it is combined with the aforementioned state-of-the-artcomponents.
Alternating minimization against vanishing gradient problem.
As revealed in the previousstudies[35, 16], vanishing gradient problem can occur while minimizing JSD. On the same line, AUBcan 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 thetrue JSD. Thus, in experiments, we have found that alternating minimization does not suffer fromvanishing 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.
Conclusion
In this paper, we propose a novel variational upper bound on the generalized JSD that leads to atheoretically grounded alignment loss.
We then show that this framework unifies previous flow-basedapproaches to distribution alignment and demonstrate the benefits of our approach compared toprior flow-based methods.
In particular, our framework allows a straightforward extension to multidistribution alignment that we show could be more parameter efficient than naively extending priorapproaches.
More broadly, we suggest that our AUB metric can be useful as a domain-agnostic metricfor comparing distribution alignment methods (analogously to how ELBO is used to evaluate densityestimation 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 distributionalignment.
We hope this paper provides one step in that direction.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[No] We fixed random seed throughout our experiments.
(d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[Yes]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]