Disentangled Representation Learning with Sequential Residual Variational Autoencoder

1 INTRODUCTION . Learning a sparse and interpretable representation of data is a critical component of a generalized , robust and explanatory intelligent system . This concept is inspired by human ’ s ability to generalize the knowledge with abstract concepts and use them to reason the unseen environments Gupta et al . ( 2018 ) . Despite recent advances on representation learning , it was shown that deep convolutional neural networks ( CNN ’ s ) have a tendency to learn superficial statistics of data associated with given tasks , rather than important generative factors embedded in the physical world Jo & Bengio ( 2017 ) ; Goodfellow et al . ( 2014 ) . One way towards this goal is disentangled representation learning which aims to capture the independent and interpretable generative factors of the data . Bengio et al . ( 2013 ) defines the disentangled representation intuitively as a representation where changes in one dimension correspond to changes in only one generative factor of the data , while being relatively invariant to changes in other factors . Recently , Higgins et al . ( 2018 ) assigned a principled definition by connecting symmetry transformations to vector representations using the group and representation theory . Based on these definitions , disentangled representation can be learned in a supervised fashion where explicit and/or implicit prior knowledge on the generative factors of data are available . However , it is ideal to achieve this in an unsupervised learning setting to take advantage of the large amount of available unlabeled data . Along with the recent development of the generative models , many unsupervised disentangled learning approaches have been proposed based on either the generative adversarial networks ( GAN ) ( proposed as InfoGAN in Chen et al . ( 2016 ) ) or the variational autoencoders ( VAE ) ( proposed as β-VAE in Higgins et al . ( 2017 ) ) . While β-VAE achieves better results and does not suffer from the training stability issue of InfoGAN , it faces a trade-off between the disentanglement and reconstruction due to its information bottleneck . The current state-of-the-art approaches extend the β-VAE with augmented objective function to reduce this trade-off Burgess et al . ( 2017 ) ; Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) ; Kumar et al . ( 2017 ) . A recent study by Locatello 1Codes available at : https : //www.dropbox.com/s/5hkfn8xy5r8w5sz/Code.zip ? dl=0 et al . ( 2018 ) carefully compared these approaches based on extensive experiments . They found that the performance of these approaches is very sensitive to the hyperparameter tuning associated with the augmented objective function and the initial random seed during training . More importantly , they proved that unsupervised learning of disentangled representation is impossible without introducing inductive bias on either the model or the data . We believe the trade-off between disentanglement and reconstruction in VAE-based approaches can be addressed by a different training approach . The idea of relying on modified training approaches , instead of augmented objective function , to encourage network behavior is commonly used for different problems . Take model over-fitting prevention for example , one way to address this is to augment the objective function with regularization terms , such as L1 or L2 regularization . An alternative solution is to apply special operations during training to enforce the generalization of the network representations , such as Dropout Srivastava et al . ( 2014 ) or Batch Normalization Ioffe & Szegedy ( 2015 ) . Our main contribution in this work is four-fold : 1 ) We propose Sequential Residual Variational Autoencoder ( SR-VAE ) that uses a novel “ Residual learning ” mechanism to learn disentangled representation with the original VAE objective . This is different from previous VAE-based approaches that merely focus on objective function design where hyperparameter tuning is crucial . 2 ) We show the proposed “ Residual learning ” mechanism defines an explicit dependency structure among the latent variables via sequential latent variable update . This encourages learning the disentangled representation . 3 ) We highlight that SR-VAE decomposes the reconstruction into residual and network decoder output via skip connection . This relaxation of reconstruction reduces the trade-off between disentanglement and reconstruction of VAE . 4 ) We demonstrate both qualitatively and quantitatively that SR-VAE improves the current state-of-the-art disentanglement representation learning performance on a variety of complex datasets . 2 CHALLENGES OF DISENTANGLING WITH AUGMENTED VAE OBJECTIVE . In this section , we first briefly review the VAE framework , followed by the β-VAE and its extensions for disentangled representation learning . We highlight the challenges of using an augmented objective function to balance the VAE ’ s trade-off between the disentanglement and reconstruction . From these discussions , we then motivate the proposed SR-VAE framework . VAE is a deep directed graphical model consisting of an encoder and a decoder Kingma & Welling ( 2013 ) . The encoder maps the input data x to a latent representation qθ ( z|x ) and the decoder maps the latent representation back to the data space qφ ( x|z ) , where θ and φ represent model parameters . The loss function of the VAE is defined as following : LV AE = Eqθ ( z|x ) [ log qφ ( x|z ) ] −KL ( qθ ( z|x ) ‖ p ( z ) ) , ( 1 ) where KL ( . ‖ . ) stands for the Kullback-Leibler divergence . By regularizing the posterior qθ ( z|x ) with a prior over the latent representation p ( z ) ∼ N ( 0 , I ) , where I is identity matrix , VAE learns a latent representation qθ ( z|x ) that contains the variations in the data . The goal of disentangled representation learning is to identify the latent representation z ∈ Rd where each latent variable only corresponds to one of the generative factors for given data x . To achieve this , β-VAE augments VAE objective with an adjustable hyperparameter β as : Lβ−V AE = Eqθ ( z|x ) [ log qφ ( x|z ) ] − βKL ( qθ ( z|x ) ‖ p ( z ) ) . ( 2 ) The addition of β encourages the posterior qθ ( z|x ) to match the factorized unit Gaussian prior p ( z ) . It enhances the independence among the latent variables thus disentangling the representation . On the other hand , it reduces the amount of information about x stored in z , which can lead to a poor reconstruction especially for high values of β . This trade-off is further discussed from the rate-distortion theory perspective in Burgess et al . ( 2017 ) . To reduce the trade-off , different augmentations of β-VAE objective are proposed Burgess et al . ( 2017 ) ; Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) ; Kumar et al . ( 2017 ) . Locatello et al . ( 2018 ) categorized these methods into three main categories of bottleneck capacity , penalizing the total correlation and disentangled priors . Burgess et al . ( 2017 ) focused on bottleneck capacity and proposed to gradually increase the average KL divergence from zero for each generative factor . This method relaxes the information bottleneck during training via increasing the encoding capacity through a parameter C that is linearly dependent on the training iteration . Kim & Mnih ( 2018 ) aimed to solve the disentangled representation learning by minimizing the total correlation ( TC ) term . They proposed FactorVAE where the objective is augmented with a TC term controlled by hyperparameter γ . Minimizing the TC term forces the hidden representation to be factorial and hence independent . Chen et al . ( 2018 ) looked into an alternative way to minimize the TC term in the augmented objective , named β-TCVAE . They used a mini-batch based alternative instead of a density-ratio-trick based method from FactorVAE . However , the results from both methods are sensitive to hyperparameter associated with the TC term . Kumar et al . ( 2017 ) studied the disentangled prior and introduced a regularizer to the objective that is associated with hyperparamter λ . This regularizer encourage the covariance of qφ ( z ) to match the identity matrix . While all aforementioned approaches have shown promising results , they rely on a careful tuning of the hyperparamater introduced in the augmented objective functions such as β in Higgins et al . ( 2017 ) , C in Burgess et al . ( 2017 ) , γ in Kim & Mnih ( 2018 ) and Chen et al . ( 2018 ) , and λ in Kumar et al . ( 2017 ) . Finding the optimal hyperparamater setting can be challenging especially in an unsupervised learning setting where the evaluation of the results mainly relies on visualization and human inspection . More importantly , Locatello et al . ( 2018 ) found that hyperparameter tuning is more important for state-of-the-art performance than the choice of augmented objective functions . In this work , we propose SR-VAE to address the aforementioned challenge . Instead of the forward pass in the original VAE , SR-VAE uses a “ Residual learning ” forward pass illustrated in Fig . 1 as an inductive bias on the model . SR-VAE consists of two components : 1 ) Explicit dependency in the latent space via a multi-step sequential forward pass where one latent variable is updated at each step ; 2 ) Decomposition of the reconstruction via skip connetion between the input and network output to relax the network reconstruction constraint . Together , these two components enable SR-VAE to address the trade-off between the reconstruction and disentanglement in VAE using the original objective in Eq . 1 . In the next section , we first describe the details of the “ Residual learning ” forward pass and SR-VAE training . We then discuss the two aforementioned components in detail . In Section 5 , we demonstrate the effectiveness of SR-VAE and investigate the effect of each component separately . The experimental results show our approach achieves better disentanglement and reconstruction results compared to the current state-of-the-art approaches , including β-VAE Higgins et al . ( 2017 ) and FactorVAE Kim & Mnih ( 2018 ) . 3 SEQUENTIAL RESIDUAL VARIATIONAL AUTOENCODER – SR-VAE . Same as the original VAE , SR-VAE consists of an encoder network noted as qθ ( ~z|x ) , and a decoder network noted as qφ ( x|~z ) . Here x and ~z stand for input data and latent representation vector ; θ and φ represent encoder and decoder network parameters . Let the dimension of the latent representation be d , SR-VAE learns ~z = [ z1 , z2 , . . . , zd ] ∈ Rd as the latent representation of the data . Its forward pass follows a “ Residual learning ” mechanism that consists of d steps . Each step only updates one latent variable . In the first step , the input data x passes through the encoder to calculate the parameterized posterior , noted as ~µ1 and ~σ1 . Instead of drawing samples for all latent variables ~z ∼ N ( ~µ1 , ~σ1 ) , we only sample the first latent variable z1 ∼ N ( ~µ1 [ 1 ] , ~σ1 [ 1 ] ) and set the remaining latent variables to 0 . The modified latent variables ~z = [ z1 , 0 , . . . , 0 ] then passes the decoder to generate the output noted Algorithm 1 SR-VAE Forward Pass Input : observation x , latent dimension d , VAE encoder qθ ( z|x ) , VAE decoder , qφ ( x|z ) Output : reconstruction x ′ , latent representation ~µ ′ , ~σ ′ 1 : γ1 ← x 2 : ~µ ′ = [ 0 , . . . , 0 ] ∈ Rd 3 : ~σ ′ = [ 0 , . . . , 0 ] ∈ Rd 4 : for i = 1 to d do 5 : { ~µi , ~σi } ← Encoder : qθ ( γi ) 6 : ~µ ′ [ i ] = ~µi [ i ] 7 : ~σ ′ [ i ] = ~σi [ i ] 8 : ~z ← Reparameterize ( ~µ ′ , ~σ ′ ) 9 : x̂i ← Decoder : qφ ( ~z ) 10 : if i < d then 11 : γi+1 ← γi − x̂i 12 : end if 13 : end for 14 : x ′ ← x̂d + γd Algorithm 2 SR-VAE Learning Input : Dataset X , batch size m , latent dimension d , Initialize VAE parameters θ , φ 1 : repeat 2 : Randomly select batch x = ( x ( i ) ) i∈B of size m , 3 : { x′ , ( ~µ , ~σ ) } ← Forward_pass ( x ) 4 : Lrecon ←MSE_loss ( x , x′ ) 5 : LKL ← − 12 B∑ j=1 d∑ i=1 [ 1 + log ( σj [ i ] ) 2 − ( µj [ i ] ) 2 − ( σj [ i ] ) 2 ] 6 : L ← LKL + Lrecon 7 : { θ , φ } ← Backward ( L ) 8 : until convergence of objective as x̂1 . We subtract the decoder output from the skip connection ( defined as an Identity function ) as the input for the second pass , noted as γ2 = γ1 − x̂1 . In the second pass , γ2 passes the same encoder to generates a new parameterized posterior ( ~µ2 and ~σ2 ) . This time , we sample only the second latent variable with this parameterized posterior as z2 ∼ N ( ~µ2 [ 2 ] , ~σ2 [ 2 ] ) . We re-sample the first latent variable with z1 ∼ N ( ~µ1 [ 1 ] , ~σ1 [ 1 ] ) while setting the remaining latent variables to 0 . The modified latent variable ~z = [ z1 , z2 , 0 , . . . , 0 ] is then used to generate the new reconstruction x̂2 . We then calculate the corresponding residual γ3 = γ2 − x̂2 as the input for the third pass . In the ith pass , the ith latent variable is sampled from the encoder thus zi ∼ N ( ~µi [ i ] , ~σi [ i ] ) . The previous updated latent variables follow their corresponding residual encoding and the remaining latent variables are set to zeros , ~z = [ z1 , z2 , . . . , zi , 0 , . . . , 0 ] . The process repeats d times such that all the latent variables are sampled . In step d , the final output of SR-VAE , x′ consists of the decoder output x̂d and the residual term γd as x′ = x̂d + γd . In the case where d = 1 , SR-VAE follows the last step and the input is connected to the output through the skip connection . Algorithm 1 shows the pseudo code of the “ Residual learning ” forward pass in SR-VAE . We train the SR-VAE with the original VAE objective defined in Eq . 1 . The parameters are updated using the standard back-propagation demonstrated in Algorithm 2 . The prior p ( z ) is set to the isotropic unit GaussianN ( 0 , I ) and posterior qθ ( z|x ) is parameterized as Gaussians with a diagonal covariance matrix . The “ reparametrization trick ” is used to transform each random variable zi ∼ qθ ( z|x ) as a differentiable transformation of a noise variable ∼ N ( 0 , 1 ) with zi = µi + σi . Due to the sequential update process , SR-VAE can generate a sequence of images during the forward pass . As we shall see in Section 5 , these images reflect image transformations corresponding to the disentangled factors at different steps . Comparing to other VAE based approach that directly generates a single reconstruction output by sampling from the joint distribution ~z ∼ qθ ( ~z|x ) for all latent variables , this step-by-step visual inspection allows for better understanding of the learned generative factor . As a result , SR-VAE provides a new way to understand the disentangled representation results . Explicit Dependency in the Latent Space : The SR-VAE forward pass defines a sequential update of latent variables : the added latent variable zi at step i learns from the residual between the input data and the previously updated latent variables zj , ∀j ∈ { 1 , . . . , i− 1 } . This procedure defines explicit dependency among the latent variables in the posterior that can be written as qθ ( z1 , z2 , ... , zd|x ) = qθ ( z1|x ) qθ ( z2|z1 , x ) ... qθ ( zd|z1 , ... , zd−1 , x ) . The KL loss term of the original VAE objective in Eq . 1 encourages the posterior qθ ( z1 , z2 , ... , zd|x ) to match the factorized unit Gaussian prior p ( ~z ) . Adding the explicit dependency by the “ Residual learning ” mechanism , the SR-VAE objective can be seen as a modified VAE objective : maximize θ , φ LSR−V AE = Eqθ ( ~z|x ) [ log qφ ( x|~z ) ] −KL ( qθ ( ~z|x ) ‖ p ( ~z ) ) , subject to p ( z1 ) ∼ qθ ( z1|x ) , p ( z2 ) ∼ qθ ( z2|z1 , x ) , . . . , p ( zd ) ∼ qθ ( zd|z1 , ... , zd−1 , x ) . ( 3 ) These constraints encourage the newly added latent variable to be independent of the ones already added , thus enhance the disentanglement of the latent representation . Moreover , the solution space of Eq . 3 is a subset of the original VAE . The constrained objective limits the optimization search space to regions where a better local optimal solution exists . We empirically verify this result in terms of both the performance and stability to random initialization in Section 5 . Decomposition of the Reconstruction : The final output of SR-VAE , x′ consists of the decoder output and the residual term as x′ = x̂d+γd . This formulation relaxes the reconstruction constraint on the network ’ s decoder output when comparing with other VAE-based approaches . More importantly , it creates a balancing measure between the data generation and reconstruction . In one extreme case , the input x directly passes through the first d− 1 steps and reaches step d as the input . In this case , SR-VAE becomes the original VAE model with added skip connection between input and output ( see the last step in Fig 1 ) . This architecture relaxes the VAE reconstruction hence reduces VAE reconstruction and disentanglement trade-off . We will show in Section 5.1 , this architecture alone can reach similar performance to FactorVAE . On the other extreme case , if the first d − 1 steps have learned a perfect disentangled representation of the data , the input for step d would be 0 . In this case , the reconstruction loss encourages SR-VAE to generate the input data from learned latent representation vectors . Combining these two extreme cases , SR-VAE can be understood as a training mechanism to balance between a VAE model with emphasis on the reconstruction quality ( as the first case ) and the data generation model given the learned latent variables ( as the latter case ) . Notice that each of the aforementioned components can be separately added to VAE as a modified model . To add the explicit dependency in the latent space component , we can apply the sequential forward pass of SR-VAE with the output x′ = x̂d . We refer to this model as SeqVAE . For the Decomposition of the Reconstruction component , as mentioned earlier , it is equivalent to adding a skip connection to the original VAE between the input and output . We refer to this model as ResVAE . Using these two models , we perform the ablation study to understand the effectiveness of each individual component in Section 5.1 . Computational Complexity : SR-VAE replaces the standard forward pass of VAE with d forward passes , thus increases the computational complexity . However , in addition to the improved stateof-the-art performance , it eliminates the hyperparameter tuning associated with prior works . As mentioned earlier , the hyperparameter tuning was shown to be critical for state-of-the-art performance . It is a difficult and time-consuming process especially for unlabeled data due to : 1 ) The large hyperparameter search space of continuous values ; 2 ) The lack of evaluation metric . As a result , we believe that the increased computational complexity by SR-VAE is reasonable . Moreover , we will show that each of the d forward passes in SR-VAE correspond to a disentangled generative factor . Visualization of these intermediate steps provides a new way to understand the result .

Summary: The authors of this paper present a novel that for unsupervised disentangled representation learning. The model, named sequential residual VAE (SR-VAE), gradually activates individual latent variables to reconstruct residuals. Quantitative and qualitative experiments show that the proposed model outperforms beta-VAE and Factor-VAE.  Since the training involves a sequence of model training,  SR-VAE certainly consumes more time than other VAEs. Minors: citations in the main text should be put in brackets.