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 focus of this work is disentangled representation in sequential data. It merges existing approaches/views of contrastive learning, disentanglement, and dynamical data modeling, into a unified framework with the goal to extract better representations of sequential data (video, audio experiments).  The paper is well written and self-contained which makes it easy to follow. The authors made an effort to present related literature (with useful remarks on differences) and in the experiments, they compare to the most relevant state-of-art approaches.  Many insightful results are redelegated to the supplementary even though the paper could have content on 9 pages (submitted version has 8).  Although this paper builds on previous works, the detailed analysis of both theoretical and empirical findings makes their contribution in terms of contrastive disentanglement valuable to the community.