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: The authors propose an autoencoder where the latent space is defined on a torus. A D-dimensional torus is represented by the tensor product of D unit circles. They argue the torus induces disentanglement as the analogous of entropy entanglement in quantum physics. Using 5 datasets, they show both the reconstruction performance and disentanglement score (Eastwood & Williams, 2018) are high.