FAVAE: SEQUENCE DISENTANGLEMENT USING IN- FORMATION BOTTLENECK PRINCIPLE

1 INTRODUCTION . Representation learning is one of the most fundamental problems in machine learning . A real world data distribution can be regarded as a low-dimensional manifold in a high-dimensional space ( Bengio et al. , 2013 ) . Generative models in deep learning , such as the variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) and the generative adversarial network ( GAN ) ( Goodfellow et al. , 2014 ) , are able to learn low-dimensional manifold representation ( factor ) as a latent variable . The factors are fundamental components such as position , color , and degree of smiling in an image of a human face ( Liu et al. , 2015 ) . Disentangled representation is defined as a single factor being represented by a single latent variable ( Bengio et al. , 2013 ) . Thus , if in a model of learned disentangled representation , shifting one latent variable while leaving the others fixed generates data showing that only the corresponding factor was changed . This is called latent traversals ( a good demonstration of which was given by Higgins et al . ( 2016a ) 1 ) . There are two advantages of disentangled representation . First , latent variables are interpretable . Second , the disentangled representation is generalizable and robust against adversarial attacks ( Alemi et al. , 2016 ) . We focus on the disentangled representation learning of sequential data . Sequential data is characterized by dynamic factors and static factors : dynamic factors are time dependent , and static factors are independent of time . With disentangled representation learning from sequential data , we should be able to extract dynamic factors that can not be extracted by disentangled representation learning models for non-sequential data such as β-VAE ( Higgins et al. , 2016a ; b ) and InfoGAN ( Chen et al. , 2016 ) . The concept of disentangled representation learning for sequential data is illustrated in Fig . 1 . Consider that the pseudo-dataset of the movement of a submarine has a dynamic factor : the trajectory shape . The disentangled representation learning model for sequential data can extract this shape . On the other hand , since the disentangled representation learning model for non-sequential data does not consider the sequence of data , it merely extracts the x-position and y-position . 1This demonstration is available at http : //tinyurl.com/jgbyzke There is a wide range of potential applications if we extend disentanglement representation to sequential data such as speech , video , and stock market data . For example , disentangled representation learning for stock price data can extract the fundamental trend of a given stock price . Another application is the reduction of action space in reinforcement learning . Extracting dynamic factors would enable the generation of macro-actions ( Durugkar et al. , 2016 ) , which are sets of sequential actions that represent the fundamental factors of the actions . Thus , disentangled representation learning for sequential data opens the door to new areas of research . Very recent related work ( Hsu et al. , 2017 ; Li & Mandt , 2018 ) separated factors of sequential data into dynamic and static factors . The factorized hierarchical variational autoencoder ( FHVAE ) ( Hsu et al. , 2017 ) is based on a graphical model using latent variables with different time dependencies . By maximizing the variational lower bound of the graphical model , the FHVAE separates the different time dependent factors such as the dynamic and static factors . The VAE architecture developed by Li & Mandt ( 2018 ) is the same as the FHVAE in terms of the time dependencies of the latent variables . Since these models require different time dependencies for the latent variables , these approaches can not be used disentangle variables with the same time dependency factor . We address this problem by taking a different approach . First , we analyze the root cause of disentanglement from the perspective of information theory . As a result , the term causing disentanglement is derived from a more fundamental rule : reduce the mutual dependence between the input and output of an encoder while keeping the reconstruction of the data . This is called the information bottleneck ( IB ) principle . We naturally extend this principle to sequential data from the relationship between x and z to xt : T and z . This enables the separation of multiple dynamic factors as a consequence of information compression . It is difficult to learn a disentangled representation of sequential data since not only the feature space but also the time space should be compressed . We created the factorized action variational autoencoder ( FAVAE ) in which we implemented the concept of information capacity to stabilize learning and a ladder network to learn a disentangled representation in accordance with the level of data abstraction . Since our model is a more general model without the restriction of a graphical model design to distinguish between static and dynamic factors , it can separate depen- dency factors occurring at the same time . Moreover , it can separate factors into dynamic and static factors . 2 DISENTANGLEMENT FOR NON-SEQUENTIAL DATA . β-VAE ( Higgins et al. , 2016a ; b ) is a commonly used method for learning disentangled representations based on the VAE framework ( Kingma & Welling , 2013 ) for a generative model . The VAE can estimate the probability density from data x . The objective function of the VAE maximizes the evidence lower bound ( ELBO ) of log p ( x ) as log p ( x ) = LVAE +DKL ( q ( z|x ) ||p ( z|x ) ) ︸ ︷︷ ︸ ≥0 , ( 1 ) where z is latent variable , DKL is the Kullback-Leibler divergence , and q ( z|x ) is an approximated distribution of p ( z|x ) . DKL ( q ( z|x ) ||p ( z|x ) ) reduces to zero as the ELBO LVAE increases ; thus , q ( z|x ) learns a good approximation of p ( z|x ) . The ELBO is defined as LVAE ≡ Eq ( z|x ) [ log p ( x|z ) ] −DKL ( q ( z|x ) ||p ( z ) ) , ( 2 ) where the first term , Eq ( z|x ) [ log p ( x|z ) ] , is a reconstruction term used to reconstruct x , and the second term DKL ( q ( z|x ) ||p ( z ) ) is a regularization term used to regularize posterior q ( z|x ) . Encoder q ( z|x ) and decoder p ( x|z ) are learned in the VAE . Next we will explain how β-VAE extracts disentangled representations from unlabeled data . β-VAE is an extension of the coefficient β > 1 of the regularization term DKL ( q ( z|x ) ||p ( z ) ) in the VAE . The objective function of β-VAE is Lβ−VAE = Eq ( z|x ) [ log p ( x|z ) ] − βDKL ( q ( z|x ) ||p ( z ) ) , ( 3 ) where β > 1 and p ( z ) = N ( 0 , 1 ) . β-VAE promotes disentangled representation learning via the Kullback-Leibler divergence term . As β increases , the latent variable q ( z|x ) approaches the prior p ( z ) ; therefore , each zi is pressured to learn the probability distribution of N ( 0 , 1 ) . However , if all latent variables zi become N ( 0 , 1 ) , the model can not reconstruct x . As a result , as long as z reconstructs x , β-VAE reduces the information of z . 3 PRELIMINARY : ORIGIN OF DISENTANGLEMENT . To clarify the origin of disentanglement , we will explain the regularization term . The regularization term has been decomposed into three terms ( Chen et al. , 2018 ; Kim & Mnih , 2018 ; Hoffman & Johnson , 2016 ) : Ep ( x ) [ DKL ( q ( z|x ) ||p ( z ) ) ] = I ( x ; z ) +DKL q ( z ) ||∏ j q ( zj ) +∑ j DKL ( q ( zj ) ||p ( zj ) ) , ( 4 ) where zj denotes the j-th dimension of the latent variable . The second term , which is called ” total correlation ” in information theory , quantifies the redundancy or dependency among a set of n random variables ( Watanabe , 1960 ) . β-TCVAE has been experimentally shown to reduce the total correlation causing disentanglement ( Chen et al. , 2018 ) . The third term indirectly causes disentanglement by bringing q ( z|x ) close to the independent standard normal distribution p ( z ) . The first term is mutual information between the data variable and latent variable based on the empirical data distribution . Minimizing the regularization term causes disentanglement but disturbs reconstruction via the first term in Eq . ( 4 ) . The shift C scheme was proposed ( Burgess et al. , 2018 ) as a means to solve this conflict : −Eqφ ( z|x ) [ log pθ ( x|z ) ] + β |DKL ( q ( z|x ) ||p ( z ) ) − C| , ( 5 ) where constant shift C , which is called ” information capacity , ” linearly increases during training . This shift C can be understood from the point of view of an information bottleneck ( Tishby et al. , 2000 ) . The VAE can be derived by maximizing the ELBO , but β-VAE can no longer be interpreted as an ELBO once this scheme has been applied . The objective function of β-VAE is derived from the information bottleneck ( Alemi et al. , 2016 ; Achille & Soatto , 2018 ; Tishby et al. , 2000 ; Chechik et al. , 2005 ) . max I ( z ; x ) s.t . |I ( x̂ ; z ) − C| = 0 , ( 6 ) where C is the information capacity and x̂ is the empirical distribution . Solving this equation by using Lagrange multipliers drives the objective function of β-VAE Eq . ( 5 ) with β as the Lagrange multiplier ( details in Appendix B of ( Alemi et al. , 2016 ) ) . In Eq . ( 5 ) , information capacity C prevents I ( x̂ , z ) from becoming zero . In the information bottleneck literature , y typically stands for a classification task ; however , the formulation can be related to the autoencoding objective ( Alemi et al. , 2016 ) . Therefore , the objective function of β-VAE can be understood using the information bottleneck principle . 4 PROPOSED MODEL : DISENTANGLEMENT FOR SEQUENTIAL DATA . Our proposed FAVAE model learns disentangled and interpretable representations from sequential data without supervision . We consider sequential data x1 : T ≡ { x1 , x2 , · · · , xT } generated from a latent variable model , p ( x1 : T ) = ∫ p ( x1 : T |z ) p ( z ) dz . ( 7 ) For sequential data , we replace x with ( x1 : T ) in Eq . 5 . The objective function of the FAVAE model is −Eqφ ( z| ( x1 : T ) i ) [ log pθ ( x1 : T |z ) ] + β |DKL ( q ( z| ( x1 : T ) i ) ||p ( z ) ) − C| , ( 8 ) where p ( z ) = N ( 0 , 1 ) . The variational recurrent neural network ( Chung et al. , 2015 ) and stochastic recurrent neural network ( Fraccaro et al. , 2016 ) extend the VAE model to a recurrent framework . The priors of both networks are dependent on time . The time dependent prior experimentally improves the ELBO . In contrast , the prior of our model is independent of time like those of the stochastic recurrent network ( Bayer & Osendorfer , 2014 ) and the Deep Recurrent Attentive Writer ( DRAW ) neural network architecture ( Gregor et al. , 2015 ) ; this is because FAVAE is disentangled representation learning rather than density estimation . For better understanding , consider FAVAE from the perspective of IB . As with β-VAE , FAVAE can be understood from the information bottleneck principle . max I ( z ; x1 : T ) s.t . |I ( x̂1 : T ; z ) − C| = 0 , ( 9 ) where x̂1 : T follows an empirical distribution . These principles make the representation of z compact while reconstruction of the sequential data is represented by x1 : T ( see Appendix A ) . 4.1 LADDER NETWORK An important extension to FAVAE is a hierarchical representation scheme inspired by the VLAE ( Zhao et al. , 2017 ) . Encoder q ( z|x1 : T ) within a ladder network is defined as hl = fl ( hl−1 ) , ( 10 ) zl ∼ N ( µl ( hl ) , σl ( hl ) ) , ( 11 ) where l is a layer index , h0 ≡ x1 : T , and f is a time convolution network , which is explained in the next section . Decoder p ( x1 : T |z ) within the ladder network is defined as z̃L = gl ( zL ) ( 12 ) z̃l = gl ( z̃l+1 + gate ( zl ) ) , ( 13 ) x1 : T ∼ r ( x1 : T , z̃0 ) ( 14 ) where gl is the time deconvolution network with l = 1 , · · · , L − 1 , and r is a distribution family parameterized by g0 ( z̃0 ) . The gate computes the Hadamard product of its learnable parameter and ( a ) 2D reaching dataset . ( b ) β-VAE , 1st z . ( c ) β-VAE , 2nd z . ( d ) FAVAE , 1st z in 3rd ladder . ( e ) FAVAE , 1st z in 1st ladder . Figure 3 : Visualization of latent traversal of β-VAE and FAVAE . On one sampled trajectory ( red ) , each latent variable is traversed and purple and/or blue points are generated . The color corresponds to the value of the traversed latent variable . 3a represents all data trajectories of 2D reaching . input tensor . We set r as a fixed-variance factored Gaussian distribution with the mean given by µt : T = g0 ( z̃0 ) . Fig . ( 2 ) shows the architecture of the proposed model . The difference between each ladder network in the model is the number of convolution networks through which data passes . The abstract expressions should differ between ladders since the time convolution layer abstracts sequential data . Without the ladder network , the proposed method can disentangle only the representations at the same level of abstraction ; with the ladder network , it can disentangle representations at different levels of abstraction .

Summary: The paper proposes the factorized action variational autoencoder (FAVAE) as a new model for learning disentangled representations in high-dimensional sequences, such as videos. In contrast to earlier work that encouraged disentanglement through a carefully-designed dynamical prior, the authors propose a different encoder-decoder architecture combined with a modified loss function (Beta-VAE with a “shift C scheme”). As a result, the authors claim that their approach leads to useful disentangled representation learning in toy video data, and in data taken from a robotic task.