Identifying through Flows for Recovering Latent Representations

Identifiability , or recovery of the true latent representations from which the observed data originates , is de facto a fundamental goal of representation learning . Yet , most deep generative models do not address the question of identifiability , and thus fail to deliver on the promise of the recovery of the true latent sources that generate the observations . Recent work proposed identifiable generative modelling using variational autoencoders ( iVAE ) with a theory of identifiability . Due to the intractablity of KL divergence between variational approximate posterior and the true posterior , however , iVAE has to maximize the evidence lower bound ( ELBO ) of the marginal likelihood , leading to suboptimal solutions in both theory and practice . In contrast , we propose an identifiable framework for estimating latent representations using a flow-based model ( iFlow ) . Our approach directly maximizes the marginal likelihood , allowing for theoretical guarantees on identifiability , thereby dispensing with variational approximations . We derive its optimization objective in analytical form , making it possible to train iFlow in an end-to-end manner . Simulations on synthetic data validate the correctness and effectiveness of our proposed method and demonstrate its practical advantages over other existing methods . 1 INTRODUCTION . A fundamental question in representation learning relates to identifiability : under which condition is it possible to recover the true latent representations that generate the observed data ? Most existing likelihood-based approaches for deep generative modelling , such as Variational Autoencoders ( VAE ) ( Kingma & Welling , 2013 ) and flow-based models ( Kobyzev et al. , 2019 ) , focus on performing latent-variable inference and efficient data synthesis , but do not address the question of identifiability , i.e . recovering the true latent representations . The question of identifiability is closely related to the goal of learning disentangled representations ( Bengio et al. , 2013 ) . While there is no canonical definition for this term , we adopt the one where individual latent units are sensitive to changes in single generative factors while being relatively invariant to nuisance factors ( Bengio et al. , 2013 ) . A good representation for human faces , for example , should encompass different latent factors that separately encode different attributes including gender , hair color , facial expression , etc . By aiming to recover the true latent representation , identifiable models also allow for principled disentanglement ; this suggests that rather than being entangled in disentanglement learning in a completely unsupervised manner , we go a step further towards identifiability , since existing literature on disentangled representation learning , such as β-VAE ( Higgins et al. , 2017 ) , β-TCVAE ( Chen et al. , 2018 ) , DIP-VAE ( Kumar et al. , 2017 ) and FactorVAE ( Kim & Mnih , 2018 ) , are neither general endeavors to achieve identifiability ; nor do they provide theoretical guarantees on recovering the true latent sources . Recently , Khemakhem et al . ( 2019 ) introduced a theory of identifiability for deep generative models , based upon which the authors proposed an identifiable variant of VAEs called iVAE , to learn the distribution over latent variables in an identifiable manner . However , the downside of learning such an identifiable model within the VAE framework lies in the intractability of KL divergence between the approximate posterior and the true posterior . Consequently , in both theory and practice , iVAE inevitably leads to a suboptimal solution , which renders the learned model far less identifiable . In this paper , aiming at avoiding such a pitfall , we propose to learn an identifiable generative model through flows ( short for normalizing flows ( Tabak et al. , 2010 ; Rezende & Mohamed , 2015 ) ) . A normalizing flow is a transformation of a simple probability distribution ( e.g . a standard normal ) into a more complex probability distribution by a composition of a series of invertible and differentiable mappings ( Kobyzev et al. , 2019 ) . Hence , they can be exploited to effectively model complex probability distributions . In contrast to VAEs relying on variational approximations , flow-based models allow for latent-variable inference and likelihood evaluation in an exact and efficient manner , making themselves a perfect choice for achieving identifiability . To this end , unifying identifiablity with flows , we propose iFlow , a framework for deep latentvariable models which allows for recovery of the true latent representations from which the observed data originates . We demonstrate that our flow-based model makes it possible to directly maximize the conditional marginal likelihood and thus achieves identifiability in a rigorous manner . We provide theoretical guarantees on the recovery of the true latent representations , and show experiments on synthetic data to validate the theoretical and practical advantages of our proposed formulation over prior approaches . 2 BACKGROUND . An enduring demand in statistical machine learning is to develop probabilistic models that explain the generative process that produce observations . Oftentimes , this entails estimation of density that can be arbitrarily complex . As one of the promising tools , Normalizing Flows are a family of generative models that fit such a density of exquisite complexity by pushing an initial density ( base distribution ) through a series of transformations . Formally , let x ∈ X ⊆ Rn be an observed random variable , and z ∈ Z ⊆ Rn a latent variable obeying a base distribution pZ ( z ) . A normalizing flow f is a diffeomorphism ( i.e an invertible differentiable transformation with differentiable inverse ) between two topologically equivalent spaces X and Z such that x = f ( z ) . Under these conditions , the density of x is well-defined and can be obtained by using the change of variable formula : pX ( x ) = pZ ( h ( x ) ) ∣∣∣∣det ( ∂h∂x ) ∣∣∣∣ = pZ ( z ) ∣∣∣∣det ( ∂f∂z ) ∣∣∣∣−1 , ( 1 ) where h is the inverse of f . To approximate an arbitrarily complex nonlinear invertible bijection , one can compose a series of such functions , since the composition of invertible functions is also invertible , and its Jacobian determinant is the product of the individual functions ’ Jacobian determinants . Denote φ as the diffeomorphism ’ s learnable parameters . Optimization can proceed as follows by maximizing log-likelihood for the density estimation model : φ∗ = arg max φ Ex [ log pZ ( h ( x ; φ ) ) + log ∣∣∣∣det ( ∂h ( x ; φ ) ∂x ) ∣∣∣∣ ] . ( 2 ) 3 RELATED WORK . Nonlinear ICA Nonlinear independent component analysis ( ICA ) is one of the biggest problems remaining unresolved in unsupervised learning . Given the observations alone , it aims to recover the inverse mixing function as well as their corresponding independent sources . In contrast with the linear case , research on nonlinear ICA is hampered by the fact that without auxiliary variables , recovering the independent latents is impossible ( Hyvärinen & Pajunen , 1999 ) . Similar impossibility result can be found in ( Locatello et al. , 2018 ) . Fortunately , by exploiting additional temporal structure on the sources , recent work ( Hyvarinen & Morioka , 2016 ; Hyvarinen et al. , 2018 ) established the first identifiability results for deep latent-variable models . These approaches , however , do not explicitly learn the data distribution ; nor are they capable of generating “ fake ” data . Khemakhem et al . ( 2019 ) bridged this gap by establishing a principled connection between VAEs and an identifiable model for nonlinear ICA . Their method with an identifiable VAE ( known as iVAE ) approximates the true joint distribution over observed and latent variables under mild conditions . Nevertheless , due to the intractablity of KL divergence between variational approximate posterior and the true posterior , iVAE maximizes the evidence lower bound on the data log-likelihood , which in both theory and practice acts as a detriment to the achievement of identifiability . We instead propose identifying through flows ( normalizing flow ) , which maximizes the likelihood in a straightforward way , providing theoretical guarantees and practical advantages for identifiability . Normalizing Flows Normalizing Flows are a family of generative approaches that fits a data distribution by learning a bijection from observations to latent codes , and vice versa . Compared with VAEs which learn a posterior approximation to the true posterior , normalizing flows directly deal with marginal likelihood with exact inference while maintaining efficient sampling . Formally , a normalizing flow is a transform of a tractable probability distribution into a complex distribution by compositing a sequence of invertible and differentiable mappings . In practice , the challenge lies in designing a normalizing flow that satisfies the following conditions : ( 1 ) it should be bijective and thus invertible ; ( 2 ) it is efficient to compute its inverse and its Jacobian determinant while ensuring sufficient capabilities . The framework of normalizing flows was first defined in ( Tabak et al. , 2010 ) and ( Tabak & Turner , 2013 ) and then explored for density estimation in ( Rippel & Adams , 2013 ) . Rezende & Mohamed ( 2015 ) applied normalizing flows to variational inference by introducing planar and radial flows . Since then , there had been abundant literature towards expanding this family . Kingma & Dhariwal ( 2018 ) parameterizes linear flows with the LU factorization and “ 1 × 1 ” convolutions for the sake of efficient determinant calculation and invertibility of convolution operations . Despite their limits in expressive capabilities , linear flows act as essential building blocks of affine coupling flows as in ( Dinh et al. , 2014 ; 2016 ) . Kingma et al . ( 2016 ) applied autoregressive models as a form of normalizing flows , which exhibit strong expressiveness in modelling statistical dependencies among variables . However , the forwarding operation of autoregressive models is inherently sequential , which makes it inefficient for training . Splines have also been used as building blocks of normalizing flows : Müller et al . ( 2018 ) suggested modelling a linear and quadratic spline as the integral of a univariate monotonic function for flow construction . Durkan et al . ( 2019a ) proposed a natural extension to the framework of neural importance sampling and also suggested modelling a coupling layer as a monotonic rational-quadratic spine ( Durkan et al. , 2019b ) , which can be implemented either with a coupling architecture RQ-NSF ( C ) or with autoregressive architecture RQ-NSF ( AR ) . The expressive capabilities of normalizing flows and their theoretical guarantee of invertibility make them a natural choice for recovering the true mixing mapping from sources to observations , and thus identifiability can be rigorously achieved . In our work , we show that by aligning normalizing flows with an existing identifiability theory , it is desirable to learn an identifiable latent-variable model with theoretical guarantees of identifiability . 4 IDENTIFIABLE FLOW . In this section , we first introduce the identifiable latent-variable family and the theory of identifiability ( Khemakhem et al. , 2019 ) that makes it possible to recover the joint distribution between observations and latent variables . Then we derive our model , iFlow , and its optimization objective which admits principled disentanglement with theoretical guarantees of identifiability . 4.1 IDENTIFIABLE LATENT-VARIABLE FAMILY . The primary assumption leading to identifiability is a conditionally factorized prior distribution over the latent variables , pθ ( z|u ) , where u is an auxiliary variable , which can be the time index in a time series , categorical label , or an additionally observed variable ( Khemakhem et al. , 2019 ) . Formally , let x ∈ X ⊆ Rn and u ∈ U ⊆ Rm be two observed random variables , and z ∈ Z ⊆ Rn a latent variable that is the source of x . This implies that there can be an arbitrarily complex nonlinear mapping f : Z → X . Assuming that f is a bijection , it is desirable to recover its inverse by approximating using a family of invertible mappings hφ parameterized by φ . The statistical dependencies among these random variables are defined by a Bayesian net : u → z → x , from which the following conditional generative model can be derived : p ( x , z|u ; Θ ) = p ( x|z ; φ ) p ( z|u ; T , λ ) , ( 3 ) where p ( x|z ; φ ) def= p ( x−h−1 ( z ) ) and p ( z|u ; T , λ ) is assumed to be a factorized exponential family distribution conditioned upon u . Note that this density assumption is valid in most cases , since the exponential families have universal approximation capabilities ( Sriperumbudur et al. , 2017 ) . Specifically , the probability density function is given by pT , λ ( z|u ) = n∏ i=1 pi ( zi|u ) = ∏ i Qi ( zi ) Zi ( u ) exp [ k∑ j=1 Ti , j ( zi ) λi , j ( u ) ] , ( 4 ) where Qi is the base measure , Zi ( u ) is the normalizing constant , Ti , j ’ s are the components of the sufficient statistic and λi , j ( u ) the natural parameters , critically depending on u . Note that k indicates the maximum order of statistics under consideration .

Summary: This paper is about learning an identifiable generative model, iFlow, that builds upon a recent result on nonlinear ICA. The key idea is providing side information to identify the latent representation, i.e., essentially a prior conditioned on extra information such as labels and restricting the mapping to flows for being able to compute the likelihood. As the loglikelihood of a flow model is readily available, a direct approach can be used for learning that optimizes both the prior and the observation model.