Unsupervised Learning of Global Factors in Deep Generative Models

1 INTRODUCTION . Since its first proposal by Kingma & Welling ( 2013 ) , Variational Autoencoders ( VAEs ) have evolved into a vast amount of variants . To name some representative examples , we can include VAEs with latent mixture models priors ( Dilokthanakul et al . ( 2016 ) ) , adapted to model time-series ( Chung et al . ( 2015 ) ) , trained via deep hierarchical variational families ( Ranganath et al . ( 2016 ) , Tomczak & Welling ( 2018 ) ) , or that naturally handle heterogeneous data types and missing data ( Nazabal et al . ( 2020 ) ) . The large majority of VAE-like models are designed over the assumption that data is i.i.d. , which remains a valid strategy for simplifying the learning and inference processes in generative models with latent variables . A different modelling approach may drop the i.i.d . assumption with the goal of capturing a higher level of dependence between samples . Inferring such kind of higher level dependencies can directly improve current approaches to find interpretable disentangled generative models ( Bouchacourt et al . ( 2018 ) ) , to perform domain alignment ( Heinze-Deml & Meinshausen ( 2017 ) ) or to ensure fairness and unbiased data ( Barocas et al . ( 2017 ) ) . The main contribution of this paper is to show that a deep probabilistic VAE non i.i.d . model with both local and global latent variable can capture meaningful and interpretable correlation among data points in a completely unsupervised fashion . Namely , weak supervision to group the data samples is not required . In the following we refer to our model as Unsupervised Global VAE ( UG-VAE ) . We combine a clustering inducing mixture model prior in the local space , that helps to separate the fundamental data features that an i.i.d . VAE would separate , with a global latent variable that modulates the properties of such latent clusters depending on the observed samples , capturing fundamental and interpretable data features . We demonstrate such a result using both CelebA , MNIST and the 3D FACES dataset in Paysan et al . ( 2009 ) . Furthermore , we show that the global latent space can explain common features in samples coming from two different databases without requiring any domain label for each sample , establishing a probabilistic unsupervised framework for domain alignment . Up to our knowledge , UG-VAE is the first VAE model in the literature that performs unsupervised domain alignment using global latent variables . Finally , we demonstrate that , even when the model parameters have been trained using an unsupervised approach , the global latent space in UG-VAE can discriminate groups of samples with non-trivial structures , separating groups of people with black and blond hair in CelebA or series of numbers in MNIST . In other words , if weak supervision is applied at test time , the posterior distribution of the global latent variable provides with an informative representation of the user defined groups of correlated data . 2 RELATED WORK . Non i.i.d . deep generative models are getting recent attention but the literature is still scarse . First we find VAE models that implement non-parametric priors : in Gyawali et al . ( 2019 ) the authors make use of a global latent variable that induces a non-parametric Beta process prior , and more efficient variational mechanism for this kind of IBP prior are introduced in Xu et al . ( 2019 ) . Second , both Tang et al . ( 2019 ) and Korshunova et al . ( 2018 ) proposed non i.i.d . exchangable models by including correlation information between datapoints via an undirected graph . Finally , some other works rely on simpler generative models ( compared to these previous approaches ) , including global variables with fixed-complexity priors , typically a multi-variate Gaussian distribution , that aim at modelling the correlation between user-specified groups of correlated samples ( e.g . images of the same class in MNIST , or faces of the same person ) . In Bouchacourt et al . ( 2018 ) or Hosoya ( 2019 ) , authors apply weak supervision by grouping image samples by identity , and include in the probabilistic model a global latent variable for each of these groups , along with a local latent variable that models the distribution for each individual sample . Below we specify the two most relevant lines of research , in relation to our work . VAEs with mixture priors . Several previous works have demonstrated that incorporating a mixture in the latent space leads to learn significantly better models . In Johnson et al . ( 2016 ) authors introduce a latent GMM prior with nonlinear observations , where the means are learned and remain invariant with the data . The GMVAE proposal by Dilokthanakul et al . ( 2016 ) aims at incorporating unsupervised clustering in deep generative models for increasing interpretability . In the VAMP VAE model Tomczak & Welling ( 2018 ) , the authors define the prior as a mixture with components given by approximated variational posteriors , that are conditioned on learnable pseudo-inputs . This approach leads to an improved performance , avoiding typical local optima difficulties that might be related to irrelevant latent dimensions . Semi-supervised deep models for grouped data . In contrast to the i.i.d . vanilla VAE model in Figure 1 ( a ) , and its augmented version for unsupervised clustering , GMVAE , in Figure 1 ( b ) , the graphical model of the Multi-Level Variational Autoencoder ( ML-VAE ) in Bouchacourt et al . ( 2018 ) is shown in Figure 1 ( c ) , where G denotes the number of groups . ML-VAE includes a local Gaussian variable Si that encodes style-related information for each sample , and global Gaussian variable CG to model shared in a group of samples . For instance , they feed their algorithm with batches of face images from the same person , modeling content shared within the group that characterize a person . This approach leads to learning a disentangled representations at the group and observations level , in a content-style fashion . Nevertheless , the groups are user-specified , hence resulting in a semisupervised modelling approach . In Vowels et al . ( 2020 ) authors use weak supervision for pairing samples . They implement two outer VAEs with shared weights for the reconstruction , and a Nested VAE that reconstructs latent representation off one to another , modelling correlations across pairs of samples . The graphical model for Nested VAE is depicted in Figure 1 ( d ) . 3 UNSUPERVISED GLOBAL VAE . We present UG-VAE , a deep generative VAE framework for modeling non-i.i.d . data with global dependencies . It generalizes the ML-VAE graphical model in Figure 1 ( c ) to i ) remove the group supervision , ii ) include a clustering-inducing prior in the local space , and iii ) propose a more structured variational family . 3.1 GENERATIVE MODEL . Figure 2 represents the generative graphical model of UG-VAE . A global variable β ∈ Rg induces shared features to generate a group of B samples X = { x1 , ... , xB } ⊆ RD , and G is the number of groups we jointly use to amortize the learning of the model parameters . During amortized variational training , groups are simply random data mini-batches from the training dataset , being G the number of data mini-batches . We could certainly take B = N ( the training set size ) and hence G = 1 , but this leads to less interpretable global latent space ( too much data to correlate with a single global random variable ) , and a slow training process . Conditioned to β , data samples are independent and distributed according to a Gaussian mixture local ( one per data ) latent variable Z = { z1 , ... , zB } ⊆ Rd , and d = { d1 , ... , dB } ⊆ { 1 , ... , K } are independent discrete categorical variables with uniform prior distributions . This prior , along with the conditional distribution p ( zi|di , β ) , defines a Gaussian mixture latent space , which helps to infer similarities between samples from different batches ( by assigning them to the same cluster ) , and thus , di plays a similar role than the semi-supervision included in Bouchacourt et al . ( 2018 ) by grouping . Our experimental results demonstrate that this level of structure in the local space is crucial to acquire interpretable information at the global space , and specially , if we fix di for all the samples within a batch , that the global variable β is able to tune different generative factors for each cluster . The joint distribution for a single group is therefore defined by : pθ ( X , Z , d , β ) = p ( X|Z , β ) p ( Z|d , β ) p ( d ) p ( β ) ( 1 ) where the likelihood term of each sample is a Gaussian distribution , whose parameters are obtained from a concatenation of zi and β as input of a decoder network : p ( X|Z , β ) = B∏ i=1 p ( xi|zi , β ) = B∏ i=1 N ( µθx ( [ zi , β ] ) , Σθx ( [ zi , β ] ) ) ( 2 ) In contrast with Johnson et al . ( 2016 ) , where the parameters of the clusters are learned but shared by all the observations , in UG-VAE , the parameters of each component are obtained with networks fed with β . Thus , the prior of each local latent continuous variable is defined by a mixture of Gaussians , where di defines the component and β is the input of a NN that outputs its parameters : p ( Z|d , β ) = B∏ i=1 p ( zi|di , β ) = B∏ i=1 N ( µ ( di ) θz ( β ) , Σ ( di ) θz ( β ) ) , ( 3 ) hence we trained as many NNs as discrete categories . This local space encodes samples in representative clusters to model local factors of variation . The prior of the discrete latent variable is defined as uniform : p ( d ) = B∏ i=1 Cat ( π ) πk = 1/K ( 4 ) and the prior over the continuous latent variable β follows an isotropic Gaussian , p ( β ) = N ( 0 , I ) . 3.2 INFERENCE MODEL . The graphical model of the proposed variational family is shown in Figure 2 ( b ) : qφ ( Z , d , β|X ) = q ( Z|X ) q ( d|Z ) q ( β|X , Z ) ( 5 ) where we employ an encoder network that maps the input data into the local latent posterior distribution , which is defined as a Gaussian : q ( Z|X ) = B∏ i=1 q ( zi|xi ) = B∏ i=1 N ( µφz ( xi ) , Σφz ( xi ) ) ( 6 ) Given the posterior distribution of z , the categorical posterior distribution of di is parametrized by a NN that takes zi as input q ( d|Z ) = B∏ i=1 q ( di|zi ) = B∏ i=1 Cat ( πφd ( zi ) ) ( 7 ) The approximate posterior distribution of the global variable β is computed as a product of local contributions per datapoint . This strategy , as demonstrated by Bouchacourt et al . ( 2018 ) , outperforms other approaches like , for example , a mixture of local contributions , as it allows to accumulate group evidence . For each sample , a NN encodes xi and the Categorical parameters πφd ( zi ) in a local Gaussian : q ( β|X , Z ) = N ( µβ , Σβ ) = B∏ i=1 N ( µφβ ( [ xi , πφd ( zi ) ] ) , Σφβ ( [ xi , πφd ( zi ) ] ) ) ( 8 ) If we denote by µi and Σi the parameters obtained by networks µφβ and Σφβ , respectively , the parameters of the global Gaussian distribution are given , following Bromiley ( 2003 ) , by : Λβ = Σ −1 β = B∑ i=1 Λi µβ = ( Λβ ) −1 B∑ i=1 Λiµi ( 9 ) where Λβ = Σ−1β is defined as the precision matrix , which we model as a diagonal matrix .

Summary: This article introduces a VAE-based method for separating local variation factors from global variation factors in the data in an unsupervised manner. It achieves so by designing a graphical model with a mix of example-local and batch-shared variables, and training it using the ELBO. The article provide an detailed experimental analysis on MNIST and CelebA, and experimental evidence that all parts of the model (notably the discrete d variable) are relevant.