Why be adversarial? Let's cooperate!: Cooperative Dataset Alignment via JSD Upper Bound

1 INTRODUCTION . 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 ) . 2 REGULARIZED 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 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 ) 2.1 GJSD UPPER BOUND . 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 can not 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 can not 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 .

Summary: 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.