Non-Synergistic Variational Autoencoders

1 INTRODUCTION . Our world is hierarchical and compositional , humans can generalise better since we use primitive concepts that allow us to create complex representations ( Higgins et al . ( 2016 ) ) . Towards the creation of truly intelligent systems , they should learn in a similar way resulting in an increase of their performance since they would capture the underlying factors of variation of the data ( Bengio et al . ( 2013 ) ; Hassabis et al . ( 2017 ) ; Botvinick et al . ( 2017 ) ) . In addition , good representations improve the performance for tasks involving transfer learning and multi-task learning ; since it will capture the explanatory factors . According to Lake et al . ( 2016 ) , a compositional representation should create new elements from the combination of primitive concepts resulting in a infinite number of new representations . For example if our model is trained with images of white wall and then is presented a boy with a white shirt , it should identify the color white as a primitive element . Intuitively , our model will be able to construct different and multiple representations from the primitives . Furthermore , a disentangled representation has been interpreted in different ways , for instance Bengio et al . ( 2013 ) define it as one where single latent variables are sensitive to changes in generative factors , while being invariant to changes in other factors . In addition , we agree with Higgins et al . ( 2017a ) , which mentions that a disentangle representation should be factorised and interpretable . Intuitevely , the model could learn generative factors such as position , scale or colour ; if it is disentangle it should be able to traverse along the position variable without changing the scale or the colour . It ’ s worth noting that disentangled representations have been useful for a variety of downstream tasks such as domain adaptation by training a Reinforcement Learning agent that uses a disentangled representation of its environment Higgins et al . ( 2017b ) ; or for learning disentangled primitives grounded in the visual domain discovered in an unsupervised manner Higgins et al . ( 2017c ) . 2 RELATED WORK . The original Variational autoencoder ( VAE ) framework ( Kingma & Welling ( 2013 ) ; Rezende et al . ( 2014 ) ) has been used extensively for the task of disentanglement by modifying the original ELBO formulation ; for instance β-VAE is presented in Higgins et al . ( 2017a ) which increases the latent capacity by penalising the KL divergence term with a β hyperparameter . In addition , Kim & Mnih ( 2018 ) achieved a more robust disentangled representation by using the model called Factor VAE which penalises the total correlation of the latent variables encouraging the independence of the latents ; a similar approach is shown in Chen et al . ( 2018 ) , where they present a clever ELBO decomposition based on Hoffman & Johnson ( 2016 ) . Other approaches rely on information bottleneck presented in Tishby et al . ( 1999 ) to model frameworks for this task such as Alemi et al . ( 2016 ) ; Achille & Soatto ( 2016 ) . Furthermore , Chen et al . ( 2016 ) describe a model based on Generative Adversarial Networks ( Goodfellow et al . ( 2014 ) ) by encouraging the mutual information between the latents and the output of the generator . Notably , in Higgins et al . ( 2016 ) , they describe comprehensively the β-VAE model using a neuroscience and information theory approaches ; they suggest that by encouraging redundancy reduction the model achieves statistical independence within the latents . This model inspired us to look into different fields for new ways to enforce disentanglement of the latents . 3 SYNERGY . To understand our model , we need first to describe Synergy ( Gat & Tishby ( 1998 ) ; Schneidman et al . ( 2003 ) ) being a popular notion of it as how much the whole is greater than the sum of its parts . It ’ s common to describe it with the XOR gate , since we need two independent variables to fully specified the value of the output . Following , we describe the synergy from two related fields . 3.1 INFORMATION THEORY APPROACH . Computing the multivariate synergistic information is an ongoing topic of research Schneidman et al . ( 2003 ) ; Williams & Beer ( 2010 ) ; Bertschinger et al . ( 2012 ) ; Griffith & Koch ( 2012 ) . Most of the current research in this topic uses the Partial information diagram described by Williams & Beer ( 2010 ) . In order to understand the importance of the Synergy information in our framework it ’ s essential to describe the relations with the Unique and Redundant Information . Introducing the notation from Williams & Beer ( 2010 ) , let ’ s consider the random variable S and a random vector R = { R1 , R2 , .. , Rn } , being our goal to decompose the information that the variable R provides about S ; the contribution of these partial information could come from one element fromR1 or from subsets of R ( ie . R1 , R2 ) . Considering the case with two variables , { R1 , R2 } , we could separate the partial mutual information in unique ( Unq ( S ; R1 ) and Unq ( S ; R2 ) ) , the information that only R1 or R2 provides about S is redundant ( Rdn ( S ; R1 , R2 ) ) , which it could be provided by R1 or R2 ; and synergistic ( Syn ( S ; R1 , R2 ) ) , which is only provided by the combination of R1 and R2 . The figure 1 depicts the decomposition ; this diagram is also called PI-diagram ( Partial information ) . Additional notation is used for a better visualisation in the case of more variables : Unique { 1 } , { 2 } ; Redundant { 1 } { 2 } and Synergistic { 12 } . For the case of 2 variables ( X1 , X2 ) , we expect four contributions to the mutual information as described in Bertschinger et al . ( 2012 ) ; Olbrich et al . ( 2015 ) : I ( S ; R1 , R2 ) = SI ( S ; R1 , R2 ) ︸ ︷︷ ︸ Redundant +Unq ( S ; R1 \R2 ) ︸ ︷︷ ︸ Unique +Unq ( S ; R2 \R1 ) ︸ ︷︷ ︸ Unique +Syn ( S ; R1 , R2 ) ︸ ︷︷ ︸ Synergistic ( 1 ) It ’ s easy to see that the number of terms increases exponentially as the number of sources increases . The best measure for synergy is an ongoing topic of research . In the subsection we are going to talk about the synergy metrics . 3.2 POPULATION CODING APPROACH . For neural codes there are three types of independence when it comes to the relation between stimuli and responses ; which are the activity independence , the conditional independence and the information independence . One of the first measures of synergy for sets of sources of information came from this notion of independence . In Williams & Beer ( 2010 ) it is stated that if the responses come from different features of the stimulus , the information encoded in those responses should be added to estimate the mutual information they provide about the stimulus . Formally : I ( S ; R1 , R2 ) = I ( S ; R1 ) + I ( S ; R2 ) ( 2 ) However , we just saw in the previous sections that the I ( S ; R1 ) and I ( S ; R2 ) could be decomposed in their unique and redundant and synergistic terms . Intuitively , this formulation only holds if there is no redundant or synergistic information present ; which means in the context of population coding that the responses encoded different parts of the stimulus . If the responses R1 and R2 convey more information together than separate , we can say we have synergistic information ; if the information is less , we have redundant information . That ’ s the reason why in Gat & Tishby ( 1998 ) , the synergy is considered as measure of information independence : Syn ( R1 , R2 ) = I ( S ; R1 , R2 ) − I ( S ; R1 ) − I ( S ; R2 ) ( 3 ) 3.3 SYNERGY METRIC . First we need to state the notation ( the same as Griffith & Koch ( 2012 ) ) : • n : Number of individual predictors Xi • Ai : subset of individual predictors ( ie . Ai = { X1 , X3 } ) • X : Joint random variable of all individual predictors X1X2 .. Xn • { X1 , X2 , ... , Xn } : Set of all the individual predictors • Y : Random variable to be predicted • y : A particular outcome of Y . The intuition behind this metric is that synergy should be defined as the ” whole beyond the maximum of its parts ” . The whole is described as the mutual information between the joint X and the outcome Y ; whereas the maximum of all the possible subsets is interpreted as the maximum information that any of the sources Ai provided about each outcome . Formally , this is stated as : Smax ( { X1 , X2 , ... , Xn } ; Y ) = I ( X ; Y ) − Imax ( { A1 , A2 .. An } ; Y ) ( 4 ) = I ( X ; Y ) − ∑ y∈Y p ( Y = y ) max i I ( Ai ; Y = y ) ( 5 ) This metric derives from Williams & Beer ( 2010 ) and Griffith & Koch ( 2012 ) , however one of the differences with this metric with the one presented in Griffith & Koch ( 2012 ) is that in this one we are considering the specific mutual information Imax in a group of latents Ai , whereas in the paper mentioned it considers only an individual latent . Notably the Imax can be expressed in terms of the KL divergence . I ( Ai ; Y = y ) = ∑ ai∈Ai P ( ai | y ) log P ( ai , y ) P ( ai ) P ( y ) ( 6 ) = DKL [ P ( Ai | y ) ‖ P ( Ai ) ] ( 7 ) Putting together the equation 7 and 5 , we have the following : Smax ( { X1 , X2 , ... , Xn } ; Y ) = I ( X ; Y ) − ∑ y∈Y p ( Y = y ) max i DKL [ P ( Ai | y ) ‖ P ( Ai ) ] ( 8 ) In the following section we are going to use the intuition provided by the above equation in the VAE framework for the task of disentanglement . 4 MODEL DERIVATION . The motivation of our contribution is inspired in this concept and driven by the belief that synergy is not desirable for the task of disentanglement , since we want the latents to be independently informative as possible about the data , instead of needing many latents to specify the data . Therefore , we argue that by penalising the synergistic information within the latents and the data , we would encourage the disentanglement of the underlying factors of variation . This hypothesis is also inspired in the information presented in Griffith & Koch ( 2012 ) ; Gat & Tishby ( 1998 ) , where it is stated that the synergy is a measure of independence information in the responses of the stimuli . First , we need to change the notation to match the VAE framework notation ( Z are the latents and X is the observations ) . Also , Ai is a subset of the latents , such that Ai ∈ { Z1 , Z2 , ... , Zn } and Z is the joint of the latents . Formally : Z = ∏d i Zi , where d is the number of dimensions of the latent variables . Besides , from the VAE standard framework , we know that the distribution p ( Ai|x ) is intractable which is why we need to use an approximate distribution qφ ( Ai|x ) parametrised by φ parameters . It ’ s important to notice that this KL divergence could be computed in the same way as in the VAE framework ; the only difference is the number of dimensions used for the random variable z . In the original VAE framework , we compute the KL divergence considering the joint Z = ∏d i Zi ; whereas for the Synergy metric we don ’ t use the joint but a subset of the latents . For instance , if Ai = Z2Z5Z8 , we have the following expression : DKL [ qφ ( z2z5z8 | x ) ‖ p ( z2z5z8 ) ] ( 9 ) Taking in account these considerations , we express the equation 8 as follows : Smax ( { Z1 , Z2 , ... , Zd } ; X ) = I ( Z ; X ) − ∑ x∈X p ( X = x ) max i DKL [ qφ ( Ai | x ) ‖ p ( Ai ) ] ( 10 ) We start with the original ELBO formulation ( Rezende et al . ( 2014 ) ; Kingma & Welling ( 2013 ) ) and add the penalised term corresponding to the synergy , where α is a hyperparameter : Lelbo ( θ , φ , x ) = Eqφ ( z|x ) [ log pθ ( x | z ) ] −DKL [ qφ ( z | x ) ‖ p ( z ) ] ( 11 ) Lnew ( θ , φ , x ) = Lelbo ( θ , φ , x ) − αSmax ( { Z1 , Z2 , ... , Zd } ; X ) ( 12 ) Expanding the Smax term , we have the llowing : Lnew ( θ , φ , x ) = Lelbo ( θ , φ , x ) −α ( I ( z ; x ) − ∑ x∈X p ( X = x ) max i DKL [ qφ ( Ai |x ) ‖ p ( Ai ) ] ) ( 13 ) From Hoffman & Johnson ( 2016 ) , we know that the KL term in the ELBO loss is decomposed in DKL [ qφ ( zn ) ‖ p ( zn ) ] + I ( xn ; z ) when we use the aggregate posterior and define the loss over the empirical distribution of the data pdata ( x ) . Taking in account that , we can express the equation 12 as follows : Lnew ( θ , φ , x ) = 1 N N∑ i=1 [ Eqφ ( z|x ) [ log pθ ( x ( i ) | z ) ] ] −DKL [ qφ ( zn ) ‖ p ( zn ) ] − I ( xn ; z ) −αI ( xn ; z ) ︸ ︷︷ ︸ Penalise +α ∑ x∈X p ( X = x ) max i DKL [ qφ ( Ai | x ) ‖ p ( Ai ) ] ) ( 14 ) If we penalise the synergy ( see Eq . 12 ) , we will be penalising the mutual information term which is not desirable for this task Kim & Mnih ( 2018 ) ; we can see this effect explicitly in Eq . 14 . Therefore , we use only the second term to perform the optimisation which means maximising the subset of latents with the most amount of MI per outcome . Lnew ( θ , φ , x ) = Lelbo ( θ , φ , x ) + α ∑ x∈X p ( X = x ) max i DKL [ qφ ( Ai | x ) ‖ p ( Ai ) ] ) ( 15 ) It ’ s easy to see in Eq . 15 that it ’ s not a guaranteed lower bound on the log likelihood px anymore , which is why we decided to penalise the subset of latents with the minimum specific mutual information ( ie . Aw ) . In practice we found that computing the maximum subset Ai for each outcome of x is too computational intensive , which is why we decided to use a mini-batch approximation , which is the version we show in the pseudo-code using a two step optimisation in the next section . The final version we are going to use is the one below , where Imax is the KL term of the synergy term . Lnew ( θ , φ , x ) = Eqφ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qφ ( z|x ) ‖p ( z ) ) ︸ ︷︷ ︸ Lelbo −αDKL [ qφ ( Aw|x ) ‖ p ( Aw ) ] ︸ ︷︷ ︸ α∗Imax ( 16 ) Algorithm 1 Non Syn VAE Input : Observations ( x ( i ) ) Ni=1 , batch sizem , latent dimension d , weight of synergy loss α , discount factor ω , optimiser optim , reparametrisation function gφ . θ , φ← Initialise VAE parameters repeat 3 : x ( i ) ← Random minibatch B of size m , i ∈ B zi← gφ ( , x ( i ) ) . Sample zi ∼ qφ ( z|x ) 6 : φ , θ← optim ( ∇θ , φLelbo ( θ , φ ; x ) ) . : Gradients of ELBO minibatch x′ ( i ) ← Random minibatch B ’ of size m , i ∈ B′ mu , logvar← Encoder ( x′ ( i ) , φ ) 9 : worst index← get index greedy ( mu , logvar , ω ) Lsyn← α ∗ Imax ( mu , logvar , worst index ) . See Eq.16 for Imax function φ← optim ( ∇φLsyn ( φ ; x′ ( i ) ) ) . : Gradients of Syn loss minibatch 12 : until convergence of objective In the algorithm shown above we see in practice that we get better results when we sample the values of mu and logvar from the encoder for the step 2 of the optimisation . We use a greedy approximation of the best latents by following a greedy policy ( See Appendix ) .

Summary: This paper proposes a new approach to enforcing disentanglement in VAEs using a term that penalizes the synergistic mutual information between the latent variables, encouraging representations where any given piece of information about a datapoint can be garnered from a single latent.  In other words, representations where there is no information conveyed by combinations of latents that is not conveyed by considering each latent in isolation.  As the resultant target is intractable to evaluate, a number of approximations are employed for practical training.