ProtoVAE: Using Prototypical Networks for Unsupervised Disentanglement
Authors: Vaishnavi S Patil
Matthew Evanusa
Joseph JaJa
Venue: Submitted to ICLR 2023 (2023)
Keywords: Unsupervised Learning
Disentangled Representations

Abstract: Generative modeling and self-supervised learning have in recent years made great strides towards learning from data in a completely \emph{unsupervised} way. There is still, however, an open area of investigation into guiding the neural network to learn useful or good representations. The problem of unsupervised \textit{Disentanglement} is of particular importance as it offers to learn interpretable representations, with disjoint subsets of the representation encoding different, meaningful factors of variation. Recent work has theoretically grounded the factors of variation, via the lens of group theory, as disentangled actions of the symmetry subgroups which transform only the correspond subspaces of the disentangled representation. We use this mathematical formalism instead to impose constraints on the representations learned by a unsupervised generative neural network, such that transformations of the representation correspond to the actions of a unique symmetry subgroup. To this end, we introduce a novel model, ProtoVAE, that leverages a deep metric learning Prototypical network trained via self-supervision to constrain the latent space of a Variational Autoencoder to decompose into independent subspaces. Further, we actively change or \textit{intervene} in the latent space during training to enforce each dimension of the representation to uniquely and consistently transform the data corresponding to some symmetry subgroup. We demonstrate and evaluate our proposed model on the benchmark DSprites and 3DShapes datasets and compare with other state of the art disentanglement methods via qualitative traversals in the latent space, as well as quantitative disentanglement metrics. We further qualitatively demonstrate the effectiveness of our model on the real-world datasets CelebA which consistently encodes the different factors.

Under review as a conference paper at ICLR 2023
PROTOVAE: USING PROTOTYPICAL NETWORKS FOR UNSUPERVISED DISENTANGLEMENT
Anonymous authors Paper under double-blind review
ABSTRACT
Generative modeling and self- supervised learning have in recent years made great strides towards learning from data in a completely unsupervised way. There is still however an open area of investigation into guiding neural network to encode the data into representations that are interpretable or explainable. The problem of un- supervised disentanglement is of particular importance it proposes to discover the different latent factors of variation or semantic concepts from the data alone, without labeled examples, and encode them into structurally disjoint latent rep- resentations. Without additional constraints or inductive biases placed in the net- work, generative model may leam the data distribution and encode the factors, but not necessarily in disentangled way. Here, we introduce a novel deep gen- erative VAE-based model, Proto ProtoVAE, that leverages deep metric learning Pro- totypical network trained using self-supervision to impose these constraints. The prototypical network constrains the mapping of the representation space to data space to ensure that controlled changes in the representation space are mapped to changes in the factors of variations the data space. Our model completely unsupervised and requires a priori knowledge of the dataset, including the number of factors. We evaluate our proposed model on the benchmark dSprites, 3DShapes, and MPI3D disentanglement datasets, showing state of the art results against previous methods via qualitative traversals in the latent space, as well as quantitative disentanglement metrics. We further qualitatively demonstrate the effectiveness of our model on the real-world CelebA dataset.
INTRODUCTION
One theory of the success of deep learning models for supervised learning revolves around their ability to learn mappings from the input space to a lower dimensional abstract representation space which are best predictive of the corresponding labels Tishby & Zaslavsky (2015). However, for the models to be robust to noise and adversarial examples, be transferable to different domains and distributions and interpretable, we need to impose additional constraints on the learning paradigm. As promising solution this, the models can be encouraged to focus all the latent "distinctive properties" of the data distribution and encode them into representation for downstream super- vised tasks. These latent distinctive properties or factors of variations are the interpretable abstract concepts that describe the data.
The intuitive notion of disentanglement, first proposed in Bengio (013), proposes to discover all the different factors of variations from the data, and encode each factor in a separate subspace or dimension of the learned latent representation. These disentangled representations are not only interpretable and give valuable insights into the data distribution but are also more robust for multiple downstream tasks Bengio (203)): [Schoekkopt et al. (2012) which might depend only on a subset of factors Suter et al. (2019).
The problem of learning these disentangled representations in a completely unsupervised way is particularly challenging we do not have access to the ground truth labels of factors nor are privy to the true number of factors or their nature. Recent works have proposed to solve this problem by training generative networks to effectively model the data distribution and in turn the factors of variations. From this generative perspective disentanglement, higher dimensional data assumed to be non-linear mapping of these factors of variation, where each factor assumes different values

---

Under review as a conference paper at ICLR 2023
to generate specific examples in the data distribution. Locatello et al. (2119) intuitively characterizes representations which encode the factors as disentangled if a change in a single underlying factor of variation in the data produces change in a single factor of the learned representation (or change in the subspace of the representation that encodes that factor). Conversely, from the generative perspective, for representation to be disentangled, change in single subspace of the learned representation, when mapped to the data space, must produce a change in a single factor of variation.
For this generative mapping between changes in the representation space to the changes in the fac- tors of variations (in the data space) to be injective, we propose constrains on the changes in the factors of variations for pre-determined changes in the representation space. Each separate subspace of the representation, when changed, must map to a change in unique factor of variation which in turn encourages information about the different factors to be encoded in separate subspaces of the representation. Moreover, each separate subspace must consistently map to change in a single factor throughout the subspace range. This encourages the different subspaces of the representation to encode information only about single factor of variation. The recent work of et al. (2021) also demonstrated empirically that the concept of local isometry was a good inductive bias for unsu- pervised disentanglement, and it can aid generative models in discovering "natural" decomposition of data into factors of variation. This local isometry constraint on the mapping enforces the changes in the data space to be proportional any changes made in the representation space. In order to effectively impose the above constraints in an unsupervised manner, we turn towards deep metric learning.
In recent years, metric learning has emerged as a powerful unsupervised learning paradigm for deep neural networks, in conjunction with self-supervised data augmentation. One of the more successful metric learning models, Prototypical Networks, projects the data into a new metric space where examples from the same class cluster around a prototype representation of the class and away from the prototypes of other classes. We use this ability of the network to cluster the different changes in the data space mapped by the corresponding changes in the representation space and thereby enforce the above described constraints.
We develop a novel deep generative model, Proto VAE, consisting of a Prototypical Network and Variational Autoencoder network (VAE). The VAE acts as the generative component, while the Prototypical Network guides the VAE separating out the representation space by imposing the constraints for disentanglementt To learn these representations in an unsupervised way, the proto- typical network needs labeled data for clustering, we train the prototypical network using generated self-supervised datasets. To produce the self supervised dataset, representation space, which change individual elements of the latent space and map the intervened representations to the data space. Owing to the self-supervised training, our model is able disen- langle without any explicit prior knowledge of the data, including the number of desired factors.
perform interventions in the
In this work, our core contributions are:
We design a self-supervised data generation mechanism using a VAE that creates new sam- ples via process of intervention to train a metric-learning prototypical network.
We design and implement a novel model, Pro,oVAE, which combines a VAE and proto- typical network to perform disentanglement without any prior knowledge of the underlying data.
We empirically evaluate VAE on standard benchmark DSprites, 3DShapes, MPI3D, and CelebA datasets, showing state of the art results.
2 PROTOVAE
Our proposed model consists of a V Welling (2014); Rezende al. (2014) as the base generative model (Section 2.1). The VAE consists of an inference network which encodes the data into lower dimensional latent representations and a generator network that maps the representations back into the data space. To implicitly encourage the inference network to encode disentangled rep- resentations, we impose constraints on the generative mapping from changes in the representation space to changes in the factors of variations in the data space. This generative mapping is deter- mined by both the generator and the inference networks. To generate self-supervised data for the prototypical network, perform interventions (Sec 2.2h which changes individual dimensions of
2

---

Under review as a conference paper at ICLR 2023
the representation. Given batch of latent representations encoded by the inference network, first intervene on dimension of the representation by changing its value to the value of another rep- resentation from the batch for the same dimension. The original representations and the intervened representations are then mapped into the data space by the generator network and concatenated to form a pair of original data and generated data from interventions. Given that the original and the in- tervened representations differ in single dimension, the generative mapping should be constrained ensure that the corresponding pair of original and generated data differs only single factor of variation.
This constraint is enforced using a Prototypical network (Appendix A.2) which based on the idea that there exist an embedding in which examples from the same class cluster around prototype representation for that class. Our proposed prototypical network (Section 2.3) takes as input pairs of data generated by the self-supervised process described above, and maps these pairs of data into metric space which pairs generated by intervening on the same dimension cluster together. These clusters which are identifiable with intervening dimensions in-turn become identifiable with the factors variation that differ in value between the pair when a dimension is intervened upon. We further augment the prototypical network with separate output head that enforces local isometry, by predicting the difference in the value of the intervened dimensions from the pair in the data space. F the diagram overview of the complete model.
Lastly, for the intervened representations to be mapped into the data space such that only a factor of variation is changed, we constrain the generated data to lie in the true data distribution. This con- straint can be effectively enforced in the representation space by minimizing the distance between the distribution of the original representations of the inference network and the intervened represen- tations such that the generator network maps both the distributions to the true data distribution. We do so by training a discriminator network (Section 2.4p in the representation space to distinguish between the original and the intervened representations. The inference network which generates the original representations then trained to fool the discriminator thus effectively bridging the distance between the distributions.
2.1 VARIATIONAL AUTOENCODER
The base generative model which consists of an inference network go RDD - Rd that encodes the data to lower dimensional representation z and a generator network po : .Rd - RDD which reconstructs the data 3 from the representations z. The inference and the generator network are trained together to maximize the evidence lower bound (ELBO) of the data log- ligeliil as in equation 1 max Lv (0, (6 = (ele po (=<((ll)))) 0,6
(1)
Maximizing the first term of equation I Jensures that the latent representation encodes all the infor- mation needed to faithfully reconstruct the data from the representation alone. This ensures that the representations encode all the different factors of variations in the data. The KL divergence term creates information bottleneck which enforces optimal, compact encoding of the data by enfore- ing the posterior distribution to be similar to the independent, non- non-informative prior distribution. For more details please refer to Appendix A.
2.2 SELF-SUPERVISED DATA GENERATION
The prototypical network works to cluster changes in the factors of variations in the data space and provides gradients the generator and inference network to better separate out the factors in the representation space. To do so, the prototypical network requires a set of supervised examples, called the support set, from each class to compute a prototype around which examples from the same class cluster. Furthermore, supervised query set is required to compute the distance of query examples from the target prototypes and the subsequent loss is used to update the network. For learning disentangled representations in an unsupervised way, we propose generate these support and query sets using self-supervision. We describe the full algorithm in Appendix B.I.
Given batch of data z, first use the inference network to encode the data into the representation space E Rd. Following Suter et al. (2019), we define interventions as the act of changing the value of single dimension of the representation k En [d] while keeping the values of the other
3

---

Figure 1: Architecture of our model consisting of a VAF, a discriminator and Prototypical network. The representation z from the inference network of the VAF (Sec (.T). is changed a particular di- mension k to get the intervened representation âx. A discriminator (Sec 2.4) is trained to distinguish between z and & and the inference network updated to fool the discriminator. and & are passed to the generator network to map it to the reconstructed data and the intervened data 2g. The orig- inal and the intervened data are concatenated to form the pair (œ, fx), which is then passed to the prototypical network. The prototypical network (Sec S.32.3) maps the pair closer to other pairs with the Under review as a conference paper at ICLR 2023
Diall R/P
2010
x3
same dimension intervened. The prototypical network is updated by it's ability to correctly predict the intervened dimension of the query examples and the magnitude of the change || åx|l-
dimensions the same. We change the value of the intervened dimension to another example's 's value for the same dimension. The result of intervening on representation & in dimension & produces the intervened representation &x- The representation and the intervened representation are then mapped by the generator network to å and fg respectively. This pair of (œ, .t)) forms the input data for the prototypical network with the intervened dimension & being the label and the difference in the representation space |z - Σg as the label for the isometry head. The representation & is intervened upon every dimension to generate d support sets and is also intervened upon one dimension chosen uniformly from [d] to generate the query set. The self-supervised data generation algorithm takes in batch of data æ and outputs d support sets S {S1.... Sa). one query set Q, labels for the query set L and labels for training the isometric head I.
2.3 PROTOTYPICAL NETWORK
Our proposed prototypical network maps a pair of data generated by intervening on single dimen- sion of the representation to a lower dimensional metric space. By mapping pair of data to the metric space, the prototypical network can focus on the factor differing value between the pair while being invariant to the values of the other non non-differing factors. Critically, the factor differing in value remains the same across pairs different examples when the same dimension is intervened upon and hence should be mapped closer in the metric space. Thus, comparing a pair of data al- lows the prototypical network to focus on the change or difference that was brought about by the intervened dimension, and makes the central focus of the losses this change.
The prototypical network first takes in elements of the generated support set (described in Section 2.2 and computes an m-dimennsional representation through the embedding function f+ RR x Rm In this mndimmensional space, the prototypical network computes prototype embedding C₁ for each element in the support set S E using eq. 2
Σ 1V
(5
C₁
SA
While the support set is used to compute the prototypes, the query set is used to compute the loss by calculating the distance of its embeddings in the metric space to the target prototypes. For each dimension of the representation to encode information about a unique factor variation, each di- mension when intervened upon and mapped to the data space must change different factor of
4

---

Under review as a conference paper at ICLR 2023
variation. Thus embeddings of pairs of data generated with the same intervening dimension of the representation must cluster closer in the metric space and away from the clusters of other dimen- sions. To enforee this, we introduce the uniqueness loss which is computed for each query q()) example by calculating the negative log- ligllikeihood of the true class / as in eq.
1
min Lu(r, = 7:0.8
log (1 = 11/()) (L
|Q1
(3)
o EQ
where the probability of each class p, (t = () is calculated as a distribution over the Euclidean distance d the prototypes as in eq.
exp c (-
(-)(f((())) (())
py (1 (1
En exp (-d(f, (g()), )(ck)
The loss for every intervening dimension of the query examples multiplied by the KL--divergence of that dimension, averaged for the batch of examples. This ensures that the loss for the intervening dimensions is scaled by amount of information encoded by that dimension. For the dimensions that do not encode any information, and hence do not change any factor of variation upon intervention, the corresponding loss is scaled by zero. This is important as we do not need any prior assumptions on the dimension of the representation needed encode all the factors and the VAE can find the right number of dimensions needed.
In addition to the uniqueness loss, we want each dimension to consistently encode only a single factor of variation. When the representation = is first intervened on dimension A: and mapped to the data space it makes certain change in factor. When the representation intervened again dimension k by a different amount and mapped to the data space it should produce change in the same factor, irrespective of the amount it was changed. When passed in to the prototypical network, the pair of data generated by the original and intervened & must be embedded in the prototypical metric space closer to the pair generated by the representation intervened in the same dimension by a different amount. To enforce this we introduce the consistency loss in eq. 5 where the prototypes are replaced by the embeddings the same example the support set.
min Lo(r, 0. 7,,0,9
Σ EQ
logr,, (t ()) KI(go
Q
(5)
With ri calculated as follows:
rt = (1)
exp
(6)
Ex exp (-d(f(
The consistency loss and the uniqueness loss are added together get a combined prototypical loss eq. 7
Lp(y, o 0) Lo + Lu
As an additional inductive bias, as proposed in |Horan et al. (2021) }. we constrain the generative mapping between original and intervened representation (z, 2 and the generated pair (#, â&) to be locally isometric Donoho Grimes (2003). Thus the factor changed in åe when compared with must differ in value proportional to the corresponding change in dimension A: of and the intervened 2g. This serves as an imperative inductive bias for unsupervised disentanglement.
The additional head he RD R Rd, when given pair of data, is trained to predict the difference the values for the all the dimensions of and E through the loss function in equation
g
min Cr 0, = ((+ #x)) - E4 $.0.0
The training data for the isometry head generated in a self supervised manner as described in section 2.2 where the support set S consists of data pairs and the set I consists of the corresponding targets. In the final implementation, f,, and he share all hidden convolutional layers and differ only in the
final fully connected layer.
5

---

Under review as a conference paper at ICLR 2023
Figure 2: A comparison of latent traversals in latent space for the 3DShapes and Dsprites dataset. Left: 3DShapes, Right: Dsprites. Proto-VAFE produces smooth, disentangled latent representations. Row and are some sample original images, and their reconstructions generated by our model, respectively. Rows downward are the traversals for each latent element, as detailed below. For 3DShapes, we actually see near-perfect traversal across all of the known factors of variation.
2.4 REPRESENTATION SPACE DISCRIMINATOR
The above described process does not guarantee that the intervened representations will lie in the same distribution, as the original representations which can lead to the intervened representations being mapped in the data space away from the true data distribution. To minimize the distance be- tween the distribution of the original representations and the intervened representations E, we train a separate discriminator network in the representation space to distinguish the intervened represen- tations from the original representations by minimizing the loss in eq.
min (D(w) - - [logl (3)) (lo((1 D (z)))]
Note here that the real samples for the discriminator are the intervened representations and the fake samples are the original representations. Similar to the training of the generator in Generative Adversarial Network (GAN) Goodfellow et al. al. (2014), the inference network, which generates these representations, is trained to fool this discriminator by minimizing the loss in equation [0 Since each dimension in the original representation can independently changed to obtain the intervened representation, the inference network is encouraged encode representations such that the total cor- relation loss (L(((=)]] II: ((2;)) is minimized as proposed in Kim Mnih (2019).
min (x(Ø) B.[log(1 Dm)(=))
(10)
The final objective (equation 11 of our method is weighted sum of the different losses, and is optimized by the network parameters of the VAE and the prototypical network corresponding to each loss.
min L - Lv 0) aLE(d) + XCr(1, 0. (...6660)))) 6.0.1.0
(11)
3 EMPIRICAL EVALUATION
To empirically evaluate our method, we perform both quantitative and qualitative evaluation on two synthetic datasets and one real dataset with known factors of variation and qualitative evaluation on CelebA dataset Liu et al. (2015). The two synthetic and one real datasets are generated from independent ground truth factors of variation; DSprites Matthey al. (017)) binary x 64 images with 5 factors of variation: 3 shapes, 6 scales, 40 orientations. 32 x-positions and 32 y-positions; 3D Shapes Burgess & Kim (2018) 64 x 64 x 3 color images with 6 factors variations: shapes, 8 scales. 15 orientations, 10 floor colors, 10 wall colors and 10 object colors; MPI3D real Gondal ct al. (2019) 64 64 x 3 color images with 7 factors variations: 6 colors, shapes, 2 sizes, cam- heights, background colors, 40 horizontal axis, 40 vertical axis, and real world dataset;

---

Under review as a conference paper at ICLR 2023
1-
object Rue
object
scale
Scale
shape orientation
shape
orientation
-6
Figure 3: Output of the prototypical network with embedding dim m 2 when the input is real pairs of data with single factor with differing values amongst the pair on the 3DShapes dataset. Each color corresponds changes correctly the pairs from the original dataset. This suggests that the prototypical network clustering pairs of images based on the changed factor of variation. Left: = 10. Right: = 5.
different uniquely changed factor of variation. The network clusters the
CelebA. The details of the architecture for the different components and the corresponding hyperpa- rameters are listed in the supplementary material (Appendix B.2) and (Appendix B.3) respectively.
DATASETS
DSPRITES
3DSHAPES
MODEL
FVAE
DCI
MIG
FVAE
DCI
MIG
B- B-VAE
0.51 110
0.23 + .10
0.15
100
0.81
1.10
0.44
.17
0.28 + .18
ANNVAE
0.70
0.28 + .10
0.23 .10
0.84 1 .09
0.46 1 .16
0.31 1 .15
B-TCVAE
110
0.35 .06
0.17 .09
0.88 .07
0.63 1100
0.40
1 .18
EVAE
0.74 .06
0.38 = .10
0.28 .09
0.81 .06
0.47 .12
0.33 14
GR-FVAE
0.75 + .08
0.41 = .07
0.31 + .06
0.79
0.49 = .06
0.43 + .11
PROTOVAE
0.70 = .06
0.51 = .04
0.37 + .09
0.90 + .06
0.84 = .07
0.71 + .11
Table 1: Various disentanglement metrics evaluated across a number of state of the art methods for the DSprites and 3Dshapes dataset. For all metrics, higher is better. The results for the other models are obtained using the hyperparameter settings and experimental conditions as described in Locatello al./ (2019). The scores for all the models were averaged across ten runs with different random seeds, with standard deviation shown as +. Gr-FVAE is the Groupify VAE variant applied to the VAE, as this is the closest to our model and the results for which are taken from Yang al.) (2021). The highest values in a column are written in bold. As we see, the Proto VAE outperforms the state the art on a majority of the metrics. Proto- VAE hyperparameters for DSprites and 3Dshapes results shown are {a- 10, À 10, K 10} and {ar 20, À = 20, 20}
We qualitatively evaluate our model by intervening on the different dimensions of the learned rep- resentations and traversing the range of values of the dimension linearly fixed range [-2 A model is better disentangled if the changes made in the data space while traversing a dimension are similar the changes factor of variation in the data space.
Our model both finds the correct number of factors and encodes them separately without any specific hyper-parameter tuning. From Figure 2 we can see that our method ProtoVAE, produces disentan- gled traversals covering both the number of factors as well as the entirety of the range of the values for the Dsprites and the 3DShapes dataset. The latent traversals on the MP13D dataset can found in figure (6) Our method effectively separates the factors thus disentangling the learned latent rep- resentation without compromising on the reconstruction quality as seen from row 2. Owing to the unsupervised nature our method struggles to exactly disentangle the non-isometric discrete factor of shape in the DSprites dataset. For the 3DShapes dataset, in our traversals in 2 we achieved near perfect disentanglement, completely unsupervised. In the Appendix D. we show the performance of the Proto VAE for a subset the Dsprites dataset with only a few factors. We show traversals from models PactorVAF and 3- VAE, along with our model, with only few isometric factors for com-
7

---

Under review as a conference paper at ICLR 2023
Model
FVAE
DCI
MIG
B-VAE
.41 = .05
.23 .04
.06 = .03
VAn
.29 + .04
=12
.07 = -07
4
3-TCVAE
.45 + .06
.27 =7
.16 = .03
4
4
4
FVAE
.40 .04
.30
.23 = .03
+
4
4
4
4
4
4
4
4
DisCo
.39 +0
.29 2 .02
.07 003
4
4
4
ProtoVAE
.46 404
=.5
.25 = .11
Figure 4: Left: Latent traversals the CelebA dataset PAo successfully captures ground-truth factors of variation on real-world data. From top to bottom: background color, hairstyle, head angle, age, hairstyle, hair color, skin color, face profile. Right: Quantitative comparative metrics the MPI3D dataset. Proto VAE performs comparatively or better consistently across multiple metrics on difficult real disentanglement dataset. See Appendix Fig. 6 for latent traversals on MPI3D. Proto-VAE hyperparameters for results shown are {a = 10, X = 2, K = 2}. The numbers for DisCo have been borrowed from their Ren Ren et al. (2021) for the VAE-based methods.
parison. Our proposed Proto VAE is the only model that does not conflate two factors and encodes them in separate dimensions of the representation.
Furthermore, we quantitatively evaluate the learned representation by calculating state-of-the-art. dis- entanglement metrics. We choose metrics from each of the three kinds of metrics described in /Zadii al. (202T): Intervention- FactorVAE Kim & (2019) Predictor-basee Disentanglement- Compltenes--nfoomaiveeess (DCI) Eastwood & Williams (2018) and Information-bast Information Gap MMG))Ceen al. (2019). The metrics were implemented as proposed in Locatello et al. (2019) with the same hyperparameters. We refer the interested readers to Zaidi et al. (2021) for an intuitive understanding of the metries. We highlight here that our model achieves higher DCI metric and a higher corresponding completeness and informativeness metric, which reflects the mode covering capabilities of the learned representation.
Mutual
From table Te see that on the DSprites dataset, our method outperforms the state of the art models in a majority of the metrics. Similar performance of our model on the 3DShapes dataset as seen in table. Also, the variance in the metrics for the different runs is significantly lower than of the previous methods, thus ensuring a more robust way to disentangle representations. For the real disentanglement dataset of MPI3D Gondal et al. (209)) consisting of camera taking photos of an object attached to a jointed arm, we see that our model consistently either matches or outperforms the state of the art. From the baselines, especially important is the Factor VAE model, which is the base model upon which we add our contributions for the ProtoVAE model and hence use it as comparison demonstrate the effectiveness of our contributions.
On the CelebA dataset, we find that across multiple runs, our model is able to find the same "natural" decompositions that correspond to numan-interpretable factors of variation consistently (Fig. 4). We notice that the model is not constrained to completely encode one factor per latent dimension and the model might encode different ranges of factor in different latents; we see this occur for example when it encodes half of the azimuth in one latent, and half in another. However, as we can see, for the most part, each latent dimension contains information only about one factor of variation and even in the unsupervised regime our model still encodes natural decompositions.
We visualize the embedding space of trained prototypical network using our method in Fig. 3 We see that when input to the prototypical network is pairs of images from the dataset, with one ground truth factor differing in value between the pair, the prototypical network effectively clusters the pairs based the differing factor. This clustering aligns with the labels based the intervened dimen- sions during training and thus points to the effectiveness of the prototypical network for encouraging disentanglement.
8

---

Under review as a conference paper at ICLR 2023
We also performed quantitative and qualitative ablation studies the 3DShapes dataset by changing the values of a, À and & to understand the effectiveness of each of the components and losses we introduce. The results of these ablations can be found in the supplementary material (Appendix C). Furthermore, we also perform ablation studies on the effect of dimension m of the metric space of the prototypical network limitation cases where the representations of the model did not axis align with a few factors but was rotated with respect to those factors. We see that smaller values of the prototypical network metric space m performs better by encoding data in tighter clusters which in turn it imposes stronger constraints on the VAE. The discriminator and the corresponding Lg helps in confining the encoding of the factors into a single dimension whereas Lp alone fails to do this effectively as seen in S
the metric scores. We also show in the Appendix (Section C) some
4 RELATED WORKS
Many state-of ahrt unsupervised disentanglement methods extend the VAE objective function to impose additional constraints on the structure of the latent space to match the assumed independent factor distribution. 8-VAE (Higgins et al. 2017) and AnnealedVAE (Burgess et al. 2018) heavily
penalize the KI. divergence term thus foreing the learned posterior distribution gs ar to be inde-
pendent like the prior. Factor- VAE (Kim & Mnih 2019) and 3-TCVAE (Chen et al.
2019) penalize
((:)) where the aggregated posterior is calculated as (2) = (ote)/((2 mx)] = Y E (o(2 using adversarial and statistical techniques respectively. DIP-VAE (Kumar al.] 2018 forces the covariance matrix of the aggregated posterior q(z) to be close to the identity matrix by method of moment match- ing. The changes that we described in the latent space are defined as intervention by |Suter et al. (2019) to study the robustness of the learned representations under the Independent Mechanisms (IM) (Schoelkopf et al., 2012) assumption. Most closely related to our work are VAE models that learn to disentangle by altering the latent code. In Jha et al. (2018), the authors use VAE or AE with a split double latent code, with a cycle consistent loss, but required that attribute labels be known a priori, which was also a requirement in FFng et al. (2018) which learned by swapping out chunks of the latent code, and [Szabó et al. (2017). which used the labels as a constraint to find unique disentanglement. The authors in Chen et al. (2020) also used cycle-consistent loss, but again required labeling. In Hu et al.] (2018) the authors attempted unsupervised disentanglement with regular (non- variaaiiiii)l Autoencoder network models, by stacking one after another, our model in- stead uses a prototypical neural network. In Lezama (2018) the authors derive a novel Jacobian loss combined with student-teacher iterative training algorithm with an Autoencoder network model. In Park et al. (2020) the authors develop a mheettmmaippllating model aimed at human-interactive image manipulation tasks.
the total correlation of the aggregated posterior (0(2). TC (L(g(>)IIIL=
Yang et al. 2021) use the group based definition by Higgins et al. (2018) and a cycle consistency loss define the elements of a group. Our work differs significantly we do not re-encode the data or the intervention ooraeoobbbbbosss..s. Zhu et al. (202T) encode the the latent space of a VAE using the commutative Lie group and enforce constraints on the latent space. A recent work Ren et al. (2021) propose to learn disentangled representations from pre-trained models using contrastive methods.
5 CONCLUSION AND FUTURE WORK
In this work, we proposed a novel generative model consisting of a VAE and a Prototypical Network for learning disentangled representations in a completely unsupervised way, inspired by recent dis- covery of sufficient inductive biases. We impose constraints on the structure of the representations learned by training the model in an self-supervised manner to encode information about the differ- ent factors in separate dimensions of the representation. Our proposed method is able to outperform other state of the art networks on a number of metrics on three prominent disentanglement datasets. For future work, our method can be easily adapted to be trained in a weakly supervised regime with pairs of data differing in known number of factors being the prototypes for the prototypical network. The results can be possibly improved by intervening on multiple dimensions of the representations simultaneously. The importance of methods that can disentangle data without labels is critical as data is plentiful and the resulting representations give interpretable insights into the variations in the

---

Under review as a conference paper at ICLR 2023
data distribution, and can be used for downstream tasks. Our hope is this work adds evidence that self-supervised generative methods are important in this endeavor.
REFERENCES
Yoshua Bengio. Deep learning of representations: Looking forward, 2013.
Chris Burgess and Hyunjik Kim. 3d shapes dataset. https///sshhllooomeeemmmm/////es 2018.
Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in B-vae, 2018.
Jiawei Chen, Janusz Konrad, and Prakash Ishwar. A cyclically-trained adversarial network for in- variant representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition Workshops, pp. 782-783, 2020.
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources disentan- glement in variational autoencoders, 2019.
David L. Donoho and Carrie Grimes. Hessian eigenmaps: Locally linear embedding techniques for high-dimensional data. Proceedings of the National Academy of Sciences, 100(10):5591-559 2003. ISSN 0027-8424,
doi: 10.10733pnas.103111 1031596100.
URL. https: httt:p:/wwww.pna.t org/
[content/10010/5555
Cian Eastwood and Christopher K. I. Williams. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
URL httpthtpss///ppppeeeeeeeeeeeeeee otw.://eeeeeeeeee.eeeoooooooooaa
n Forum?id=By-7de 7dz-22
Zunlei Feng, Xinchao Wang, Chenglong Ke, An-Xiang Zeng, Dacheng Tao, and Mingli Song. Dual swap disentangling. Advances in neural information processing systems, 31, 2018.
Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Mar- tin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation to the real world: new disentan- glement dataset. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buce Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019. URL wtt.tttsse./pooceecs.........aaa d97c404b6 011.:444666111211211011111111119
pdf
lan J. Goodfellow, Jean Pouget-Abaal Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks, 2014.
I. Higgins, Loic Matthey, Pal, C. Burgess, Xavier Glorot, M. Botvinick, S. Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In ICLR, 2017.
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. Towards definition of disentangled representations, 2018.
Daniella Horan, Eitan Richardson, and Yair Weiss. When is unsupervised disentanglement possible? Advances in Neural Information Processing Systems, 34, 2021.
Qiyang Hu, Attila Szabó, Tiziano Portenier, Matthias Zwicker, and Paolo Favaro. Disentangling factors variation by mixing them, 2018
Ananya Harsh Jha, Saket Anand, Maneesh Singh, and VSR Veeravasarapu. Disentangling factors of variation with cycle-consistent variational auto-encoders. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 805- 820, 2018.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, 2019.
Diederik Kingma and Max Welling, Auto-encoding variational bayes, 2014.
10

---

Under review as a conference paper at ICLR 2023
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disen- tangled latent concepts from unlabeled observations. In International Conference on Learning Representations, 2018. URL https: / http://:////eeeeeee.eeeo net/ wtt../.teteeeeteee-/llc11
José Lezama. Overcoming the disentanglemenn reconstruction trade-off via jacobian supervision. In International Conference on Learning Representations, 2018.
Ziwei Liu, Ping Luo. Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Francesco Locatello. Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. International Conference on Machine Learning, pp. 4114- 4124, 2019.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset. http://ggtiii.cooomeceppppaiipps 2017.
Taesung Park, Jun-Yan Zhu, Oliver Wang, Jingwan Lu, Shechtman, Alexei Efros, and Richard Zhang. Swapping autoencoder for deep image manipulation. Advances in Neural Information Processing Systems, 33:7198-722 12111, 2020.
Xuanchi Ren, Tao Yang, Yuwang Wang, and Wenjun Zeng. Do generative models know disentan- glement? contrastive learning all you need. arXiv preprint arX1v:2102.10543, 2021.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models, 2014.
Bernhard Schoelkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning, 2012.
Jake Snell, Kevin Swersky, and Richard S. Zemel. Prototypical networks for few-shot learning, 2017.
Raphael Suter, Dorde Miladinović, Bernhard Schölkopf, and Stefan Bauer. Robustly disentangled causal mechanisms: Validating deep representations for interventional robustness, 2019.
Attila Szabó, Qiyang Hu, Tiziano Portenier, Matthias Zwicker, and Paolo Favaro. Challenges in disentangling independent factors of variation. arXiv preprint arX1v:1711.02245, 2017.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle, 2015.
Tao Yang, Xuanchi Ren, Yuwang Wang, Wenjun Zeng, and Nanning Zheng. Towards building group-based unsupervised representation disentanglement framework. In International Confer- ence on Learning Representattions 2021.
Julian Zaidi, Jonathan Boilard, Ghyslain Gagnon, and Marc-André Carbonneau. Measuring disen- tanglement: A review of metrics, 2021.
Xinqi Zhu, Chang Xu, and Dacheng Tao. Commutative lie group vae for disentanglement learning. In International Conference on Machine Learning, pp. 12924-12934, PMLR, 2021.
A APPENDIX
A.1 VARIATIONAL. AUTOENCODER NETWORKS
Variational Autoencoder networks (VAEs) Kingma & Welling (2014); Rezende et al. (2014) are generative models, which consist of an Inference network and Generator network, instantiated as deep neural networks, which aim to maximize the data likeliioooo calculated according to the generative model po (x/22 with ~ p(z):
log logp(z) = log (
(12)

---

Under review as a conference paper at ICLR 2023
where 2 is lower dimensional latent variable than and is sampled from known prior. However, calculating this exact log-likelihood is intractable due to higher dimensional latent space, hence, VAEs instead maximize bower-bound by approximating the true posterior distribution po(2 with tractable distribution qo (2|2r learned by the Inference network:
max Li (0, 6) = Ba(a)) 0
- -#Dxa(go(el))))lttet
(13)
Where q(r) is the empirical distribution of the dataset given by q(z) = * D 8(x(0) and
Dx (xl(a(2|)))ll(e) (x log 2a(z)r)
p(s)
and @ parameterize the distributions learned by the Generator network and Inference network of VAE respectively. The top-down generator p=(2))) and bottom up Inference network network (=|2)) and trained jointly to maximize lower bound of the empirical distribution of the training data, when passed through the information bottleneck Burgess et al. (2018).
where N is the total number of training points in the dataset, the Inference network network maps (*)) to lower dimensional representation 2(9), which sampled from the Inference network distribution qa (2(2))) ,and the Generator network network maps the representation back to the higher dimensional data space و(((), In the objective function in first term measures the difference between the example x(() and the reconstructed example ((0) and the second term measures the distance between the approximate posterior distribution (0(=)z:) of the learned representation and the assumed prior distribution p(z) ~ (N,, 7). VAEs, and variants thereof, have shown to disentangle data by imposing structural constraints in the latent space and owing to the information bottleneck principle Buurgess et al. (2018). Our network builds on this base capability of the VAE.
Given training example r(() : E
A.2 PROTOTYPICAL NETWORKS
Prototypical Networks Snell et al. (017), originally used for the few-shot supervised classification regime, are used to learn a metric embedding space where examples belonging to the same class cluster around a prototype representation of the class. The base prototypical network is based around supervised loss, where the labels are known priori and the network is used to project the data into a new clustered metric space. The network trained by training episodes consisting of labeled support and query sets, where the support sets are used to find the class prototypes and the query sets are used to evaluate the loss. Mathematically, they are given a labeled support set {SS, ... ..5k), (N)
where K' is the number of classes, and where each set S ((=))) ..")). ( 1/ (()))) contains
k E IR such that {()) = k. Prototypical networks compute an M-dimensional representation Cj E Rm or prototype, of each class, E (1..... K}. through an embedding function fn R.D
Rm usually deep neural network with parameters 7-
Each prototype Cę is the mean vector of the embedded support points belonging to its class
C
Σ (xL 1E8x
fx(2(²)
(14)
Sp
For given query point Eg. classification is performed by calculating the distance of the query point's embedding f- (22,) with the prototypes of the different classes and assigning the label of the closest prototype. While any distance metric AM x RM + [0, +oo) can be used, the Euclidean distance metrics is used in practice which corresponds to the spherical Gaussian density. A distribution over the classes is calculated by applying a softmax function on these distances:
exp (=d)((((()))
py = klm,)) = Σn cxp (2)) cx)))
(15)
The weights of the embedding function are updated by minimizing the negative log- log-probability of the true class:
min L2(7) log p- k: (..g)
(16)
12

---

Under review as a conference paper at ICLR 2023
The above loss is minimized when examples belonging the same class cluster around the class prototype sufficiently far from the prototypes of the other clusters. For our work, we convert this into an unsupervised training regime by creating labels the fly in a self-supervised manner, as described in the text.
A.3 DISENTANGLED REPRESENTATIONS VIA GROUP THEORY
We base our work on the proposed formalized definition proposed by Higgins et al. (2018). which connect the disentangled representations to symmetry transformations using group theory. They define the notion symmetry transformations as transformations which affect only certain aspects properties of the data while leaving the rest unchanged. The set of these symmetry transformations constitute group G and their corresponding effect on some abstract space are defined as the actions of the symmetry group G x X X. Furthermore, they define actions disentangled with respect to some group decomposition G = G1 x x Gn if there exist a decomposition of the abstract space X1 .x,,,, and actions & G. X - X₁ such that
(. gg)) ( (01,,..... (01 U1,--'9h In Un)
(17)
Using these definitions Higgins al. (2018) characterize representations disentangled with respect to a particular decomposition of a symmetry group into subgroups if decomposes into independent subspaces, where each subspace is affected by only the action of single subgroup.
Importantly, this definition also brings in an agent as critical aspect of the definition. Disentangle- ment is longer disembodied here: an agent takes a specific action symmetry transformation), and observes the result of the transformation. Formally, we can say that agent's s representation of the Z
data is disentangled with respect a particular group decomposition if: (1) An action G exists, (2) there exists a map f W + Z which is equivariant between actions on W and Z, and (3) each element the decomposition the representation = Zi X Zn is affected by only one transformation G:
B IMPLEMENTATION DETAILS
In this section we give more details and descriptions of our proposed ProtoVAE architecture which consists of. The visual depiction of the model shown in Fig. T The layers and layer sizes are given in Table 2] The hyperparameters used are given in Table 4 The architecture of the latent discriminator a Multi-Layer Perceptron (MLP) with 6 layers with 1000 hidden dimension per layer and 2 output logits as in Kim & Mnih (201).. We normalize the data between [0.1] as input to the Inference network and the output of the Generator network is the mean of a Bernoulli distribution po(x12). The reconstruction cost in the first term of [ is calculated the negative cross- cross-entropy between the output of the Generator network and the input to the Inference network. Code available at: https//ggthuh.cmmmooaaaaa//aoott
13

---

Under review as a conference paper at ICLR 2023
B.1 SELP-SUPERVISED DATA GENERATION
Algorithm 1: Self-Supervised Dataset Creation for Prototypical Network Learning
input Data examples {z
} {,, Batch size B, latent dimension d
given Inference network: ())2 Generator network: p=(x]=)
output: Support Set: (S1, S2, ...: Sa} s.t. = {st,-.. ... o J. Query Set:
Q = - .gl.... q J. Query set labels L, Isometry Loss Labels: {1iVj [d]
Encode data examples into representation space {2(0) (0(2 with E Rd Generate each Support Set Sj Vj ed:
for j to d do
f random permutation of [B]
(8) - 3j
(c())
E
[B]: swap the j dimension of example with that of ((()))))
J - (zf
20))
2 (0)
Vi e E [B]
(1) p
Vi E [B] generate data from intervened latent
(+(1, ,4)
Vi E [B] assign examples to an element of a support set with label j
Sj (8).
1 (1x(0)
20
E [B]]
end
+ [S ... Sa 1 (4, ... Ia}
Generate Query Set Q: for = to B do
X random permutation of [B]
n [d]
(x(i))
4 2/
2 + (zl 1..-%
(i)
2k
Zd
(0) ( (1))
fg
(1)) (#6) 200
100
end
Q (2) (B)}
(10), 11813
B.2 ARCHITECTURE DETAILS
INFERENCE NETWORK
GENERATOR NETWORK
PROTONET
INPUT: CHANNELS 64 CONV. LRELU. STRIDE
INPUT: E TO FC. 128 LRELU.
INPUT: X
64 64
CONY. LRELU, STRIDE
4 4 CONY. 32 LRELU, STRIDE
4 x 64 LRELU.
x 4 CONV, 32 LRELU, STRIDE 2
4 4 CONV. 64 LRELU. STRIDE
4 4 UPCONV. 64 LRELU. STRIDE 2
4 x CONV. 64 LRELU, STRIDE 2
4 4 CONV. 64 LRELU. STRIDE FC. 128 LRELU. FC. x 10.
4 4 UPCONV. 64 LRELU. STRIDE 2
4 CONV. 64 LRELU, STRIDE 2
UPCONY. 64 LRELU, STRIDE
FC. 128 LRELU.
4 4 UPCONV, SIGMOID, STRIDE 2
FC. 2, FC, 10
Table 2: Network Architecture of the VAE Inference network, VAE Generator network, and the Prototypical Network.
B.3 NETWORK HYPERPARAMETERS
We detail the hyperparameters used for each dataset. Table 1 shows the full list values for each component.
14

---

Under review as a conference paper at ICLR 2023
PARAMETER (DSPRITES)
VALUE
PARAMETER (3DSHAPES, CELEBA)
VALUE
LATENT SPACE DIMENSION
10
LATENT SPACE DIMENSION
10
METRIC SPACE DIMENSION
3
METRIC SPACE DIMENSION
2
OPTIMIZER
ADAM
OPTIMIZER
ADAM
OPTIMIZER (BETAl)
0.9
OPTIMIZER (BETA1)
0.9
OPTIMIZER (BETA2)
0.999
OPTIMIZER (BETA2)
0.999
LEARNING RATE (VAE)
le-1
LEARNING RATE (VAE)
1c-04
LEARNING RATE (GAN)
le-1
LEARNING RATE (GAN)
5c-05
LEARNING RATE (PROTO)
le-1
LEARNING RATE (PROTO)
1c-04
BATCH SIZE
128
BATCH SIZE
EPOCHS (TRAIN)
50
EPOCHS (TRAIN)
30
(TOTAL CORRELATION (GAN)
10
(TOTAL CORRELATION (GAN)
20
À (PROTOTYPICAL Loss)
10
À (PROTOTYPICAL Loss)
20
(ISOMETRY Loss)
10
(ISOMETRY Loss)
20
Table 3: Hyperparameters values for the different components of the network for the different datasets.
PARAMETER (MP13D)
VALUE
LATENT SPACE DIMENSION
10
METRIC SPACE DIMENSION
3
OPTIMIZER
ADAM
OPTIMIZER (BETAI)
0.9
OPTIMIZER (BETA2)
0.999
LEARNING RATE (VAE)
1e-4
LEARNING RATE (GAN)
5e-5
LEARNING RATE (PROTO)
le-4
BATCH SIZE
128
EPOCHS (TRAIN)
30
(TOTAL CORRELATION (GAN)
10
(PROTOTYPICAL LOsS)
2
R (ISOMETRY LOSS)
2
Table 4: Hyperparameters values for the different components of the network for the MPI3D dataset
C ABLATION STUDIES (METRICS AND IMAGES)
We demonstrate qualitative and quantitative results showing the effectiveness of the different hyper- parameters for disentanglement by varying them across a range while keeping the others the same. We perform these ablations studies the 3D shapes dataset with hyperparameters given in Table except the one hyperparameter which is varied. in 5 we vary the dimension of the metric space of the prototypical network and it's effect the various metrics. In Table 5 we vary the strength of the total correlation loss as measured by the latent space GAN to train the VAE by varying the hy- perparameter a. In Table O we perform similar experiments by varying the hyperparameter À which scales the combined loss of the consistency and uniqueness loss and the regularization strength of the Isometry loss which is governed by the hyperparameter k.
Qualitatively we show the latent traversals of results where learned latent representations are not axis aligned with the ground-truth factors of variations. In Figure 7 we show the effects of having lower value of and see that even though each dimension encodes information related to one factor, the complete range of the factor encoded by multiple dimensions. of having lower regularization for the Isometry loss (x) and see that learn subset of factors which are rotated in space and are not axis aligned with the ground truth factors.
Figure we show the effect
In addition, we perform ablation study by removing the GAN discriminator and seeing the result of just the VAE and prototypical network Fig. 9 We found that without the GAN, a lower dimensional prototype metric space is needed (we dimension d 1), but that the network is
15

---

Under review as a conference paper at ICLR 2023
METRIC
= 2
m 5
m 8
METRIC
=
10
20
FVAE
.90 1 .06
.88 1 .04
.87 1 .04
FVAE
.93 1 .04
.85 1 .03
.88 .05
DCI
884 + 007
.77 + 05
.74 05
DCI
.78 - 05
.81 = .06
.81 .07
MIG
771 11
.62
- 08
57 06
MIG
.59 00
.63 = 00
.65 .08
8-VA
.92
- 05
91 .04
B-VAE
.92 .06
= 04
.90 .05
Table 5: Ablation study showing the performance different metrics for our proposed PAO,oVAE, with different values of the prototype dimension space m (left) a (right) on 3D shapes dataset. In some instances, these versions outperform the results posted in the main paper. We see that smaller values the prototypical network dimension actually perform better as it imposes stronger constraints Factor VAE metrics but does not perform as well on the MIG and the DCI scores as the Latent Space GAN helps in confining the encoding of the factors into single dimension whereas losses do not do this as effectively.
the VAE. We also that a
performs well for the B VAE and the
METRIC
0
10
K = 40
METRIC
==
= 10
À = 20
FVAE
4: .03
.72 .04
FVAE
+ .04
.88 .05
DCI
.57 + .66
.72 + 00
.04
.69 .05
DCI
.64 + 005
.71 .02
.81 = .07
MIG
.43
066
.56
.88
.58 .08
MIG
.58
+ .66
.57
+ .06
.65 .08
3-VAE
.80
+ .05
884
+ .06
.88
.06
B- B-VAE
.88 - 004
.88
.05
.90 .05
Table 6: Ablation study showing the performance on different metrics for our proposed Proto with different values of the hyperparameter λ. (left) and X (right) on the 3D shapes dataset.
still able disentangle well. We leave for future work more investigations into removing the GAN component entirely.
Another promising future work would be to impose constraints the structure of the metric space to obtain more separable clusters. Moreover, curriculum learning could be used while performing interventions ensure that the intervened dimension is changed by significant value and then to slowly decrease this value as training proceeds.
N2.2 7852
Figure 5:
Latent traversals on the MPI3D real world disentanglement dataset. The data is collected via camera that observes jointed arm with known changed ground truth factors of variation. From top to bottom: original data, reconstruction, arm angle left/right, arm angle top/bottom, background height, end color, size.
16

---

Under review as a conference paper at ICLR 2023
Figure 6: More latent traversals the CelebA dataset. We that the traversals similar the ones discovered in Fig. 4 of the main paper.
Figure 7: Results on 3dShapes with a lower value of X, showing theoretically correct disentangle- ment that disagrees with a "natural" decomposition namely, that the factors need not be completely capture by one latent, and may be split into several; for example, here, color change need not cover the entire spectrum in one traversal. In row 7, 8, left and row 6,8 right, we can see that color traversal was split into two latents. This still follows the group theoretical definition of disentangled repre- sentations. And given a high value k, the mapping between the dimensions and the corresponding traversals is isometric.
C.I PROTOTYPICAL. NETWORK EMBEDDINGS
We take a closer look at the workings of the prototypical network component, and analyze its output. To do so, trained network using prototypical metric dimension d 2, that plot the
17

---

Under review as a conference paper at ICLR 2023
Figure 8: Similar to Fig. 7 non-unique decomposition can be found by 'rotation' of that factors such that facets of each factor are mixed together. We can see this here with the wall and shape color being intermingled. For this run, had lower value of the Isometry loss scaling (x). This shows that this loss was important in ensuring a more natural decomposition.
Figure 9: Results of the ablation study removing the GAN discriminator and just training a VAE with the prototypical network, find that removing the GAN requires lower dimensional metric space. However, even without the GAN discriminator we find that the network performs remarkably well disentangling the factors of variation into natural decompositions. We also found that we needed higher values of À and x for better disentanglement. Due this however, with the GAN removed not all factors are captured and the reconstructions are slightly blurrier. We leave to future work further development of this variant without the GAN discriminator.
low dimensional prototypical network metric space d
resulting points in 2D space. We pass into the prototypical network network two data samples from the 3DShapes dataset that differ in one ground-truth factor of variation (mimicking the intervention process, but with ground truth examples), and plotted the corresponding projection into the R² space. This tests to see the trained prototypical network will generalize the clustering to ground truth pairs of data which differ in only factor shown in Fig. 12 If the prototypical network
18

---

Under review as a conference paper at ICLR 2023
Figure 10: Results of latent traversals for the DSprites dataset for two baseline methods of 3- and FactorVAE. left The latent traversals of representations learned by a 3-VAE with same hyperparameters as listed in and 3 = 4. right The latent traversals for the FactorVAE model with y 40. We see that our method qualitatively outperforms both these methods on the DSprites dataset with better reconstruction quality.
Figure 11: Results of latent traversals for the 3Dshapes dataset for the baseline method of Factor- VAE. left The latent traversals of representations learned by the FactorVAE for 20 with the same hyperparameters as listed in right The latent traversals for the FactorVAB model with 7 = We see that our method qualitatively outperforms both these methods the 3Dshapes dataset with better reconstruction quality and more disentangled factors encoded in the different latents.
has been trained correctly and aids in disentanglement of the factors, it would cluster changes to the same latent in similar clusters in IR, Notably, here we supply the prototypical network directly with two real dataset images, which is not the same distribution that the prototypical network is trained on. During training the prototypical network trained on a pair of one reconstruction sample and one intervention-observatooo Furthermore, we also see the effect of the regularization effect of the À parameter on the clustering for model with the other parameters the same. In Fig. 3 the clusters formed by trained model on the 3D shapes dataset.
real
This suggests that during training the prototypical network clusters the pair of reconstructed and intervention-observation according to the index of the dimension of the representation that was in- tervened on to create that pair. This in turn forces the Generator network the VAE to make consistent and unique changes to the data based the intervened dimension. Furthermore, it is apparent that clusters of some factors are well separated whereas clusters are closer in the metric space based on the entity that they affect.
19

---

Under review as a conference paper at ICLR 2023
Inputs to the prototypical network. The first two represent pairs (each pair is column) where the floor hue the only factor that is differing amongst the pair of images. Similarly the bottom rows represent pairs where the object hue is the differing factor.
Figure 12:
D COMPARISON OF MODELS WITH FEW KNOWN FACTORS
To show the effectiveness the Prototypical network in disentangling the factors of variation we perform ablation studies with smaller number of factors of variation, with the latent space being equal to the number of factors. We show the comparison of our method with the 3-VAE and the Factor-VAL that we trained ourselves for fairness of comparison. As an ablation, we created subset of synthetic datasets from DSprites containing only two factors of variation: one with only position x and position y, and one with the orientation (o) and position y factors of variations. We show the reconstructions and the latent traversals of our method in the in Figure ? and Figure ?2 for the position and position y factors and the orientation and position y factors respectively. From the figure we can see that the Pro effectively disentangles the different factors into separate dimensions of the latent representation. We can see that for subdataset consisting
and y, the FactorVAB and 3-VAE correctly disentangle the two factors (Fig. ?? and Fig. ?? as does proposed ProtoVAE However, with y and orientation, we see that the Factor (Fig. ??) and the 3-VAE (Fig. ?? do not fully "naturally" disentangle the factors, leading mixed representations in the latents.
For this, used prototypical metric dimension of d 1. We found that the GAN is only necessary for larger metric dimension space, such as the ones we use for the main paper results. In Figure ??? and Figure ?? we show the Factor-VATs performance and in Figure ?? and Figure ?? we show the 3-VAEs performance on the same datasets. We see that the Proto outperforms these methods.
10, 10)without the latent discriminator, on the sub-dataset from DSprites created by only varying x and y (same shape, orientation, and size.) The first row consists examples from the dataset, the second row consists of the corresponding reconstructions. The third row shows the latent traversals for the 1 dimensions which encodes the position x factor and the 2nd row which encodes the position y factor.
Figure 13: ProtoVAE & y)Trraerrsal results for (a 0.
20

---

Under review as a conference paper at ICLR 2023
Figure 14: ProtoVAL VAE (orientation & y-position) Traverssl results for d 1 without the latent dis- criminator, on the sub-dataset from DSprites created by only varying y and orientation (same shape, position x, and size.) The third shows that the dimension 1 of the latent representation encodes the position y factor and dimension encodes the orientation factor. Unlike FactorVAE and 3-VAE (Fig. ?? Fig. ??). our proposed model fully disentangles the factors into separate latents without rotation.
Figure 15: FactorVAE (x position & position): Latent traversal results for FactorVAE with hyper- parameter value 10. The FactorVAB correctly disentangles these factors without any mixing.
Figure 16: FactorVAE (orietnation y y-position): Traversal results for FactorVAE on the orientation and y-position dataset, showing only partial disentanglement, rotation of the factors. Top four rows: Hyperparameter value for the total correlation loss Bottom four rows: The value of hyperparameter is 10. We see that rotational subspace of the factors is learned. In particular, we can see in row top, the y position and orientation is mixed.
Figure 17: 3-VAE on the subdataset of (x- (x-position & y-position): Traversal results with 9 = 10 We that the 8-VAE did train successfully to disentangle and
21

---

Under review as a conference paper at ICLR 2023
Figure 18: 3-VAE on orientation and y-position: Traversal results for 3- 3-VAE on the orientation and
position dataset. Top four rows: 8 5. Bottom four rows: 3 = 10. We can see the network
performs more poorly on these factors than our proposed model (row 8) in generating smooth traversals.
22


Reviews:
[{"review_id": "fQu9PdE3B_I", "reviewer": "Reviewer_jevi", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an approach, called ProtoVAE, that learns disentangled representations from unsupervised data. The overarching goal is to incorporate two inductive biases into the learned generative model: unique and consistent changes to the latent representations, and local isometry. ProtoVAE is based on a variational autoencoder, with several modifications. First, interventions are performed in the latent space such that both a reconstruction of the original image and a generated image corresponding to the intervened latent representation are produced. Second, a prototypical network is trained with pairs of images as input to predict both the intervention dimension and the magnitude of the change. Third, a GAN discriminator attempts to predict whether a latent representation was produced by the original inference network or if it was the result of an intervention. Experiments on dSprites, 3DShapes, MPI3D, and CelebA datasets show that the method is successfully able to learn disentangled representations.", "summary_of_the_review": "Overall, the paper investigates an important goal, and there are potentially novel aspects contained in the proposed ProtoVAE model. However, the precise nature of the relationship between the ProtoVAE and previous baselines is still unclear in my mind. The disentanglement results are strong but the experimental justification for each of the components is also underdeveloped in the current version of the paper. I am hoping that these issues can be resolved during the discussion period.\n", "strength_and_weaknesses": "Strengths\n- The overall goal of learning disentangled representations from unsupervised data is important for interpretability and controllable generation.\n- Strong experimental results on disentanglement metrics.\n\nWeaknesses\n- The proposed model is complex: it combines a VAE, a GAN, and a prototypical network. In order to justify each of these components, the experiments should isolate their relative contributions but I did not see such results in the paper.\n- The paper is difficult to read. The description of the model is not sufficient to understand how all the components fit together. Specific  matters that remain unclear include: how the interventions were performed, how the support and query sets are generated, and how the traversal dimensions for the visualizations are selected.\n- One of the stated key inductive biases is local isometry, but what exactly this means, why it is valuable, and its corresponding mathematical expressions require further elaboration.\n- The relationship between the proposed ProtoVAE and previous VAE-style approaches to disentangled representation learning is not clearly explained. What aspects of the ProtoVAE are novel and what aspects are included in previous works?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "aYKsTmp18Sg", "reviewer": "Reviewer_69qY", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new approach for training disentangled variational autoencoder (VAE). The approach builds on the top of FactorVAE, and adds a prototypical network that clusters pairs of generated samples that differ in one latent dimension. The encoder and decoder are trained to be easy to be classified by the prototypical network, thus encouraging a disentangled latent space. Experiments are conducted on dSprites, 3DShapes, MPI3D, and CelebA datasets.", "summary_of_the_review": "\nOverall, the idea is interesting, but the current writing quality and experimental results are below the threshold. I would recommend the authors polish the writing and add more clarifications on the experimental results.\n", "strength_and_weaknesses": "\nStrength:\n\n* The idea is technically sound.\n\nWeaknesses:\n\n* The writing needs a lot of improvement. The main text is not self-contained: some important technical details and definitions of notations are deferred to the appendix. In addition, even after reading both the main text and the appendix, some important details are still unclear to me. Also, there are many typos. Detailed questions are below.\n\n    - Page 2: Given that locally isometric is the motivation of one key component of the proposed algorithm, it is important to explain what it means instead of only mentioning the term. \n    - Page 2: what does \"the interpretability of GANs\" mean here?\n    - Page 3: equation 13 should be equation 2\n    - Page 3: The core algorithm of generating samples for prototypical networks is deferred to the appendix. The important notations that the following paper depends on (e.g., support set, query set) are also in the appendix and are not defined anywhere in the main text. \n    - Page 3: \"The the\" should be \"The\"\n    - The notations are not consistent throughout the paper: z and \\hat{z}_k are sometimes in bold and sometimes not. The proposed algorithm is named ProtoVAE or Proto-VAE in different places.\n    - Page 4: Again, the important details of the prototypical network (needed for understanding the paper) are deferred to Appendix A.2.\n    - Equation 3: I understand the KL coefficient is used for scaling down the loss when the dimension is uninformative. However, why do you choose KL instead of other distance metrics? This is not explained in the paper.\n    - Equation 6 does not match what is described in the text \"the prototypes are actions of the different dimensions on a particular example and the query example is the action of a randomly chosen dimension on that example\". Because of that, I am not able to check the correctness/soundness of this part of the algorithm.\n    - Page 6: \"To ensure that interventions in these dimensions do not contribute to the losses, we scale equations 5 and 3 with the KL divergence values of the intervened latent dimension.\" is already explained before. Maybe you can remove it.\n    - Appendix B.1: Is \"k_{(k\\not= i)}\" a typo? k is not used in this line.\n\n* The scores of the proposed approach are much worse than state-of-the-art, and the baseline scores reported in the paper are worse than what was reported in prior work. Details are below.\n\n    - Most of the reported scores of baselines are worse than what was reported in prior work (e.g., https://arxiv.org/pdf/1906.06034.pdf). More importantly, the proposed ProtoVAE is built upon FactorVAE, but the ProtoVAE's scores are worse than FactorVAE's scores reported in prior work (e.g., https://arxiv.org/pdf/1802.05983.pdf).\n    - I understand that different implementations may lead to different scores. As the paper states, the hyperparameter settings and experimental conditions are taken from Locatello et al. But the reported scores are still not consistent with what was reported in Locatello et al.\n\n    Therefore, I am not convinced by the experimental results and have concerns about how much ProtoVAE improves upon prior work.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "XV9Luocrao1", "reviewer": "Reviewer_L6Gs", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed to combine variational autoencoder, prototypical network, and generative adversarial network in a way to solve unsupervised disentanglement. The proposed method was evaluated only empirically.", "summary_of_the_review": "This work tried to solve a challenging problem but did not properly use existing methods and theories. The author misinterpreted and miscited some existing works. The mechanism of the proposed method is unclear due to the confusing writing style of this paper. Therefore I'm afraid I  cannot recommend acceptance of this paper to ICLR.\n", "strength_and_weaknesses": "The author attempted to solve an important and challenging problem -- unsupervised disentanglement -- but seems to have several fundamental misunderstandings of existing theories and tools:\n\n- The author misinterpreted the result of [Locatello et al., 2019] and claimed that unsupervised disentanglement is possible *as long as there is a proper inductive bias*. However, the author did not mathematically prove the proposed method can actually lead to disentanglement.\n- It seems the author does not know the difference between VAE and beta-VAE. They cited the original VAE paper (well there is even no citation in Section 2.2) but used beta-VAE instead, showing that they are unaware of the basic derivation of the learning objective (log-likelihood vs. information bottleneck).\n- The method is overcomplicated. It is unclear which part is essential and which part is optional. For example, the effect of \"latent space GAN\" is confusing.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JKKZPpf0lE", "reviewer": "Reviewer_Fvdi", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an unsupervised approach to learning disentangled latent factors within a VAE-like model. The core idea of the approach appears to be the swapping of a single value within the latent code between batch items during training; the decoded original latent and latent with the swapped value can then be compared allowing the network to isolate specific changes in the data space that the swapping induces. The training proceedusre forces the model to learn to ensure that each single latent dimension only changes one factor in the data space. Experimental results seem to indicate that the proposed approach does a significantly better job at disentanglement than existing approaches when compared using a range of measures. \n", "summary_of_the_review": "As I said in the strengths section, I do think that there might be a good, novel, idea buried within this paper. However as currently written that idea is not clear enough to justify acceptance. Hopefully the authors will work on this aspect. In terms of scoring for correctness, novely and significance I've tried to be positive, scoring on what I think the paper is trying to say rather than necessarily what it does say.\n", "strength_and_weaknesses": "### Strengths\n\n- Experimental results appear to show good results. An appropriate range of measures is used, and the method is compared against a range of other approaches.\n- I _think_ the overall idea is sound, and is novel\n\n### Weaknesses\n\n- Clarity and presentation leaves a lot to be desired (see below)\n- Many of the experimental aspects are rather unclear\n- It's not clear to me that the \"latent space GAN\" is actually a GAN; is it not just a discriminator network that's used in the loss? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject