Improving the Unsupervised Disentangled Representation Learning with VAE Ensemble

Variational Autoencoder ( VAE ) based frameworks have achieved the state-of-the- 1 art performance on the unsupervised disentangled representation learning . A re- 2 cent theoretical analysis shows that such success is mainly due to the VAE im- 3 plementation choices that encourage a PCA-like behavior locally on data sam- 4 ples . Despite this implied model identifiability , the VAE based disentanglement 5 frameworks still face the trade-off between the local orthogonality and data re- 6 construction . As a result , models with the same architecture and hyperparameter 7 setting can sometimes learn entangled representations . To address this challenge , 8 we propose a simple yet effective VAE ensemble framework consisting of multi- 9 ple VAEs . It is based on the assumption that entangled representations are unique 10 in their own ways , and the disentangled representations are “ alike ” ( similar up to a 11 signed permutation transformation ) . In the proposed VAE ensemble , each model 12 not only maintains its original objective , but also encodes to and decodes from 13 other models through pair-wise linear transformations between the latent repre- 14 sentations . We show both theoretically and experimentally , the VAE ensemble 15 objective encourages the linear transformations connecting the VAEs to be triv- 16 ial transformations , aligning the latent representations of different models to be 17 “ alike ” . We compare our approach with the state-of-the-art unsupervised disen- 18 tangled representation learning approaches and show the improved performance . 19 1 INTRODUCTION 20 . Disentangled representation learning aims to capture the semantically meaningful compositional 21 representation of data ( Higgins et al. , 2018 ; Mathieu et al. , 2018 ) , and is shown to improve the 22 efficiency and generalization of supervised learning ( Locatello et al. , 2019 ) , reinforcement learning 23 ( Watters et al. , 2019 ) , and reasoning tasks ( van Steenkiste et al. , 2019 ) . The current state-of-the- 24 art unsupervised disentangled representation learning deploy the Variational Autoencoder ( VAE ) 25 ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . The main challenge is to reduce the trade-off 26 between learning a disentangled representation and reconstructing input data . Most of the recent 27 works extend the original VAE objective with carefully designed augmented objective to address 28 this trade-off ( Higgins et al. , 2017 ; Burgess et al. , 2017 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ; 29 Kumar et al. , 2017 ) . A recent study in ( Locatello et al. , 2018 ) compared these methods and showed 30 that their performance is sensitive to initialization and hyperparameter setting of the augmented 31 objective function . 32 Recently , Duan et al . ( Duan et al. , 2019 ) developed an unsupervised model selection method named 33 Unsupervised Disentanglement Ranking ( UDR ) to address the challenge of hyperparameter search 34 and model selection . UDR leverages the finding in ( Rolinek et al. , 2019 ) that the implementation 35 choices of VAE encourage a local PCA-like behavior locally on data samples . As a result , disen- 36 tangled representations by VAEs are “ alike ” as they are similar up to signed permutation transfor- 37 mations . On the contrary , the entangled representations by VAEs are “ unique ” as they are similar 38 at least up to non-degenerate rotation matrices . UDR uses multiple models trained with different 39 initializations and hyperparameter settings , and builds a similarity matrix measuring the pair-wise 40 similarity between the latent variables from different models . A higher score is given to the model 41 that can match its representations to many others models . The results show close match between 42 UDR and commonly used supervised metrics , as well as the performance of downstream tasks using 43 the latent representations . 44 Inspired by the findings from these studies , we propose a simple yet effective VAE ensemble frame-45 work to improve the disentangled representation by VAE . The proposed VAE ensemble consists of46 multiple VAEs . The latent variables in every pair of these models are connected through linear lay-47 ers to force the latent representations in the ensemble to be similar up to a linear transformation.48 We show that the VAE ensemble objective encourages these pair-wise linear transformations to con-49 verge to trivial transformations , making latent representations of different VAEs in the ensemble to50 be “ alike ” , thus disentangled . In this paper , we make the following contributions : ( 1 ) We introduce51 a simple yet effective VAE ensemble framework to improve the disentangled representation learning52 using the original VAE . ( 2 ) We show in theoretical analysis that the linear transformations connect-53 ing the latent representations of the individual models in the ensemble tend to converge to trivial54 transformations thus encourage disentangled representation , and verify this result with experiments.55 ( 3 ) We evaluate our approach using the original VAE model , and show the improved state-of-the-art56 performance across different datasets.57 2 RELATED WORK58 Variatioanl Autoencoder is a deep directed probabilistic graphical model consisting of an encoder59 and a decoder ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . The encoder qφ ( z|x ) maps the60 input data x ∈ Rn to a probabilistic distribution as the latent representation z ∈ Rd , and the decoder61 qθ ( x|z ) maps the latent representation to the data space noted as qθ ( x|z ) , where φ and θ represent62 model parameters . The VAE objective is to maximize the marginalized log-likelihood of data . Direct63 optimization of this objective is not tractable and it is approximated by the evidence lower bound64 ( ELBO ) as:65 LV AE = Eqφ ( z|x ) [ log qθ ( x|z ) ] −KL ( qφ ( z|x ) ‖ p ( z ) ) , ( 1 ) In practice , the first term is estimated by reconstruction error . The second term is the Kullback-66 Leibler divergence between the posterior qφ ( z|x ) and the prior p ( z ) commonly chosen as an67 isotropic unit Gaussian p ( z ) ∼ N ( 0 , I ) .68 Disentangled representation by VAE has achieved the state-of-the-art performance ( Higgins et al.,69 2017 ; Burgess et al. , 2017 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ; Kumar et al. , 2017 ) , despite the70 fact that the VAE objective only models the marginal distribution of the data instead of the desired71 joint distribution over data and latent variables . The reason for this success is the implementation72 choices of the VAE framework ( Rolinek et al. , 2019 ) . In practice , the latent variables in VAE often73 work in “ polarized ” modes . The “ passive ” mode is defined by µ2j ( x ) 1 and σ2j ( x ) ≈ 1 , while the74 “ active ” mode is defined by σ2j ( x ) 1 . The “ passive ” latent variables closely approximate the prior75 and have little effect on the decoder . The “ active ” latent variables , on the other hand , are closely76 related to both the per sample KL loss and the decoder output . The “ polarized regime ” enables a77 reformulated VAE objective showing that VAEs optimize a trade-off between data reconstruction78 and orthogonality of the linear approximation of decoder Jacobian locally around a data sample.79 This PCA-like behavior near data points encourages an identifiable disentangled latent space by80 VAE . Furthermore , it was suggested that finding an appropriate “ polarized regime ” is dependent81 on the initialization and the hyperparameter tuning of the state-of-the-art approaches . In this study,82 we show that the proposed VAE ensemble aligns the “ polarized regime ” of individual VAE models83 towards the disentangled representation.84 Model selection In practice , we often observe neural networks achieve similar performance with85 different internal representations when trained with the same hyperparameters ( Raghu et al. , 2017 ; 86 Wang et al. , 2018 ; Morcos et al. , 2018 ) . For the unsupervised disentanlged representation , as dis-87 cussed in ( Locatello et al. , 2018 ; Duan et al. , 2019 ) , we often observe high variance in the perfor-88 mance from the model trained with the same architecture and hyperparameter setting . This poses89 a challenge for choosing the model in practice . Duan et al . ( 2019 ) proposed Unsupervised Disen-90 tanglement Ranking ( UDR ) to address this challenge . The extensive empirical evaluations on UDR91 using both the supervised metric measurement and the performance of downstream tasks validates92 its effectiveness . They also confirm that disentangled representations are “ alike ” and entangled rep-93 resentations are unique in their own ways . The proposed VAE ensemble leverages this finding.94 Identifiable VAE Built on the recent breakthroughs in nonlinear Independent Component Analy-95 sis ( ICA ) literature ( Hyvarinen & Morioka , 2016 ; 2017 ; Hyvarinen et al. , 2019 ) , Khemakhem et96 al . show that the identification of the true joint distribution over observed and latent variables is97 possible up to very simple transformations ( Khemakhem et al. , 2019 ) . They proposed identifiable 98 VAE ( iVAE ) that requires a factorized prior distribution over the latent variables conditioned on an 99 additionally observed variable , such as a class label or almost any other observation . We believe the 100 proposed VAE ensemble is related to such framework where the latent representation from one VAE 101 model can be regarded as the auxiliary observations for another . 102 Ensemble learning The idea of ensemble learning is to combine multiple learning models ( poten- 103 tially weak learners ) to improve the task performance or robustness over a single model ( Schapire , 104 1990 ) . It achieves the improved performance by averaging the bias , reducing the variance thus pre- 105 venting the over-fitting ( Drucker et al. , 1994 ; Breiman , 1996 ) . Early works in neural networks have 106 used the ensemble learning to achieve top performance in the related competition ( Krizhevsky et al. , 107 2012 ; Simonyan & Zisserman , 2014 ) . In this work , we apply the ensemble learning to enforce the 108 alignment among the latent representations of different models . This results in latent representations 109 that are similar among each other in the ensemble up to a trivial transformation . 110 3 THE VAE ENSEMBLE FRAMEWORK 111 As illustrated in Figure 1 , the proposed VAE ensemble consists of n original VAE models with the 112 same architecture but different initializations . It also consists of n× ( n−1 ) linear layers connecting 113 the latent representations of every two VAE models . Each model in the ensemble maintains its 114 original VAE objective as Eq . 1 . In addition , l2 loss is used to force mapping between latent 115 representations via pair-wise linear layers ( cross-model linear transformation ) . The decoder of each 116 VAE model generates the input reconstruction from not only their corresponding encoder ( within- 117 model reconstruction ) , but also the linearly transformed encodings from other encoders ( cross-model 118 reconstruction ) . Overall , the VAE ensemble is trained with the following objective : 119 L ( θ , φ ) = n∑ i=1 n∑ j=1 Eqφij ( zij |x ) [ log qθj ( x|zij ) ] − n∑ i=1 KL ( qφii ( zii|x ) ‖ p ( zii ) ) − γ n∑ i=1 n∑ j=1 Eqφij ( zij |x ) ‖zjj − zij‖ 2 , ( 2 ) where n is the number of models in the ensemble ; φ : = ( φij ) is the encoders parameters where φij120 represents the encoder of VAEi and its associated linear layer mapping the latent representation from121 VAEi to VAEj ( notice that φii represents the encoder parameters of VAEi only and its associated122 linear transformation can be be regarded as an identity transformation ) ; θ : = ( θi ) represents the de-123 coder parameters of VAEi ; and zjj represents the latent representation of VAEj while zij represents124 the linearly transformed latent representation from VAEi to VAEj . γ is a hyperparameter to balance125 the effect of the estimation error between the latent representaitons . p ( zii ) ∼ N ( 0 , I ) is assume to126 be the prior as defined in the original VAE objective.127 Comparing to the original VAE objective in Eq . 1 , the objective of each individual VAE model in128 the ensemble , L′V AE , can be written as:129 L′V AE ( θ , φ ) = LV AE ( θ , φ ) + n−1∑ j=1 { Eqφj ( zj |x ) [ log qθ ( x|zj ) ] − γEqφj ( zj |x ) ‖zj − z‖ 2 } ︸ ︷︷ ︸ Ensemble Regularization , ( 3 ) where n is the number of models in the ensemble , φj stands for the parameters of the encoder and130 its linear transformation layers from other VAEs , and zj represents the linear transformed latent131 representation of the encoding from other encoders.132 In this form , the VAE ensemble regularizes each VAE model with additional terms on the encoder133 as γEqφj ( zj |x ) ‖zj − z‖ 2 , and on the decoder as ∑n j=1 Eqφj ( zj |x ) [ log qθ ( x|zj ) ] . These regulariza-134 tions directly constrain the latent representations among different VAE models in the ensemble to135 be similar . In particular , for a given input data , ‖zj − z‖2 encourages the encoders to generate136 similar encodings up to the linear transformations ; and ∑n j=1 Eqφj ( zj |x ) [ log qθ ( x|zj ) ] emphasizes137 the similar effect on the data reconstruction from the latent variables such that the decoders can138 reconstruct the input data with both the original encoding z and the linearly transformed encoding139 zj . As we shall discuss in the next section , together these regularizations encourage similar latent140 representation up a trivial transformation by different models in the ensemble.141 We also introduce the hyperparameter γ to balance the trade-off between these two regularizations:142 higher value forces closer mapping between the encoders and reduce the cross-model reconstruction143 error of the decoders ; lower value relaxes the mapping between the encoders and increases the cross-144 model reconstruction error of the decoders . As we show in Section 5 , both components are important145 and balancing the trade-off between them is important as the ensemble size increases.146 Computational complexity It is a common practice to train a number of seeds per hyperparameter147 setting for the current state-of-the-art unsupervised disentanglement VAE models ( Locatello et al.,148 2018 ; Duan et al. , 2019 ) . Comparing to training n original VAEs , the proposed VAE ensemble149 requires additional n × ( n − 1 ) linear layers . While this addition does not increase the size of the150 model much , the estimation of the linear transformations loss and the cross-model reconstruction151 losses grow with n× ( n− 1 ) , which may be computationally expensive especially when n is large.152 That being said , the results in Section 5 show that the VAE ensemble achieves more stable results153 comparing to the current state-of-the-art models . Also , its computation is highly parallelisable.154 4 THEORETICAL JUSTIFICATION155 In this section , we present the theoretical analysis on why the proposed VAE ensemble can improve156 the disentangled representation . We start with analysing the l2 objective in Eq . 2 of the pair-wise157 liner transformations in the VAE ensemble , and show that : ( 1 ) the pair-wise linear transformations158 encourage similar “ polarized ” regime ( see Sec . 2 ) among the VAEs in the ensemble ; ( 2 ) the linear159 transformations are close to the orthogonal transformations . Based on these two properties , we then160 discuss how the cross-model reconstructions by the VAE ensemble encourage learning a disentan-161 gled representation over its entangled counterpart.162 4.1 THE EFFECT OF LINEAR TRANSFORMATION BETWEEN LATENT REPRESENTATIONS163 Let VAEi and VAEj be two different VAE models in the ensemble , and Mji be the linear trans-164 formation that maps the latent representation of a given input x by VAEj to the one by VAEi , as165 zj ( x ) ∼ N ( µj ( x ) , diag ( σj ( x ) 2 ) ) to zi ( x ) ∼ N ( µi ( x ) , diag ( σi ( x ) 2 ) ) . In the following we re- 166 move the input notation from the VAE latent representations for simplicity ( i.e . zj ( x ) is simplified 167 as zj ) , while keeping in mind that the analysis is based on the local latent representation of a given 168 input x . 169 For VAEj , the l2 term of the VAE ensemble loss in Eq . 2 aims to find Mji and zj that minimize 170 E ‖zi −Mjizj‖2 , where the expectation is over the stochasticity of VAEj . We can write zi and zj 171 as zi = µi + i and zj = µj + j , where i ∼ N ( 0 , diag ( σ2i ) ) and j ∼ N ( 0 , diag ( σ2j ) ) . Hence 172 using bias-variance decomposition , the l2 term can be written as : 173 min Mji , zj∼N ( µj , diag ( σ2j ) ) E ‖zi −Mjizj‖2 = min Mji , zj∼N ( µj , diag ( σ2j ) ) { ‖µi −Mjiµj‖2 + E zj ‖Mjiµj −Mjizj‖2 + E zi ‖µi − zi‖2 } = min Mji , zj∼N ( µj , diag ( σ2j ) ) { ‖µi −Mjiµj‖2 + E zj ‖Mjiµj −Mjizj‖2 } + C1 = min Mji , zj∼N ( µj , diag ( σ2j ) ) { ‖µi −Mjiµj‖2 + E j ‖Mji j‖2 } + C1 ( 4 ) where the constant C1 arises from the fact E zi ‖µi − zi‖2 does not depend on Mji and zj . Eq . 4 con- 174 sists of a deterministic component of ‖µi −Mjiµj‖2 and a stochastic component of E j ‖Mji j‖2 . 175 The deterministic component can be minimized by adjusting the parameters in VAEj such that its 176 mean encoding µj is optimized for any given Mji . This simplifies our analysis to focus on the 177 stochastic component . Between Mji and j in this stochastic component , we separately optimize 178 one while having the other fixed . 179 We start with fixed Mji and optimizing for j ∼ N ( 0 , diag ( σ2j ) ) . Notice that σ2j is associated with 180 VAEj objective . In ( Rolinek et al. , 2019 ) , the VAE objective is reformulated into the deterministic 181 reconstruction , the stochastic reconstruction and the KL loss . The last two components define the 182 stochastic loss of VAE . It is formulated as : 183 min V , σ2j ∑ X E j∼N ( 0 , diag ( σ2j ) ) ‖D j‖2 s.t . ∑ X LKL = C1 , ( 5 ) where X represents the dataset , D represents the local linear approximation of the Jacobian of the 184 decoder with singular value decomposition as D = UΣV T . Furthermore , the KL loss LKL = 185 1 2 ∑d k=1 ( µ 2 jk +σ 2 jk− log σ2jk− 1 ) can be simplified as L≈KL = 12 ∑ k∈ “ active ” ( µ 2 jk− log σ2jk− 1 ) 186 based on the “ polarized ” regime of VAE . ( Rolinek et al. , 2019 ) shows that σ2j act as the precision 187 control allowed for each latent variable where more influential ones receive more precision . Com- 188 bining the stochastic loss of the linear transformation in Eq . 4 and the stochastic loss of the original 189 VAE in Eq . 5 , the overall stochastic loss on σ2j can be formulated as : 190 min σ2j E j∼N ( 0 , diag ( σ2j ) ) [ ‖Mji j‖2 + ‖D j‖2 ] s.t . ∑ k − log σ2jk = C2 , ( 6 ) where σ2jk is the kth element of σ 2 j . Here we further simplify LKL with the L≈KL up 191 to additive constants C2 when µj is fixed . In addition to the precision control of σ2j on 192 VAEj , this objective also aims to find an optimal distribution of σ2j that aligns the “ polarized 193 regime ” among different VAEs . To see why , let ck be the kth column of Mji , we then have 194 E ‖Mji j‖2 = ∑ k ‖ck‖ 2 σ2jk . The Arithmetic-Mean–Geometric-Mean ( AM/GM ) inequality sug- 195 gests that ∑ k ‖ck‖ 2 σ2jk ≥ n ( ∏ k ‖ck‖ 2 σ2jk ) 1/n = n ( ∏ k ‖ck‖ 2 ) 1/n exp ( −C ) , where the equal- 196 ity is achieved when ‖cm‖2 σ2jm = ‖cn‖ 2 σ2jn for any m 6= n. This suggests that latent variables 197 with high ‖ck‖2 mapping from zj to zi will have smaller variance . Hence , these latent variables in zj 198 are encouraged to stay in the “ active ” mode . On contrary , the latent variables that do not share sim- 199 ilar generative factors between zj and zi will be assigned larger variance , and being pushed towards 200 the “ passive ” mode . 201 Now we fix the optimal distribution of σ2j , and optimize for Mji . Since j ∼ N ( 0 , diag ( σ2j ) ) , this202 objective can be understood as optimally rotating the latent space of VAEj such that the stochastic203 component in Eq . 4 is minimized . Specifically , we have the following objective:204 min Mji ‖Mji j‖2 = min R ∥∥MjiRT j∥∥2 ( 7 ) where R is an orthogonal transformation . Let c′k be the kth column of MjiR . Similar as before , the205 AM/GM inequality suggests ∥∥MjiRT j∥∥2 = ∑k ‖c′k‖2 σ2jk ≥ ∏k ‖c′k‖2 exp ( −C3 ) . Hadamard ’ s206 inequality suggests that ∏ k ‖c′k‖ 2 ≥ |det ( MjiR ) | , and the equality is satisfied when c′ks are pair-207 wise orthogonal . This can be understood from the geometric perspective where ∏ k ‖c′k‖ 2 gives an208 upper bound on Volume ( { MjiR Tx : x ∈ [ 0 , 1 ] d } ) . As a result , the optimization of Mji will lead to209 an orthogonal transformation . Together the optimization of Eq . 6 and Eq . 7 encourages the align-210 ment of the “ polarized regime ” among different models under orthogonal linear transformations.211 They force different models in the ensemble to capture the same mixture of the generative factors.212 In the next section , we discuss the effect of the cross-model reconstruction in the VAE ensemble that213 encourages the disentangled representation over the entangled ones.214 4.2 THE EFFECT OF CROSS-MODEL RECONSTRUCTION215 In an entangled representation , each latent variable captures a mixture of generative factors in its216 unique way . Since different generative factors typically have different effects on data variations217 ( Duan et al. , 2019 ) , the orthogonal transformation from one entangled representation zj to another218 one zi introduces different encoding variance . Some of the transformed latent variables in Mjizj219 carry larger variance comparing to the corresponding ones in zi . This discrepancy leads to larger220 cross-model reconstruction of VAEi than the within model reconstruction . This error forces both221 VAEi and VAEj to adjust their representations until the effect on the data reconstruction by indi-222 vidual latent variables matches between Mjizj and zi . The process applies to all models in the223 ensemble and eventually leads to a one-to-one mapping of latent variables between different mod-224 els , where Mji becomes a trivial transformation ( signed permutation matrix ) . In particular , if one of225 the models in the VAE ensemble learns a disentangled representation , other models in the ensemble226 will converge to it . This is because the orthogonal transformation from an entangled representation227 to a disentangled representation introduces larger cross-model encoding variance due to the mixture228 of different generative factors in the former , thus a larger cross-model reconstruction by the disen-229 tangled model . On contrary , the orthogonal transformation from a disentangled representation to an230 entangled representation would not introduce larger cross-model encoding variance than the within231 model encoding , thus similar cross-model reconstruction as within model reconstruction by the en-232 tangled model . Such a gap encourages the entangled representations to align with the disentangled233 representation . We illustrate the geometric interpretation of such a case in Appendix C.234 From these discussions , we conclude that the VAE ensemble encourages different individual mod-235 els to capture similar generative factors , thus learn representations that are “ alike ” up to a trivial236 transformation . In the next section , we verify these analytic results with experiments.237 5 EXPERIMENTS238 Our experiments are designed to confirm the discussions in the previous sections . Particularly we239 ask the following questions : Q1 : Do the linear transformations in the ensemble converge to trivial240 transformation ? Q2 : Do the VAEs in the ensemble work in similar “ polarized ” regime ? Q3 : Does241 VAE ensemble improve the unsupervised disentangled representation learning , and what is the ef-242 fect of ensemble size ? Q4 : What are the effects of the cross-model reconstruction loss , the linear243 transformation loss and the hyperparameter γ in the VAE ensemble objective ? 244 We analyze the inner working of the proposed VAE ensemble using the benchmark dSprite dataset245 ( Matthey et al. , 2017 ) with fully known generative processes , and the real-world CelebA dataset ( Liu246 et al. , 2015 ) with unknown generative process . Furthermore , for dSprites dataset , we compare our247 results with the original VAE model and the state-of-the-art disentanglement VAE models including248 β-VAE ( Higgins et al. , 2017 ) , FactorVAE ( Kim & Mnih , 2018 ) , TC-VAE ( Chen et al. , 2018 ) and249 DIP-VAE ( Kumar et al. , 2017 ) . We use two widely used supervised metrics including FactorVAE250 metric ( Kim & Mnih , 2018 ) and DCI Disentanglement scores ( Eastwood & Williams , 2018 ) as the251 quantitative measurements . They are shown to correlate with other common supervised metrics 252 ( Locatello et al. , 2018 ) . For example , FactorVAE metric and β-VAE metric ( Higgins et al. , 2017 ) 253 capture similar notions , while DCI Disentanglement and Mutual Information Gap ( MIG ) ( Chen 254 et al. , 2018 ) capture similar notions . In addition , DCI Disentanglement is closely related to the 255 unsupervised model selection method UDR ( Duan et al. , 2019 ) . For CelebA dataset , we show the 256 latent traversal visulization as a qualitative measurement in Appendix E. We provide the details of 257 the experiments in Appendix D. 258 Q1 : We use the Distance to Orthogonality ( DtO ) ( Rolinek et al. , 2019 ) to check if the linear trans- 259 formations in the ensemble converge to a signed permutation matrix during training . DtO is the 260 Frobenius norm of the difference between a matrix M and its closest signed permutation matrix 261 P ( M ) . It is solved with mixed-integer linear programming ( MILP ) formulation . The details on 262 DtO can be found in Appendix B . In Figure 2 , we show the DtO estimation of the linear transfor- 263 mations in the VAE ensemble of different ensemble size for the dSprite dataset . We show the mean 264 and standard deviation of DtO across all linear transformations over 10 different runs . Furthermore , 265 we compare these results with a baseline measurement where DtO is estimated for the linear trans- 266 formations between the mean latent representations of well-trained individual models . Specifically , 267 we use ten well-trained individual models and report the mean and standard deviation of the DtO 268 estimations . As seen in the figure , the VAE ensemble models with different ensemble size all ap- 269 proach to trivial transformations between the individual models , while other VAE models do not 270 have such property . In Fig . 6 , we show that during training , the VAE ensemble remains maintains 271 low DtO while the original VAEs do not have such property . A similar result for models trained on 272 the CelebA dataset with different latent dimensions is shown in Fig . 7 . Further discussion on these 273 results are provided in Appendix E. 274 Q2 : To check if the models in the VAE ensemble work in similar “ polarized ” regime , we estimate 275 the relative error between LKL and L≈KL as ∆ = LKL−L≈KLLKL for each latent variable . Smaller ∆ 276 indicates closer matching between LKL and L≈KL of a latent variable , thus more “ active ” . Figure 277 3 ( a ) and 3 ( b ) show log ( ∆ ) of the 10 latent variables of individual models in VAE E2 and VAE E3 278 with different γ settings trained on the dSprite dataset . The results show that individual models 279 in the ensemble do work in similar ” polarized regime ’ . In Figure 3 ( a ) , we also compare the VAE 280 ensemble with the β-VAE where β = 4 . This setting was found previously to be the optimal setting 281 for the dSprite data for β-VAE ( Higgins et al. , 2017 ) . We see that the VAE ensemble encourages 282 more “ active ” latent variables than β-VAE . When we compare Fig . 3 ( a ) and 3 ( b ) , we see that as 283 the ensemble size increases , individual models are forced to have more “ active ” latent variables by 284 decomposing the generative factors . This can be observed in the latent traversals shown in Appendix 285 E. The dSprites dataset contains five ground truth generative factors . The VAE E2 models can 286 have up to eight “ active ” latent variables depending on input , and these representations capture a 287 decomposition of the ground truth generative factors . 288 Q3 : In Table 1 , we compare the disentangled representation performance between the proposed289 VAE ensemble and the state-of-the-art models . For the VAE ensemble , we report the performance290 of the first model in the ensemble . We also report the results for the VAE ensemble with different291 ensemble size and γ values . As shown in the table , the VAE ensemble significantly improves the292 performance over the original VAE model . In many settings , the VAE ensemble achieves similar293 or better performance over the state-of-the-art models . In Table 2 , we evaluate the consistency294 among the models in the ensemble by reporting the standard deviation of the evaluation metrics295 using different models in the same ensemble . The small values confirm that different models in296 the ensemble learn similar latent representations . Furthermore , Table 1 shows the joint effect of297 ensemble size and γ setting . When γ = 1 , the performance of VAE ensemble increases as the298 ensemble size increases , indicated by the higher mean and smaller variance of both the FactorVAE299 and DCI Disentanglement metrics . This behavior is consistent with the characteristic of ensemble300 learning where the increase in performance becomes smaller as the size of ensemble increases.301 However , as γ increases , having larger ensemble size can reduce the performance . We believe this302 is due to the increased difficulty of balancing between the cross-model and within model objectives 303 of VAE ensemble for larger ensembles . The reduced alignment of the latent representations among 304 different models can also be seen in Table 2 where difference in the performance among individual 305 models in the ensemble increases as ensemble size increases . 306 Q4 : We conduct the ablation study to further understand the effect of the linear transformation loss 307 and the cross-model reconstruction loss in the VAE ensemble objective . As shown in Fig . 4 , remov- 308 ing either component leads to a lower FactorVAE metric for the VAE ensemble . Without the linear 309 transformation loss , the performance of VAE ensemble decreases significantly across different en- 310 semble sizes . Without the cross-model reconstruction loss , the performance of VAE ensemble also 311 decreases but the gap becomes smaller as γ increases . This matches the discussion in Section 3 that 312 higher γ forces closer mapping between the encoders and reduce the cross-model reconstruction er- 313 ror of the decoders . However , this also reduces the effect of cross-model reconstruction as discussed 314 in Section 4.2 . A similar result is also found for the DCI Disentanglement metric as shown in Fig . 315 8 in Appendix E. Overall , the results from the ablation study confirms the importance of both the 316 linear transformation loss and the cross-model reconstruction loss in the VAE ensemble objective . 317 6 CONCLUSION 318 In this study , we propose a simple yet effective VAE ensemble framework consisting of multiple 319 original VAEs to learn disentangled representation . The individual models in the ensemble are 320 connected through linear layers that regularize both encoders and decoders to align the latent repre- 321 sentations to be similar up to a linear transformation . We show in theory that the regularization by 322 the VAE ensemble forces the linear transformations to be trivial transformations and show improved 323 performance on the unsupervised disentangled representation learning . The theoretical discussion 324 in Section 4 is based on the original VAE objective , and our experiments also focus on the ensemble 325 with original VAE . We believe such framework can be extended to other disentangled VAE models , 326 or even a mixture of different VAE models , as long as the regularization by the ensemble does not 327 conflict with the augmented objective of these models . 328 REFERENCES329 Leo Breiman . Bagging predictors . Machine learning , 24 ( 2 ) :123–140 , 1996.330 Christopher P Burgess , Irina Higgins , Arka Pal , Loic Matthey , Nick Watters , Guillaume Desjardins,331 and Alexander Lerchner . Understanding disentangling in β -VAE . ( NIPS ) , 2017.332 Tian Qi Chen , Xuechen Li , Roger Grosse , and David Duvenaud . Isolating Sources of Disentangle-333 ment in Variational Autoencoders . ( ICLR ) :1–13 , 2018.334 Harris Drucker , Corinna Cortes , Lawrence D Jackel , Yann LeCun , and Vladimir Vapnik . Boosting335 and other ensemble methods . Neural Computation , 6 ( 6 ) :1289–1301 , 1994.336 Sunny Duan , Loic Matthey , Andre Saraiva , Nick Watters , Chris Burgess , Alexander Lerchner , and337 Irina Higgins . Unsupervised model selection for variational disentangled representation learning.338 In International Conference on Learning Representations , 2019.339 Cian Eastwood and Christopher K. I. Williams . A framework for the quantitative evaluation of340 disentangled representations . ICLR , 2018.341 Irina Higgins , Loic Matthey , Arka Pal , Christopher Burgess , Xavier Glorot , Matthew Botvinick,342 Shakir Mohamed , and Alexander Lerchner . β-Vae : Learning Basic Visual Concepts With a Con-343 strained Variational Framework . ICLR 2017 , ( July ) :1–13 , 2017.344 Irina Higgins , David Amos , David Pfau , Sebastien Racaniere , Loic Matthey , Danilo Rezende , and345 Alexander Lerchner . Towards a Definition of Disentangled Representations . arXiv , pp . 1–29,346 2018.347 Aapo Hyvarinen and Hiroshi Morioka . Unsupervised feature extraction by time-contrastive learning348 and nonlinear ica . In Advances in Neural Information Processing Systems , pp . 3765–3773 , 2016.349 Aapo Hyvarinen , Hiroaki Sasaki , and Richard Turner . Nonlinear ica using auxiliary variables and350 generalized contrastive learning . In The 22nd International Conference on Artificial Intelligence351 and Statistics , pp . 859–868 , 2019.352 AJ Hyvarinen and Hiroshi Morioka . Nonlinear ica of temporally dependent stationary sources.353 Proceedings of Machine Learning Research , 2017.354 Ilyes Khemakhem , Diederik P Kingma , and Aapo Hyvärinen . Variational autoencoders and nonlin-355 ear ica : A unifying framework . arXiv preprint arXiv:1907.04809 , 2019.356 Hyunjik Kim and Andriy Mnih . Disentangling by factorising . Proceedings of Machine Learning357 Research , 80:2649–2658 , 10–15 Jul 2018.358 Diederik P. Kingma and Max Welling . Auto-encoding variational bayes . CoRR , abs/1312.6114,359 2013.360 Alex Krizhevsky , Ilya Sutskever , and Geoffrey E Hinton . Imagenet classification with deep convo-361 lutional neural networks . In Advances in neural information processing systems , pp . 1097–1105,362 2012.363 Abhishek Kumar , Prasanna Sattigeri , and Avinash Balakrishnan . Variational inference of disentan-364 gled latent concepts from unlabeled observations . CoRR , abs/1711.00848 , 2017.365 Ziwei Liu , Ping Luo , Xiaogang Wang , and Xiaoou Tang . Deep learning face attributes in the wild.366 In Proceedings of International Conference on Computer Vision ( ICCV ) , December 2015.367 Francesco Locatello , Stefan Bauer , Mario Lucic , Sylvain Gelly , Bernhard Schölkopf , and L G Dec.368 Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representa-369 tions . arXiv , pp . 1–33 , 2018.370 Francesco Locatello , Gabriele Abbati , Thomas Rainforth , Stefan Bauer , Bernhard Schölkopf , and371 Olivier Bachem . On the fairness of disentangled representations . In Advances in Neural Informa-372 tion Processing Systems , pp . 14584–14597 , 2019.373 Emile Mathieu , Tom Rainforth , N Siddharth , and Yee Whye Teh . Disentangling disentanglement in 374 variational autoencoders . arXiv preprint arXiv:1812.02833 , 2018 . 375 Loic Matthey , Irina Higgins , Demis Hassabis , and Alexander Lerchner . dsprites : Disentanglement 376 testing sprites dataset . https : //github.com/deepmind/dsprites-dataset/ , 2017 . 377 Ari Morcos , Maithra Raghu , and Samy Bengio . Insights on representational similarity in neural 378 networks with canonical correlation . In Advances in Neural Information Processing Systems , pp . 379 5727–5736 , 2018 . 380 Maithra Raghu , Justin Gilmer , Jason Yosinski , and Jascha Sohl-Dickstein . Svcca : Singular vector 381 canonical correlation analysis for deep learning dynamics and interpretability . In Advances in 382 Neural Information Processing Systems , pp . 6076–6085 , 2017 . 383 Danilo Jimenez Rezende , Shakir Mohamed , and Daan Wierstra . Stochastic backpropagation and 384 approximate inference in deep generative models . arXiv preprint arXiv:1401.4082 , 2014 . 385 Michal Rolinek , Dominik Zietlow , and Georg Martius . Variational autoencoders pursue pca direc- 386 tions ( by accident ) . In Proceedings of the IEEE Conference on Computer Vision and Pattern 387 Recognition , pp . 12406–12415 , 2019 . 388 Robert E Schapire . The strength of weak learnability . Machine learning , 5 ( 2 ) :197–227 , 1990 . 389 Karen Simonyan and Andrew Zisserman . Very deep convolutional networks for large-scale image 390 recognition . arXiv preprint arXiv:1409.1556 , 2014 . 391 Sjoerd van Steenkiste , Francesco Locatello , Jürgen Schmidhuber , and Olivier Bachem . Are disen- 392 tangled representations helpful for abstract visual reasoning ? In Advances in Neural Information 393 Processing Systems , pp . 14222–14235 , 2019 . 394 Liwei Wang , Lunjia Hu , Jiayuan Gu , Zhiqiang Hu , Yue Wu , Kun He , and John Hopcroft . Towards 395 understanding learning representations : To what extent do different neural networks learn the 396 same representation . In Advances in Neural Information Processing Systems , pp . 9584–9593 , 397 2018 . 398 Nicholas Watters , Loic Matthey , Matko Bosnjak , Christopher P Burgess , and Alexander Lerchner . 399 Cobra : Data-efficient model-based rl through unsupervised object discovery and curiosity-driven 400 exploration . arXiv preprint arXiv:1905.09275 , 2019 . 401 A TECHNICAL LEMMAS402 In this section , we give the lemmas used in the theoretical discussion in Section 4.403 Lemma 1 . ( Jensen ’ s inequality ) If g ( x ) is a convex transformation on x , then this convex transfor-404 mation of a mean g [ E ( x ) ] is less than or equal to the mean of the convex transformed value E [ g ( x ) ] ; 405 it is a simple corollary that the opposite is true of concave transformations.406 Lemma 2 . ( AM-GM inequality ) As an extension of Jensen ’ s inequality , given a list of non-negative407 real numbers x1 , x2 , . . . , xn , the arithmetic mean of this list 1n ∑n i=1 xi is greater than or equal408 to the geometric mean of the same list ( ∏n i=1 xi ) 1 n ; and further , the equality holds if and only if409 x1 = x2 = · · · = xn.410 Lemma 3 . ( Hadamard ’ s inequality ) . if M is the matrix having columns ci , then |det ( M ) | ≤411 ∏n i=1 ‖ci‖ ; and the equality in Hadamard ’ s inequality is achieved if and only if the vectors are412 orthogonal.413 B DISTANCE TO ORTHOGONALITY ( DTO ) 414 In this section , we introduce the detail of Distance to Orthogonality ( DtO ) that is used in our exper-415 iment to check if the linear transformations in the VAE ensemble approach trivial transformations.416 This measurement is also used in ( Rolinek et al. , 2019 ) for a similar purpose . DtO is the Frobe-417 nius norm of the difference between a square matrix M and its closest signed permutation matrix418 P ( M ) . Finding P ( M ) can be formulated as a mixed-integer linear programming ( MILP ) problem419 as following:420 min P ∑ i , j |Mi , j − P ( M ) i , j | s.t . P ( M ) i , j ∈ { −1 , 0 , 1 } , ∀ ( i , j ) ∑ i |Pi , j | = 1 , ∀j∑ j |Pi , j | = 1 , ∀i ( 8 ) By introducing new variables P+i , j , P − i , j ∈ { 0 , 1 } and Di , j = |Mi , j − P ( M ) i , j | , we can reformulate421 the above optimization problem as:422 min P ∑ i , j Di , j s.t . ( P+i , j − P − i , j ) −Mi , j ≤ Di , j , ∀ ( i , j ) Mi , j − ( P+i , j − P − i , j ) ≤ Di , j , ∀ ( i , j ) ∑ i ( P+i , j + P − i , j ) = 1 , ∀j∑ j ( P+i , j + P − i , j ) = 1 , ∀i ( 9 ) Using this optimization formulation , DoT of a given matrix M ∈ Rn×n is defined as:423 DoT = 1 n2 ∑ i , j |Mi , j − P ( M ) i , j | ( 10 ) C GEOMETRIC INTERPRETATION OF THE EFFECT OF CROSS-MODEL 424 RECONSTRUCTION 425 Given a disentangled and entangled latent representation space , Fig . 5 illustrates the effect of the 426 cross-model reconstruction by VAE ensemble . The left part shows the orthogonal transformation 427 from a disentangled representation to an entangled space , and the right part shows the transfor- 428 mation in the opposite direction . As shown in the figure , the orthogonal transformation from the 429 disentangled representation to the entangled space does not introduce larger variance than the en- 430 tangled representation . Hence , we can expect similar cross-model reconstruction and within model 431 reconstruction . However , the transformation from the entangled representation to the disentangled 432 space introduces larger variance ( yellow shaded area over blue area on the right ) than the disentan- 433 gled representation . This leads to larger cross-model reconstruction by the disentangled model . 434 D MODEL ARCHITECTURE AND TRAINING DETAILS 435 We conducted our experiments , including training and evaluating the current state-of-the- 436 art disentanglement models as well as evaluating the proposed VAE ensembles , using the 437 disentanglement lib 1 open-source library ( Locatello et al. , 2018 ) . 438 Table 3 shows the encoder and the decoder architecture of the VAE model used in our experiments . 439 This architecture is the same as the one used in the original β-VAE Higgins et al . ( 2017 ) . 440 Table 4 shows the hyperparameters setting used throughout the experiments . These parameters are 441 fixed for all the experiments . 442 E ADDITIONAL EXPERIMENTAL RESULTS 443 In this section , we present the additional results including the DtO and “ polarized regime ” analysis 444 on the models trained on the CelebA dataset similar to the ones conducted on dSprite dataset in 445 Section 5 ; the ablation results with DCI-Disentanglement metric and the DtO estimation ; and the 446 1https : //github.com/google-research/disentanglement_lib latent traversal of the trained model on both dSprite and CelebA dataset along with further discuss447 the effect of VAE ensemble on the latent representation.448 E.1 CHARACTERIZATION OF THE LINEAR TRANSFORMATION IN VAE ENSEMBLE449 In Figure 6 , we show the DtO estimation of the linear transformations in the ensemble during train-450 ing for the dSprite dataset . We report the mean and standard deviation of DtO across all linear451 transformations over 10 different runs . Furthermore , we compare these results with a VAE baseline452 where DtO is estimated for the linear transformations between original VAEs . Specifically , we train453 ten different VAEs separately and estimate the DtO of the pairwise linear transformation among454 these models during training . Similarly we report the mean and standard deviation of these DtO es-455 timations . As seen in the figure , the VAE ensemble models with different ensemble size all approach456 to trivial transformations between the individual models , while the original VAEs do not have such457 property . A similar result is also found in models trained for CelebA dataset . Similar to the results458 in Figure 6 , we observe decreased DtO of the linear transformations in the VAE ensemble during459 training.460 We also compare models trained with different latent dimension size . We observe decreased DtO461 as the latent dimension of the model increases in Figure 7 . This is because , as discussed in the462 main paper , the VAE ensemble encourages more “ active ” latent variables . Models with higher latent463 dimension likely to learn a decomposition of generative factors . As a result , the alignment of the464 latent variable between different models are easier thus the linear transformations between the latent465 representations is closer to the trivial transformation . On the contrary when there are less latent466 variables in the model than the generative factors , some of the latent variables will capture more467 than a single generative factor . As a result , the one-to-one mapping between the latent variables of468 different models will not lead to a trivial transformation.469 In Figure 9 and Figure 10 , we show the “ polarized regime ” estimation for models in VAE E2 and 470 VAE E3 trained for CelebA datasets , respectively . Similar to the results in Figure 3 , individual 471 models in the VAE ensemble tend to have similar ‘ polarized regime ” , and higher γ enforces the “ po- 472 larized regime ” by separating “ passive ” latent variables from the “ active ” ones . When we compare 473 between VAE E2 and VAE E3 , we observe increased “ active ” latent variables similar to the result 474 on dSprite dataset in Section 5 . More importantly , as discussed earlier , latent variables in a model 475 with limited latent dimensions need to capture more than a single generative factor , especially for 476 a complicated real-world dataset such as CelebA . This makes the linear transformation between the 477 latent representations less trivial . As the latent dimension size grows , such constraint is relaxed and 478 the linear transformations are closer to trivial . 479 These additional results confirm the conclusion in Section 5 : ( 1 ) as the ensemble size increases , 480 DtO increases due to the difficulty of aligning the latent representations among different models ; 481 ( 2 ) as the model latent dimension increases , DtO decreases due to the increased model capacity , and 482 encourages the one-to-one mapping between latent variables in different models ; ( 3 ) hyperparameter 483 γ does not affect DtO significantly , but plays an important role on separating “ active ” and “ passive ” 484 latent variables , especially when the latent dimension is large enough . 485 Furthermore , we believe the DtO measurement of the linear transformation in VAE ensemble could 486 be a useful indicator for latent dimension size . As shown in Figure 7 and Figure 9 , when the latent 487 dimension is sufficient for a given dataset , the DtO of the linear transformation is small and some 488 latent variables are pushed to “ passive ” mode . 489 E.2 ABLATION STUDY 490 Similar to the ablation result shown in Section 5 , here we show the same ablation study using the DCI 491 Disentanglement metric in Fig . 8 ( a ) as well as the DtO measurement 8 ( b ) . Similar as the results of 492 the FactorVAE metric in Fig . 4 , removing either component leads to a lower DCI Disentanglement 493 metric for the VAE ensemble . Without the linear transformation loss , the performance of VAE 494 ensemble decreases significantly across different ensemble sizes . Fig . 8 ( b ) shows that for VAE 495 ensemble without the cross-model reconstruction , the linear transformations among models are close496 to a trivial transformation ( signed permutation ) . This implies the orthonormal transformation of the497 linear transformations . This result further supports our intuitive justification in Appendix C that the498 cross-model objective encourages entangled models to align to disentangled models . Indeed , we see499 that adding the cross-model reconstruction can further reduce the DtO of the linear transformations500 among the models in the ensemble.501 E.3 LATENT TRAVERSAL502 In this section we show the latent traversal of models trained on both dSprites and CelebA datasets.503 For a fixed input image , to extract the latent traversal we change the value of a single latent variable504 zi in the corresponding encoding , and observe the generated output image to understand the effect505 of zi . The range of the value are usually chosen to be from -3 to 3 due to the standard Gaussian506 prior.507 In Figure 11 , we show the latent traversal for both VAE E2 and a single VAE model with 10 latent508 dimensions trained on dSprites dataset . Three images as shown in the last column of each block509 are used as input . Both models are able to capture certain generative factors of the data including 510 position , shape , rotation and scale . In Section 5 , we argue that the representation by VAE ensemble 511 encourages more “ active ” latent variables , thus can capture a decomposition of the ground truth gen- 512 erative factors . Especially from the “ polarized regime ” estimation in Figure 3 , we observe that some 513 latent variables in the VAE ensemble are in-between “ active ” and “ passive ” modes . This suggests 514 that the VAE ensemble model generates input-dependent factors based on the input complexity . In 515 Figure 11 , we observe this behavior highlighted with color boxes . The traversal on the second la- 516 tent variable z2 shows that an ellipse shape does not lead to an “ active ” latent variable . However , 517 both heart and square shape lead to an “ active ” latent variable that changes the output . In contrast , 518 the single VAE model does not have such behavior where the “ active ” modes are consistent across 519 different input data . 520 In Figure 12 and Figure 13 , we show the latent traversal for both VAE E2 and a single VAE model 521 with 16 latent dimensions trained on CelebA dataset , respectively . In this real-world dataset , the 522 generative factors are unknown . We observe different factors including background , azimuth , gen- 523 der , hair style being captured by both models . Similar as before , the single VAE model maintains 524 similar “ active ” mode for all latent variables where similar traversal patterns are observed for both 525 input images . However , VAE E2 shows semantically consistent but input-dependent “ active ” mode . 526 This is translated into different traversal effects and more realistic and sharper images by VAE E2 , 527 especially for the first input image that is less common in the dataset . We believe this is important 528 towards a meaningful compositional latent representation learning . 529 Overall , the latent traversal results in this section confirm the findings on the inner working of the 530 VAE ensemble shown in the previous section as well as the discussion in Section 5 . 531

Summary: This paper proposes a simple and effective technique to improve disentanglement by coupling the latent spaces of different VAE models. It builds on Duan et al. (2019)’s proposed method to rank the representations of different models. By learning a VAE ensemble with linear transformations between the latent spaces and an additional “cross-model” reconstruction loss, the authors show that they can achieve significantly better disentangling.