Unsupervised Disentanglement with Tensor Product Representations on the Torus

1 INTRODUCTION . Unsupervised learning methods of disentangled representations attempt to recover the explanatory factors z of variation in the data . The recovered representation , c , is expected to be : ( i ) disentangled , i.e. , each element in c should vary one of the generative factors in z ( ii ) complete in the sense that all generative factors z can be controlled by the latent representation c , and ( iii ) informative , that is the ability to relate via a low capacity model ( e.g. , a linear model ) the generative factors z and the latent representation c. It has been shown that solving this task without further assumptions is impossible , since there is an infinite number of bijective functions f that map between c and z = f ( c ) ( Locatello et al. , 2019 ) with all corresponding generative models having the same marginal distributions of the observations . In this work , we propose to use representations that are a tensor product of elements , which take values on unit circles S1 , claiming that this structure is suitable for the effective recovery of the underlying modes of variation . This claim follows from entropy considerations . One wishes to have a representation that has a low entropy of entanglement . The entropy of entanglement is measured by evaluating the number of non-zero eigenvalues of the Schmidt decomposition of the representation . If a representation can be described by only one term , which is an outer product of the orthogonal basis functions , it has the lowest possible entanglement entropy , zero . If a representation requires the use of more than one term in the decomposition , it is entangled and its entanglement entropy can be quantified , e.g . by the Von Neumann entropy constructed from non-zero eigenvalues . The tensor product of n unit circles , S1 , takes the shape of an n-torus Tn = ( S1 ) n , and has a low entanglement property . This representation has the advantage that it can capture the periodicity structure of the generative factors , where each circle controls a different aspect of the original distribution . Unlike other generative models , which rely on the Gaussian distribution as a prior , Tn is a compact manifold , and therefore any function that acts on it , such as a decoder , has to interpolate only , but not extrapolate when generating new instances . In this work , we present the TD-VAE , a Variational Auto Encoder whose latent space resides on the TD torus manifold . In an extensive set of experiments with four different datasets , we compare the torus latent representations with others proposed in the literature . We show a clear advantage of the torus representation in terms of various measures of disentanglement , completeness , informativeness and propose a new metric , the DC-score , which assesses the combined disentanglement and com- pleteness performance . We present a detailed quantitative and qualitative analysis that supports our method . 2 RELATED WORK . Generative models aim to create new instances that are indistinguishable from a given distribution . Auto encoders ( Kramer , 1991 ) achieve this by constructing the identity function using a composition of two components , an encoder and a decoder . The encoder is tasked with compressing the input into a latent representation with a dimension smaller than the input ’ s , whereas the decoder is responsible for reconstructing the original input from the latent space . Such auto encoders usually overfit to their training data and are very limited when tasked with generating new samples . Instead of mapping each input instance to a constant vector , Variational Auto Encoders ( Kingma & Welling , 2013 ) ( VAE ) map each instance to a pre-defined distribution , by minimizing both the encoder-decoder reconstruction loss and a KL-divergence term that depends on the prior distribution . VAE ’ s are capable of generating new instances by sampling a latent vector from the predefined distribution ; however , the interpretation of latent space components is typically obscure and the mapping between them and the dataset properties may be complex . In order to increase the latent space interoperability , various modifications of the original VAE were introduced . These variants try to achieve a factorization of the latent space components with each component corresponding to a specific simple characteristic of the dataset . In β-VAE ( Higgins et al. , 2016 ) , the weight of the KL-divergence term is increased w.r.t . the reconstruction term , which results in better factorization of the latent space . However , this introduces a new hyper-parameter and departs from the theoretical derivation of the ELBO term . For example , large β parameters may prevent the conditional distribution from modeling the data faithfully . DIP-VAE ( Kumar et al. , 2017 ) introduces a regularizer that constrains the covariance matrix of the posterior distribution . This term is added to the original VAE loss function , along with two hyper-parameters , and helps achieve better factorization . Factor-VAE ( Kim & Mnih , 2018 ) adds a Total Correlation ( TC ) approximation penalty over the produced latent codes and introduces an additional hyper-parameter . 2.1 DISENTANGLEMENT METRICS . In order to evaluate the expressiveness of different generative models , a plethora of disentanglement definitions and metrics were suggested in the literature ( Do & Tran , 2019 ; Eastwood & Williams , 2018 ) , see Locatello et al . ( 2019 ) for a comprehensive review . In this work , we adopt the definitions and metrics introduced in Eastwood & Williams ( 2018 ) which provides a successful set of disentanglement metrics . The authors of Eastwood & Williams ( 2018 ) introduce three metrics , Disentanglement , Completeness and Informativeness ( DCI ) , to quantify the disentanglement properties of the latent space and its ability to characterize the generative factors of a dataset . Given a dataset that was generated using a set of K generative factors , ~z ∈ RK , and that it is required to learn latent codes , ~c ∈ RD of D-dimensional representation . Ideally , in an interpretable representation , each generative factor zi would correspond to only one latent code ca . It is also beneficial if the mapping is linear , as the correlation between generative factors and codes can then be easily assessed . Generating new instances that are indistinguishable from the original dataset distribution requires that the latent codes cover the whole range of the generative factors . Furthermore , such a representation provides the ability to modify a specific property of a generated instance by directly tuning the corresponding latent code . The DCI metrics aim to quantify the relationship between codes and generating factors by having a single number that characterizes the relative importance of each code,1 ca , in predicting a factor , zi , which in turn defines an importance matrix Rai . To construct the importance matrix , K regressors are trained to find a mapping between zi and ~c , ẑi = fi ( ~c ) . In this work , we follow Eastwood & Williams ( 2018 ) and infer the importance matrix using a lasso linear regressor ’ s weights , Wia , by Ria = |Wia| . 1throughout the paper we use a , b , c , . . . and i , j , k , . . . letters for indices of the codes ~c and the factors ~z components respectively . Once the importance matrix Rai is obtained , the DCI metrics can be defined explicitly . The disentanglement is given by D = rank ( R ) K D∑ a=1 ρaDa , ( 1 ) where Da = 1−HK ( Pa ) = 1 + K∑ k=1 Pak logK Pak , Paj = Raj∑K k=1Rak , ρa = ∑ j Raj∑ bk Rbk . ( 2 ) High disentanglement means that each entry of the latent vector ~c corresponds to only one element in ~z ( in a linear sense ) , that is , each code element , ca , affects one generating factor . The disentanglement metric defined above differs from Eastwood & Williams ( 2018 ) by a correction factor rank ( R ) K . When the rank of R is equal to the generative factor ’ s dimension , this correction equals 1 , and does not affect the metric . However , it does make a difference when the number of codes is smaller than what is needed to account for all generating factors . Typically , one assumes that there are at least as many expressive codes as the number of factors . When this is not the case , it means that even if we have disentanglement among the expressive codes , their number is not sufficient , and our correction accounts for this . Note that the correction has no influence if there are irrelevant code elements in addition to a sufficient number of expressive ones . The ρ factors handle cases where some of the code dependence on the factors is very weak ( namely irrelevant codes ) , while other codes depend mainly on one factor and hence should not be of equal importance . The completeness is defined by C = 1 K K∑ j=1 Cj , Cj = 1−HD ( P̃j ) = 1 + D∑ a=1 P̃aj logD P̃aj , P̃aj = Raj∑D b=1Rbj . ( 3 ) In contrast to the disentanglement metric , which considers a weighted mean over the generative factors , motivated by the fact that irrelevant units in ~c should be ignored , here all codes are treated equally . A situation in which each factor zi is explained by only one element of the code ca results in high completeness . The informativeness is the MSE between the ground-truth factors , z , and the predicted values , fj ( ~c ) , I = 1 K K∑ j=1 Edataset [ |zj − fj ( ~c ) |2 ] . ( 4 ) Disentanglement and completeness values range between zero and one , where one is the best value , while the best value for informativeness is zero . Disentanglement and completeness alone do not provide sufficient information on whether a representation factorizes properly . For example , consider the case of a representation where only one generative factor is described by all the codes . While this representation is completely disentangled , it is not complete , and is therefore meaningless . In order to have a meaningful representation , both disentanglement and completeness need to be high . Thus , we introduce a new score , called the DC-score , which accounts for both disentanglement and completeness . We define the DC-score as the geometric mean of the two metrics , DC-score = √ DC . This way , only cases where both scores are high will result in a high DC-score . Furthermore , the score will favor cases where both disentanglement and completeness are comparable rather than having different values while having the same arithmetic mean . 3 METHOD . Assume a pre-determined set of vectors sampled from an unknown distribution χ , each vector containing K independent factors , ~z ∈ RK . The training dataset contains samples xI = F ( ~zI ) , I = 1 , . . . , N that are generated using a generative function F . In unsupervised disentanglement , the learner has no access to ~z nor to F , and receives only the set of samples , { xI } NI=1 . The learner has to recover a set of representation vectors ~cI that are linked to ~zI in a way that is bijective and disentangled , see Section 2.1 . Furthermore , in order to generate new samples , it is also necessary to be able to sample new representation vectors ~cnew , and to obtain a generative function G such that the new generated samples are from the same distribution of F ( ~z ) , where ~z ∼ χ . In this work , we view the latent representation vector ~c as the angles associated with a list of D two-dimensional unit vectors ma ∈ R2 , 1 ≤ a ≤ D , i.e. , ∀a , ‖ma‖ = 1 . Using the two-dimensional vectors we define vprod ≡ vec ( vα1 ... αDprod ) , vα1 ... αDprod = m α1 1 ⊗ · · · ⊗mαDD , vorient ≡ ( m01 , . . . , m 0 D ) , ( 5 ) where ⊗ is the outer product operator , αa ∈ { 0 , 1 } and vec is the vectorization operation ( equivalent to flattening the tensor ) . Next , define the operator V , that given m1 , . . . , mD , concatenates vprod together with vorient , V ( m1 , . . . , mD ) = [ vprod ; vorient ] . The vector v = V ( m1 , . . . , mD ) resides in a vector subspace of R2D+D , defined by only a set of D parameters . The additional D elements of vorient are required to ensure that the mapping V is bijective . For example , this can be seen by examining the case of D = 2 : let m1 ≡ ( cos θ1 , sin θ1 ) , m2 ≡ ( cos θ2 , sin θ2 ) , m′1 ≡ −m1 and m′2 ≡ −m2 . Then m1 ⊗m2 = m′1 ⊗m′2 . A natural way of acquiring a random point on the circle , S1 , is by sampling two independent Gaussians , m̂αka ∼ N ( µαka , σαka ) . Each tuple of these vectors is then normalized to have a unit norm , mαka = m̂αka√ ( m̂0a ) 2 + ( m̂1a ) 2 . ( 6 ) Assume that the elements , m̂αka , follow the normal distribution with a zero mean and a standard deviation of 1 , m̂αka ∼ N ( 0 , 1 ) , then the vectors ma follow the uniform distribution on S1 . For the purpose of obtaining the distribution parameters of m̂αka , the encoder e is applied on an instance , x , e ( x ) = [ µ01 µ11 , σ01 σ11 . . . µ0D µ1D , σ0D σ1D ] . ( 7 ) The reparametization trick ( Kingma & Welling , 2013 ) is then used to obtain a set of normally distributed vectors . Denote by S the sampling operator , we sample the coding vector as ~MI ≡ { ( m01 m11 ) , . . . , ( m0D m1D ) } = S ( e ( xI ) ) , ( 8 ) where each component , ( a , αa ) of m , a = 1 . . . D , α = 0 , 1 is sampled from N ( µαa , σαa ) . The decoder G then acts on V ( m1 , . . . , mD ) to generate a new instance , x̃ . The requirement that the generated instance is indistinguishable from the original distribution F ( χ ) translates to minimizing the ELBO loss function . First , it maximizes the log-likelihood probability of x̃ to be similar to the training sample it reconstructs by having an l2 loss term . Second , to encourage a uniform distribution on the torus , it contains a KL-divergence term with respect to N ( 0,1D ) . The overall expression for the loss function used during the optimization of the networks e and G is L = ∑ I [ ∣∣∣G ( V ( ~MI ) ) − xI ∣∣∣ 2 − βDKL ( p ( m̂ ) ||rI ) ] , ( 9 ) where β is a hyperparameter introduced in β-VAE ( Higgins et al. , 2016 ) and rI ∼ N ( 0,1D ) . As the latent representation resides on the torus , sub-vectors 2D components of ma may be described by D angles , θa ∈ [ 0 , 2π ] ( the angles θa are identified with the codes ca ) . In order to generate new instances , identify , m0a = cos θa , m 1 a = sin θa . and apply the decoder G ( V ( m ) ) to these elements . The torus topology , TD , is both compact and periodic . Since it is compact , for every sampled θa , there is a nearby θ′a from the training set , thus , the network G can be viewed as an interpolator . This is in contrast to vectors sampled over RD , in which G has to extrapolate . Furthermore , being periodic , this topology is able to exploit a periodic structure in the generative factors . Consider a common generative factor associated with rotations . A rotation can not be represented by a single non-compact variable ; therefore encoding a rotation in a latent space requires the entanglement of two components . However , on the torus , only one compact dimension can be used to identify this generative factor .

Summary: This paper proposes a new autoencoder architecture that achieves better recovery of the latent factor structure in artificial datasets according to previously established metrics: disentanglement, completeness, and informativeness.  The key idea is to nonlinearly project a small number of latent factors into a higher-dimensional space, based on the topology of the torus, in such a way that prevents an arbitrary linear rotation of the latent factors from yielding an equivalent representation.  The approach is compared against existing approaches for variational autoencoders focused on learning disentangled representation, and achieves higher disentanglement and completeness, as well as competitive or best informativeness (reconstruction) on a majority of datasets considered.