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