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: The paper tackles the identifiability problem in generative modeling, i.e., recovering the true latent representations from which the observed data originates. The paper argues that identifiable variational autoencoder (iVAE) suffers from intractability issue which leads to suboptimal solutions. The paper instead proposes an identifiable normalizing flows (iFlow) method as an alternative. The proposed iFlow outperforms iVAE in experiments using synthetic data. The paper is very well motivated and well supports its claim.