Contrastively Disentangled Sequential  Variational Autoencoder

1 Introduction . The goal of self-supervised learning methods is to extract useful and general representations without any supervision , and to further facilitate downstream tasks such as generation and prediction [ 1 ] . Despite the difficulty of this task , many existing works have shed light on this field across different domains such as computer vision [ 2 , 3 , 4 , 5 ] , natural language processing [ 6 , 7 , 8 ] and speech processing [ 9 , 10 , 11 , 12 , 13 ] ( also see a huge number of references in these papers ) . While the quality of the learnt representations improves gradually , recent research starts to put more emphasis on learning disentangled representations . This is because disentangled latent variables may capture separate variations of the data generation process , which could contain semantic meanings , provide the opportunity to remove unwanted variations for a lower sample complexity of downstream learning [ 14 , 15 ] , and allow more controllable generations [ 16 , 17 , 18 ] . These advantages lead to a rapidly growing research area , studying various principles and algorithmic techniques for disentangled representation learning [ 19 , 20 , 21 , 22 , 23 , 24 , 25 , 26 ] . One concern raised in [ 24 ] is that without any inductive bias , it would be extremely hard to learn meaningful disentangled representations . On the other hand , this concern could be much alleviated in the scenarios where the known structure of the data can be exploited . In this work , we are concerned with the representation learning for sequence data , which has a unique structure to utilize for disentanglement learning . More specifically , for many sequence data , the variations can be explained by a dichotomy of a static ( time-invariant ) factor and dynamic ( timevariant ) factors , each varies independently from the other . For example , representations of a video recording the movements of a cartoon character could be disentangled into the character identity ( static ) and the actions ( dynamic ) . For audio data , the representations shall be able to separate the speaker information ( static ) from the linguistic information ( dynamic ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We propose Contrastively Disentangled Sequential Variational Autoencoder ( C-DSVAE ) , a method seeking for a clean separation of the static and dynamic factors for the sequence data . Our method extends the previously proposed sequential variational autoencoder ( VAE ) framework , and performs learning with a different evidence lower bound ( ELBO ) which naturally contains mutual information ( MI ) terms to encourage disentanglement . Due to the difficulty in estimating high dimensional complex distributions ( e.g. , for the dynamic factors ) , we further incorporate the contrastive estimation for the MI terms with systematic data augmentation techniques which modify either the static or dynamic factors of the input sequence . The new estimation method turns out to be more effective than the minibatch sampling based estimate , and introduces additional inductive biases towards the invariance . To our knowledge , we are the first to synergistically combine the contrastive estimation and sequential generative models in a principled manner for learning disentangled representations . We validate C-DSVAE on four datasets from the video and audio domains . The experimental results show that our method consistently outperforms the previous state-of-the-art ( SOTA ) methods , both quantitatively and qualitatively . 2 Method . We denote the observed input sequence as x1 : T = { x1 , x2 , ... , xT } where xi represents the input feature at time step i ( e.g. , xi could be a video frame or the spectrogram feature of a short audio segment ) , and T is the sequence length . The latent representations are divided into the static factor s , and the dynamic factors z1 : T where zi is the learnt dynamic representation at time step i . 2.1 Probabilistic Model . We assume that in the ground-truth generation process , zi depends on z < i = { z0 , z1 , ... , zi−1 } where z0 = 0 , and the observation xi is independent of other frames conditioned on zi and s. Furthermore , we assume the static variable s and the dynamic variables z1 : T are independent from each other , i.e. , p ( s , z1 : T ) = p ( s ) p ( z1 : T ) . Formally , let z = ( s , z1 : T ) and we have the following complete likelihood p ( x1 : T , z ) = p ( z ) p ( x1 : T |z ) = [ p ( s ) T∏ i=1 p ( zi|z < i ) ] · T∏ i=1 p ( xi|zi , s ) ( 1 ) where p ( z ) is the prior for sampling z . Our formulation captures the general intuition that we can separate the variations of the sequence into the time-dependent dynamic component ( described by zi ’ s ) , and a static component ( described by s ) which remains the same for different time steps in the same sequence but differs between sequences . For example , in the cartoon character video ( see Figure 3 ) , the hair , shirt and pants should be kept the same when the character walks around , but different videos can have different characters . Similarly , for a speech utterance , the phonetic transcription controls the vocal tract motion and the sound being produced over time , but the speaker identity remains the same for the whole utterance . In this work , we refer to the dynamic component as `` motion '' and the static component as `` content '' . For the prior distributions , we choose p ( s ) to be the standard Gaussian N ( 0 , I ) , and p ( zi|z < i ) to be N ( µ ( z < i ) , σ2 ( z < i ) ) where µ ( · ) and σ ( · ) are modeled by LSTMs [ 27 ] . In operations , the drawn sample s is shared throughout the sequence . To draw a sample of zt , we first take a sample of zt−1 as the input of the current LSTM cell . Then forwarding the LSTM for one step gives us the distribution of zt , from which we draw a sample with the reparameterization trick [ 28 ] . To extract the latent representations given only the observed data x1 : T where the motion and content are mixed together , we hope to learn a posterior distribution q ( z|x1 : T ) where the two components are disentangled . That is , similar to the prior , our posterior should have a factorized form : q ( z|x1 : T ) = q ( z1 : T , s|x1 : T ) = q ( z1 : T |x1 : T ) q ( s|x1 : T ) = q ( s|x1 : T ) T∏ i=1 q ( zi|z < i , x≤i ) . ( 2 ) The posterior distributions are also modeled by LSTMs . A nested sampling procedure , resembling that of the prior , is applied to the posterior of dynamic variables . Similar parameterizations of dynamic variables by recurrent networks have been proposed in prior works [ 29 , 30 , 31 ] . The standard loss function for learning the latent representations is an ELBO [ 32 , 33 ] ( also see a derivation in Appendix A.1 ) : max p , q Ex1 : T∼pDEq ( z|x1 : T ) [ log p ( x1 : T |z ) −KL [ q ( z|x1 : T ) ||p ( z ) ] ] ( 3 ) where pD is the empirical data distribution . Under the parameterization of the posterior where s and z1 : T are mutually independent , the KL-divergence term reduces to KL [ q ( z|x1 : T ) ||p ( z ) ] = KL [ q ( s|x1 : T ) ||p ( s ) ] +KL [ q ( z1 : T |x1 : T ) ||p ( z1 : T ) ] ( 4 ) where the second term is approximated with the sampled trajectories of the dynamic variables z1 : T . 2.2 Our approach : Mutual information-based disentanglement . Several existing sequence representation learning methods [ 32 , 33 , 18 ] are built on top of the loss function ( 3 ) . However , this formulation also brings several issues . The KL-divergence regularizes the posterior of the static or dynamic factors to be close to the corresponding priors . When modeling them with powerful neural architectures like LSTMs , it is possible for the KL-divergence to be close to zero , yet at the same time , the posteriors become non-informative of the inputs ; this is a common issue for deep generative models like VAE [ 34 ] . Techniques have been proposed to alleviate this issue , including adjusting the relative weights of the loss terms to regularize the capacity of posteriors [ 20 ] , replacing the individual posteriors in the KL terms with aggregated posteriors [ 35 , 36 , 37 ] and enforcing latent structures ( such as disentangled representations ) in the posteriors [ 23 , 22 , 38 ] . Our principled approach for learning useful representations from sequences is inspired by these prior works , and at the same time incorporates the unique sequential structure of the data . Without inductive biases , the goal of disentanglement can hardly be achieved since it is possible to find entangled s and z1 : T that explain the data equally well ( in fact , by Theorem 1 of [ 24 ] , there could exist infinitely many such entangled factors ) . The problem may seem less severe in our setup , since s is shared across time and it is hard for such s to capture all dynamics . Nevertheless , since both s and zi are used in generating xi , it is still possible for zi to carry some static information . In the extreme case where each zi encompasses s , s would no longer be indispensable for the generation . This issue motivated prior works to optimize the ( estimate of ) mutual information among latent variables and inputs [ 33 , 18 , 39 ] . Our method seeks to achieve clean disentanglement of s and z1 : T , by optimizing the following objective function , which introduces additional MI terms to the vanilla ELBO in ( 3 ) : max p , q Ex1 : T∼pDEq ( z|x1 : T ) [ log p ( x1 : T |z ) ] −KL [ q ( z ) ||p ( z ) ] = Ex1 : T∼pDEq ( z|x1 : T ) [ log p ( x1 : T |z ) ] − ( KL [ q ( s|x1 : T ) ||p ( s ) ] +KL [ q ( z1 : T |x1 : T ) ||p ( z1 : T ) ] ) + Iq ( s ; x1 : T ) + Iq ( z1 : T ; x1 : T ) − Iq ( z1 : T ; s ) ( 5 ) where the aggregated posteriors are defined as q ( z ) = q ( s , z1 : T ) = EpD [ q ( s|x1 : T ) q ( z1 : T |x1 : T ) ] , q ( s ) = EpD [ q ( s|x1 : T ) ] , q ( z1 : T ) = EpD [ q ( z1 : T |x1 : T ) ] , and the MI terms are defined as Iq ( s ; x1 : T ) = Eq ( s , x1 : T ) [ log q ( s|x1 : T ) q ( s ) ] , Iq ( z1 : T ; x1 : T ) = Eq ( z1 : T , x1 : T ) [ log q ( z1 : T |x1 : T ) q ( z1 : T ) ] ( and Iq ( z1 : T ; s ) is defined similarly ) . The intuition behind ( 5 ) is simple : besides explaining the data and matching the posteriors with priors , z1 : T and s shall contain useful information from x1 : T while excluding the redundant information from each other . We further justify ( 5 ) by showing that it still forms a valid ELBO . Theorem 1 . With our parameterization of ( s , z1 : T ) , ( 5 ) is a valid lower bound of the data loglikelihood Ex1 : T∼pD log ( x1 : T ) . The full proof can be found in Appendix A.2 . With this guarantee , we can follow the spirit of [ 20 , 23 ] to add and adjust the additional weight coefficients α to the KL terms , β to the MI terms Iq ( s ; x1 : T ) , Iq ( z1 : T ; x1 : T ) , and γ to Iq ( z1 : T ; s ) , Ex1 : T∼pDEq ( z|x1 : T ) [ log p ( x1 : T |z ) ] − α ( KL [ q ( s|x1 : T ) ||p ( s ) ] +KL [ q ( z1 : T |x1 : T ) ||p ( z1 : T ) ] ) + β ( Iq ( s ; x1 : T ) + Iq ( z1 : T ; x1 : T ) ) − γIq ( z1 : T ; s ) . ( 6 ) It remains to estimate the objective ( 6 ) for optimization . The KL term for z1 : T is estimated with the standard Monte-Carlo sampling , using trajectories of z1 : T [ 32 ] . The KL term for s can be estimated analytically . For the MI terms , we attempt two estimations . The first estimation uses a standard mini-batch weighted sampling ( MWS ) following [ 23 , 33 ] . Due to the high dimensionality of z1 : T and the complex dependency among time steps , it may be hard for MWS to estimate the distributions accurately , so we also explore non-parametric contrastive estimations for Iq ( s ; x1 : T ) and Iq ( z1 : T ; x1 : T ) with additional data augmentations , which we will detail below .

Summary: The paper makes contributions to the research on disentangled sequential representation learning based on Variational Autonecoders (VAES). Towards this end, it  1)  proposes an objective function, which theoretically constitutes a lower bound of the marginal likelihood, while containing interpretable and intuitive mutual information loss terms for encouraging disentanglement between a static and a sequence of dynamic latent factors.  2) proposes a non-parametric estimation of the aforementioned mutual information terms that leverages data-augmentation to further improve performance.  3) demonstrates both quantitative and qualitative performance gains compared to other works on a variety of disentanglement metrics and video datasets.