DynamicVAE: Decoupling Reconstruction Error and Disentangled Representation Learning

1 INTRODUCTION . The goal of disentangled representation learning is to encode input data into a low-dimensional space that preserves information about the salient factors of variation , so that each dimension of the representation corresponds to a distinct factor in the data ( Bengio et al. , 2013 ; Locatello et al. , 2020 ; van Steenkiste et al. , 2019 ) . Learning disentangled representations benefits a variety of downstream tasks ( Higgins et al. , 2018 ; Lake et al. , 2017 ; Locatello et al. , 2019c ; a ; Denton et al. , 2017 ; Mathieu et al. , 2019 ) , including abstract visual reasoning ( van Steenkiste et al. , 2019 ) , zero-shot transfer learning ( Burgess et al. , 2018 ; Lake et al. , 2017 ; Higgins et al. , 2017a ) and image generation ( Nie et al. , 2020 ) , just to name a few . Due to its central importance in various downstream applications , there is abundant literature on learning disentangled representations . Roughly speaking , there are two lines of methods towards this goal . The first category includes supervised methods ( Chen & Batmanghelich , 2019 ; Locatello et al. , 2019c ; Shu et al. , 2019 ; Bouchacourt et al. , 2018 ; Nie et al. , 2020 ; Yang et al. , 2015 ) , where external supervision ( e.g. , data generative factors ) is available during training to guide the learning of disentangled representations . The second line of works focus on unsupervised methods ( Chen et al. , 2016 ; 2018 ; Burgess et al. , 2018 ; Kim & Mnih , 2018 ; Denton et al. , 2017 ; Kumar et al. , 2018 ; Fraccaro et al. , 2017 ) , which substantially relieve the needs to have external supervisions . For this reason , in this paper , we mainly focus on unsupervised disentangled representation learning . One major challenge of unsupervised disentanglement learning is that there exists a trade-off between reconstruction quality of the input signal and the degree of disentanglement in the latent representations . Let us take β-VAE and its variants ( Burgess et al. , 2018 ; Chen et al. , 2018 ; Higgins et al. , 2017a ) as an example . These methods assign a large and fixed weight β in the objective function to improve the disentanglement at the cost of reconstruction quality , which is highly correlated with accuracy in downstream tasks ( van Steenkiste et al. , 2019 ; Locatello et al. , 2020 ) . In order to improve the reconstruction quality , researchers have proposed a dynamic learning approach , ControlVAE ( Shao et al. , 2020 ) , to dynamically adjust the weight on the KL term in the VAE objective to better balance the quality of disentangled representation learning and reconstruction error . However , while ControlVAE allows better control of the trade-off between disentangled representation learning and reconstruction error , it does not eliminate it . One is still achieved at the expense of the other . The contribution of this paper , compared to the above state of the art , lies in demonstrating that with the proper design , the trade-off between disentangled representation learning and reconstruction error is completely eliminated . Both objectives can be attained at the same time in a decoupled fashion , without affecting each other . More specifically , we observe that if β was kept high in the beginning of training then lowered later in the process , the two objectives are decoupled allowing each to be independently optimized . To the authors ’ knowledge , this work is the first to attain such decoupled optimization of both quality of disentanglement and reconstruction error . Our Contributions : In this paper , we propose a novel unsupervised disentangled representation learning method , dubbed as DynamicVAE , that turns the weight of β-VAE ( β > 1 ) ( Burgess et al. , 2018 ; Higgins et al. , 2017a ) into a small value ( β ≤ 1 ) to achieve not only good disentanglement but also a high reconstruction accuracy via dynamic control . We summarize the main contributions of this paper as follows . • We propose a new model , DynamicVAE , that leverages an incremental PI controller and moving average to evolve the desired KL-divergence along a trajectory that enables decoupling of two objectives : high-quality disentanglement and low reconstruction error . • We provide the theoretical conditions on parameters of the PI controller to guarantee stability of DynamicVAE . • We experimentally demonstrate that our approach turns the weight of β-VAE ( β > 1 ) to β ≤ 1 , achieving higher reconstruction quality yet comparable disentanglement compared to prior approaches ( e.g. , FactorVAE ) . Thus , our results verify that the proposed method indeed decouples disentanglement and reconstruction accuracy without hurting each other ’ s performance . 2 PRELIMINARIES . β-VAE and its Variants : β-VAE ( Higgins et al. , 2017b ; Chen et al. , 2018 ) is a popular unsupervised method for learning disentangled representations of the data generative factors ( Bengio et al. , 2013 ) . Compared to the original VAE , β-VAE incorporates an extra hyperparameter β ( β > 1 ) as the weight of the KL term in the VAE objective : Lβ = Eqφ ( z|x ) [ log pθ ( x|z ) ] − βDKL ( qφ ( z|x ) ‖ p ( z ) ) . ( 1 ) In order to discover more disentangled factors , in other variants , practitioners further add a constraint on the total information capacity , C , to control the capacity of the latent channels ( Burgess et al. , 2018 ) to transmit information . The constraint can be formulated as an optimization method : Lβ = Eqφ ( z|x ) [ log pθ ( x|z ) ] − β · |DKL ( qφ ( z|x ) ‖p ( z ) ) − C| , ( 2 ) where β is a large and fixed hyperparameter . As a result , when the weight β is large ( e.g . 100 ) , the algorithm tends to optimize the second term in ( 2 ) , leading to much higher reconstruction error . PID Control Algorithm : The PID is a simple yet effective control algorithm that can stabilize system output to a desired value via feedback control ( Stooke et al. , 2020 ; Åström et al. , 2006 ) . The PID algorithm calculates an error , e ( t ) , between a set point ( in this case , the desired KL-divergence ) and the current value of the controlled variable ( in this case , the actual KL-divergence ) , then applies a correction in a direction that reduces that error . The correction is the weighted sum of three terms , one proportional to the error ( called P ) , one that is the integral of error ( called I ) , and one that is the derivative of error ( called D ) ; thus , the term PID . The derivative term is not recommended for noisy systems , such as ours , reducing the algorithm to PI control . The canonical form of a PI controller ( applied to control β ( t ) ) is the following : β ( t ) = Kpe ( t ) +Ki t∑ j=0 e ( j ) , ( 3 ) where β ( t ) is the output of a controller , which ( in our case ) is the used β during training at time t ; e ( t ) is the error between the output value and the desired value at time t ; Kp , Ki denote the coefficients for the P term and I term , respectively . Eq . ( 3 ) may be rewritten in incremental form , as follows : β ( t ) = ∆β ( t ) + β ( t− 1 ) , ( 4 ) where β ( 0 ) can be set as needed ( as we show later ) , and : ∆β ( t ) = Kp [ e ( t ) − e ( t− 1 ) ] +Kie ( t ) . ( 5 ) This paper adopts a nonlinear incremental form of the PI controller , described later in Section 3 . 3 THE DYNAMICVAE ALGORITHM . The goal of disentangled representation learning ( Burgess et al. , 2018 ) is to maximize the log likelihood and simultaneously stabilize the KL-divergence to a target value C. It can be formulated as the following constrained optimization problem : max φ , θ Eqφ ( z|x ) [ log pθ ( x|z ) ] , s.t . DKL ( qφ ( z|x ) ‖ p ( z ) ) = C ( 6 ) In order to achieve a good trade-off between disentanglement and reconstruction accuracy , we attempt to design a controller to dynamically adjust β ( t ) in the following VAE objective to stabilize the KL-divergence to the desired value C : Ld = Eqφ ( z|x ) [ log pθ ( x|z ) ] − β ( t ) DKL ( qφ ( z|x ) ‖ p ( z ) ) . ( 7 ) The contribution of DynamicVAE is to evolve β ( t ) along a good trajectory to achieve decoupling between disentanglement and reconstruction error . To reach this goal , we need to address the following two challenges : 1. β ( t ) should dynamically change from a large value to small one . Specifically , at the beginning of training , β ( t ) should be large enough to disentangle latent factors . After that , β ( t ) is required to gradually drop to a small value to optimize the reconstruction . 2. β ( t ) should not change too fast or oscillates too frequently . When β ( t ) drops too fast or oscillates , it may cause KL-divergence to grow with a large value . Consequently , some latent factors may emerge earlier so that they can potentially be entangled with each other . In this paper , we propose methods to deal with these two challenges , summarized below . A non-linear incremental PI controller : Fig . 1 ( a ) shows the designed non-linear PI controller that dynamically adjusts the weight β ( t ) in the KL term of the β-VAE , based on the actual KL-divergence , yKL ( t ) . Specifically , it first samples the output KL-divergence , yKL ( t ) , at training step t. Then we use the difference e ( t ) between the sampled KL-divergence at time t with the desired value , C , as the feedback to PI controller to tune β ( t ) . The corresponding PI algorithm is denoted by β ( t ) = Kpσ ( −e ( t ) ) −Ki t∑ j=0 e ( j ) , ( 8 ) where σ ( . ) is a sigmoid function ; Kp , Ki are positive hyper-parameters for P and I term respectively . As mentioned earlier , we need a large β ( t ) in the beginning to control the KL-divergence from a small value to a large target value so that the information can be transmitted through the latent channels per data sample . Accordingly , we adopt an incremental form of the PI controller in Eq . ( 8 ) , and initialize it to a large value : β ( t ) = ∆β ( t ) + β ( t− 1 ) , ( 9 ) where ∆β ( t ) = Kp [ σ ( −e ( t ) ) − σ ( −e ( t− 1 ) ) ] −Kie ( t ) . ( 10 ) and β ( 0 ) is a large initial value . When the PI controller is initialized to a large value β ( 0 ) , it can quickly produce a ( small ) KL-divergence during initial model training , preventing emergence of entangled factors . Moving average : Since our model is trained with mini-batch data , it often contains noise that causes β ( t ) to oscillate . In particular , when β ( t ) plunges during training , it would cause KL-divergence to rise too quickly . This may lead to multiple latent factors coming out together to be entangled . To mitigate this issue , we adopt moving average to smooth the output KL-divergence as the feedback of PI controller below . y ( t ) = αtyKL ( t ) + αt−1yKL ( t− 1 ) + · · ·+ αt−T yKL ( t− T ) = t∑ i=t−T αiyKL ( i ) , ( 11 ) where αi denotes weight and T denotes the window size of past training steps . Hybrid annealing : Control systems with step ( input ) function ( i.e. , those where the set point can change abruptly ) often suffer from an overshoot problem . An overshoot is temporary overcompensation , where the controlled variable oscillates around the set point . In our case , it means that the actual KL-divergence may significantly ( albeit temporarily ) exceed the desired value , when set point is abruptly changed . This effect would cause some latent factors to come out earlier than expected , and be entangled , thereby producing poor-quality disentanglement . To address this problem , we develop a hybrid annealing method that changes the set point more gradually , as illustrated in Fig . 7 in Appendix . It combines step function with ramp function to smoothly increase the target KL-divergence in order to prevent overshoot and thus better disentangle latent factors one by one . The combination of the above three methods allows DynamicVAE to evolve β ( t ) along a favorable trajectory to separate disentanglement learning and reconstruction optimization . We summarize the proposed incremental PI algorithm in Algorithm 1 in Appendix B .

Summary: This paper introduces a strategy for controlling the beta value of a beta-VAE during training using approaches from control theory that allow it to target a designated level of KL divergence between the encoder distribution and the prior.  This is done in a way that aims to achieve good reconstructions while maintaining disentangling performance.  It can be viewed as a refinement of the ControlVAE approach of Shao et al 2020, varying only in the specifics of the strategy used.