DualContrast: Unsupervised Disentangling of Content and Transformations with Implicit Parameterization

## Abstract

Abstract Unsupervised disentanglement of content and transformation has recently drawn much research, given their efficacy in solving downstream unsupervised tasks like clustering, alignment, and shape analysis. This problem is particularly important for analyzing shape-focused real-world scientific image datasets, given their significant relevance to downstream tasks. The existing works address the problem by explicitly parameterizing the transformation factors, significantly reducing their expressiveness. Moreover, they are not applicable in cases where transformations can not be readily parametrized. An alternative to such explicit approaches is self-supervised methods with data augmentation, which implicitly disentangles transformations and content. We demonstrate that the existing self-supervised methods with data augmentation result in the poor disentanglement of content and transformations in real-world scenarios. Therefore, we developed a novel self-supervised method, DualContrast, specifically for unsupervised disentanglement of content and transformations in shape-focused image datasets. Our extensive experiments showcase the superiority of DualContrast over existing self-supervised and explicit parameterization approaches. We leveraged DualContrast to disentangle protein identities and protein conformations in cellular 3D protein images. Moreover, we also disentangled transformations in MNIST, viewpoint in the Linemod Object dataset, and human movement deformation in the Starmen dataset as transformations using DualContrast.

## 1 Introduction

Figure: Figure 1: Our proposed contrastive learning-based unsupervised content-transformation disentanglement pipeline. (a) The variational inference of content and transformation factors with $L_{\text{VAE}}$. (b,c) The contrastive strategy to create positive and negative pairs w.r.t. content (b) and transformation factors (c), respectively. DualContrast combines the three parts and is trained end-to-end.
Refer to caption: x1.png

In many real-world image datasets, particularly in scientific imaging domains, the object shapes being visualized may undergo multiple transformations [, ]. These shape-focused image datasets can thus be regarded as samples generated from two independent factors, one representing the semantic attribute termed content and the other representing the transformations. Content is regarded as the form of an object being visualized in the image [] that is unaltered after applying any nuisance transformation. On the other hand, transformation dictates the specific state or realization of that form of the object. Taking proteins inside the cells as an example, changing the transformation factor changes the conformation of the protein where the protein identity (*i.e*., protein composition) does not change (Figure DualContrast: Unsupervised Disentangling of Content and Transformations with Implicit Parameterization). On the other hand, changing the content factor is analogous to changing the proteins from one identity to another through compositional change. Such a phenomenon holds for many real-world shape-focused datasets. Disentangling content and transformation factors of images can significantly facilitate several downstream tasks, including clustering [], alignment [, ], zero-shot transfer [], bias removal [, ], cross-domain retrieval [, ], shape analysis [, , , , ], etc. [].

A number of works have been done to directly or indirectly solve the content-transformation disentangling problem [, , , , , ], whereas many of them use annotated factors for training. However, such annotation is hard to obtain, and thus, unsupervised disentanglement methods are desired. Only a few works [, , ] have addressed the problem in a completely unsupervised setting. Among them, VITAE [] performed explicit parameterization of transformation factors as diffeomorphic transformations and achieved noteworthy disentanglement of content and transformation in several image datasets. SpatialVAE [] and Harmony [] focus on representing known transformation types as transformation factors. SpatialVAE [] performed disentanglement of two-dimensional translation and rotation, whereas Harmony [] demonstrated disentanglement of a broader set of transformations that can be explicitly parameterized.

Though these unsupervised explicit parameterization methods have demonstrated several successes, they have several major limitations in general: (1) These method uses the transformation factor to generate only a few parameterized transformations. In reality, many transformations do not have a well-defined parameterized form,   e.g., protein conformation changes, viewpoint change (LineMod), human deformation during movement (Starmen), etc. These explicit methods can not disentangle such transformation by design. (2) They use Spatial Transformer Networks (STN) [] for inferring the transformation factor, which requires a continuous parameterization of the transformation to be disentangled. The continuity of many transformations (*e.g*., SO(3) rotation) in neural networks is often a concern []. Moreover, (3) these methods often infer spurious contents during disentanglement in real-world datasets(Fig. 2). Due to these limitations, an unsupervised content-transformation disentangling method without explicit transformation representation is much desired.

Consequently, in this work, we focused on the more generalized setting of unsupervised content-transformation disentanglement without any explicit parameterization of the transformation factor. Given the theoretical study by Von et al. [] showing the implicit disentanglement ability of standard Contrastive Learning (CL) methods, e.g., SimCLR [], to disentangle content from style, a plausible solution could be using SimCLR with geometric data augmentation for our implicit content-transformation disentanglement task. However, these popular CL methods [] have rarely achieved disentanglement in practice [], and we also found that SimCLR with geometric data augmentation performs poorly in our shape-focused image datasets of interest (Table 1).

To this end, we develop a novel contrastive learning method for implicit content-transformation disentanglement, termed “DualContrast”. Similar to previous content-transformation disentangling methods [, , , ], it is a deep generative model based on a variational autoencoder (VAE) architecture. However, unlike previous methods, it does not require labels or explicit parameterization of any latent factor. Instead, it uses a novel contrastive strategy that creates “positive” and “negative” contrastive pairs w.r.t. both content and transformation latent factor. We create negative pairs w.r.t. transformation factor and positive pairs w.r.t. content factor by simply rotating the samples. In addition, we randomly permute the samples in the batch to create a negative pair for a batch of samples w.r.t. content latent factor. Without prior knowledge, creating positive pairs w.r.t. the transformation latent factor from the input samples is not possible. To this end, we generate samples from an identical distribution of the transformation representations in the embedding space and treat them as positive pairs for the transformation latent factor. Upon creating both “positive” and “negative” pairs w.r.t. both factors, we penalize similarity between negative pairs and distance between positive pairs for each factor. In DualContrast, we combine these contrastive losses with VAE reconstruction loss and train the model in an end-to-end manner.

With such a strategy, by just performing simple rotation to create contrastive pairs, DualContrast disentangles a much wider range of more complicated transformations, *e.g*., human deformation in starmen, viewpoint in Linemod, the writing style in MNIST, etc. With DualContrast, we, for the first time, solved the scientific problem of disentangling protein identities from protein conformations in simulated 3D protein images (Figure 6) inside the cell, known as cellular cryo-ET subtomograms. DualContrast could successfully disentangle the protein identities as content and protein conformations as transformation factors in simulated subtomograms, which is not achievable with previous methods. Overall, our qualitative and quantitative experimental results show the efficacy of DualContrast in isolating content from transformation compared to existing relevant methods.

A summary of our contributions is as follows:

- •
We work on a challenging setting of learning unsupervised disentangled representations of content and transformation in shape-focused images by removing the requirement of explicit parameterization.
- •
We develop a novel contrastive generative modeling strategy that creates positive-negative pairs for both content and transformation latent factors. We introduced a novel way of designing contrastive pairs for transformation factors.
- •
We show that DualContrast is highly effective in disentangling a wide range of transformations from the content in various shape-focused image datasets by only using simple rotation for creating contrastive pairs.
- •
Leveraging DualContrast, we, for the first time, disentangled protein identities (*i.e*.composition) from protein conformations in simulated 3D cellular cryo-ET subtomograms, as a proof of principle.

## 2 Related Works

Disentangled Representation Learning: Our method can be regarded as a disentangled representation learning method that disentangles generative factors of variation as content-specific and transformation-specific factors. There exist many generic disentangled representation learning approaches [, , , , , ] where the factors of variations are interpreted by traversing through the learned factors and reconstructing samples. However, these methods do not perform very well in specifically disentangling content and transformation factors []. To this end, specialized methods that deal with such two factors have been introduced recently. Many of them [, , ] are partially or weakly supervised and use labels associated with one of the latent factors to train their models. However, these methods are not applicable when such supervision is unavailable beforehand, which is often the case for many real-life datasets. There are a few [, , , ] approaches that deal with the unsupervised disentanglement of content and transformation. However, these methods perform explicit parameterization of the transformation factor as affine or some specific parameterized transformation. In contrast, our method does not impose any explicit parameterization on the transformation latent factor and does not face the issues of the explicit methods (Section 1).

Contrastive Learning based Disentanglement: Contrastive learning is the primary building block of our method. Contrastive Learning with data augmentation with existing approaches, *e.g*., InfoNCE loss-based SimCLR [], etc., has been previously used by [, ] to isolate content from style. Though in several scenarios, style in shape focused images can be referred to as transformations, these works do not specifically focus on shape-focused real-world images. Moreover, our experiments demonstrate that SimCLR with rotation augmentation in a standard encoder-decoder framework leads to very poor disentanglement of content and transformation (Table 1). Moreover, [] only creates positive pairs w.r.t content and negative pair w.r.t. style with data augmentation in SimCLR. We have observed this scheme does not work well in our scenario. Unlike these works, our approach uses a novel strategy to create “negative" and “positive" pairs for both content and transformation of latent factors without any InfoNCE loss.

Unsupervised Content-style disentanglement: Apart from [], several methods [, , , , , ] exist that perform unsupervised content-style disentanglement, focusing on natural images. These methods focus on the controllable generation of realistic natural images, mostly using ImageNet pretrained models. Unlike these methods, our work primarily focuses on disentangling content and transformation in shape-focused image datasets. Moreover, we do not depend on any ImageNet pretrained models as our images of interest differ greatly from the natural images of ImageNet.

Protein Composition-Conformation Disentanglement: One of the major contributions of our work is disentangling protein identity characterized by large compositional variability as the content factor and protein states characterized by conformational variability or subtle compositional variability as transformation factor from 3D cellular cryo-electron tomography (cryo-ET) subimages or subtomograms. Previously, Harmony [] was used to disentangle large compositional variability in cellular subtomograms. HEMNMA-3D [] and TomoFlow [] were used to disentangle subtle compositional or conformational variability in subtomograms using known templates. SpatialVAE [] disentangled subtle conformational variability from 2D cryo-EM images. There exists several works on single particle cryo-EM and cryo-ET reconstruction, *e.g*., cryoDRGN2 [], cryoFIRE [], cryoAI [], etc. that performs amortized inference of transformation ($SO(3)\times d^{2}$) and latent space representing content. However, these works mainly focus on 2D-to-3D reconstruction instead of content-transformation disentanglement. Ours is the first work to deal with both large and subtle compositional variability in two different latent spaces for 3D cryo-ET subtomograms.

## 3 Method

### 3.1 Disentangling Content and Transformation

Disentangling content and transformation refers to learning one mapping $h:\mathcal{X}\rightarrow\mathcal{C}\times\mathcal{Z}$, where $\mathcal{X}$ is an input space, $\mathcal{C}$ and $\mathcal{Z}$ are the intermediary content and transformation representation spaces respectively. Here, $h$ can be decomposed as $h_{c}:\mathcal{X}\rightarrow\mathcal{C}$ and $h_{z}:\mathcal{X}\rightarrow\mathcal{Z}$ representing the content and transformation mapping, respectively. In an unsupervised setting, distinguishing between content and transformation can be tricky. In this paper, we distinguish between content and transformation with respect to a family of transformations $\mathcal{T}:\mathcal{X}\rightarrow\mathcal{X}$ such that the following two conditions hold:

- •
Condition 1: $\forall$ $T\in\mathcal{T}$, s.t., for any $\mathbf{x}\in\mathcal{X}$, $h_{c}(T(\mathbf{x}))=h_{c}(\mathbf{x})$.
- •
Condition 2: $\exists$ $T\in\mathcal{T}$, s.t., for any $\mathbf{x}^{(1)},\mathbf{x}^{(2)}\in\mathcal{X}$, $h_{z}(T(\mathbf{x}^{(1)}))=h_{z}(\mathbf{x}^{(2)})$.

Condition 1 defines content space which is invariant w.r.t. $\mathcal{T}$. On the other hand, condition 2 states that there is a $T\in\mathcal{T}$ that can make the transformation of any two samples equal when applied to any one of them. Thus it defines transformation space which is informative of $\mathcal{T}$.

For content-transformation disentangling, changes in $\mathcal{C}$ must be unaffected by changes in $\mathcal{Z}$ and vice versa. In our problem setting, no labels associated with content space $C$ or transformation space $Z$ are known. The exact family of transformation $\mathcal{T}$ is also not known beforehand. Thus, we aim to simultaneously learn the family of transformations $\mathcal{T}$ and the content factors unaltered by $\mathcal{T}$.

### 3.2 Method Overview and Notation

Our method is built upon a variational auto-encoder (VAE) based architecture (Figure 1), where we perform inference of the content and transformation factors in the data. Consider, a sample space $\mathcal{X}$, two latent distribution spaces $\Phi_{Z}$ and $\Phi_{C}$, two latent spaces $\mathcal{Z}$ and $\mathcal{C}$. We apply an encoder network $q_{\psi}:\mathcal{X}\rightarrow\mathcal{C}\times\mathcal{Z}$ on a sample space $bx\in\mathcal{X}$ to encode content distribution $\phi_{c}\in\Phi_{c}$ and transformation distribution $\phi_{z}\in\Phi_{z}$, where $(\phi_{c},\phi_{z})=q_{\psi}(\mathbf{x})$. Latent parameters $\mathbf{c}\in\mathcal{C}$, $\mathbf{z}\in\mathcal{Z}$ are drawn from distributions $\phi_{c}$ and $\phi_{z}$ respectively. These parameters are then decoded using a decoder network $p_{\theta}:\mathcal{C}\times\mathcal{Z}\rightarrow\mathcal{X}$ to produce $x_{\text{recon}}=p_{\theta}(\mathbf{c},\mathbf{z})$.

### 3.3 Variational Inference of Content and Transformation Factors

Similar to a variational auto-encoder (VAE) setting, the encoder and decoder network parameters are simultaneously optimized to maximize the variational lower bound to the likelihood $p(\mathbf{x})$ called the evidence lower bound (ELBO). ELBO maximizes the log-likelihood of the data and minimizes the KL divergence between latent distributions and a prior distribution. The objective function for this can be written as:

$$ $\begin{split}p(\mathbf{x})\geq&\mathbb{E}_{q_{\psi}(\mathbf{c},\mathbf{z}| \mathbf{x})}[\log\frac{p_{\theta}(\mathbf{x},\mathbf{c},\mathbf{z})}{q_{\psi}( \mathbf{c},\mathbf{z}|\mathbf{x})}]\\ &=\mathbb{E}_{q_{\psi}(\mathbf{c},\mathbf{z}|\mathbf{x})}\log p_{\theta}( \mathbf{x}|\mathbf{c},\mathbf{z})-KL(q_{\psi}(\mathbf{c}|\mathbf{x})||p( \mathbf{c}))-KL(q_{\psi}(\mathbf{z}|\mathbf{x})||p(\mathbf{z})).\end{split}$ (1) $$

The log-likelihood term is maximized using a reconstruction loss $L_{\text{recon}}$. The prior for $\mathbf{c}$ and $\mathbf{z}$, $p(\mathbf{c})$ and $p(\mathbf{z})$ respectively are assumed to be standard normal distribution. In short, we write the KL terms as $L_{\text{KL}(c)}$ and $L_{\text{KL}(z)}$. To regularize the effect of prior on the final objective separately for two latent factors, we scale the two losses with two hyperparameters $\gamma_{c}$ and $\gamma_{z}$. Our loss function for variational inference becomes:

$$ $\displaystyle L_{\text{vae}}=L_{\text{rec}}+\gamma_{c}L_{\text{KL}(\mathbf{c}) }+\gamma_{z}L_{\text{KL}(\mathbf{z})}.$ (2) $$

We observe that by deactivating the KL term associated with any one of the factors, the likelihood tends to be optimized through the other latent factor. For instance, if $\gamma_{c}$ is set as 0 0 and $\gamma_{z}$ has a large value, then the $L_{\text{rec}}$ is optimized through $\mathbf{c}$ only. As a consequence, $\mathbf{c}$ captures all the information in $\mathbf{x}$, whereas $\mathbf{z}$ becomes totally independent of the data. In such a scenario, $\mathbf{z}$ will be a random factor with no relation to the transformation factor of the data, which is not desirable.

### 3.4 Creating Contrastive Pair with respect to Content Latent Factor

We adopt a contrastive strategy for content-transformation disentanglement. We first develop a strategy to impose a structure to the content latent factor with contrastive learning. To this end, we create negative and positive pairs for each data sample w.r.t. the content factor.
Hence, we randomly pick two samples $\mathbf{x}$ and $\mathbf{x}_{\text{neg}(\mathbf{c})}$, where $\mathbf{x},\mathbf{x}_{\text{neg}(\mathbf{c})}\in\mathcal{X}$ and regard them as negative pairs of each other w.r.t the content factor. Such strategy is commonly used in traditional self-supervised learning approaches [, ]. Considering the high heterogeneity present in the dataset, evidenced by a large number of classes, this approach accounts for the likelihood that any two randomly selected samples will have different content. To generate positive pairs w.r.t. the content factor, we pick any sample $\mathbf{x}$ and transform it with a transformation $T\in\mathcal{T}$ to generate $\mathbf{x}_{\text{pos}(\mathbf{c})}$. We use the encoder $q_{\psi}$ to generate content factors $\mathbf{c}$, $\mathbf{c}_{\text{pos}}$, and $\mathbf{c}_{\text{neg}}$ respectively from $\mathbf{x}$, $\mathbf{x}_{\text{pos}(\mathbf{c})}$, and $\mathbf{x}_{\text{neg}(\mathbf{c})}$. We use a contrastive loss to penalize the distance between $\mathbf{c}$ and $\mathbf{c}_{\text{pos}}$ and the similarity between $\mathbf{c}$ and $\mathbf{c}_{\text{neg}}$. Our contrastive loss for content latent factor can be written as:

$$ $\displaystyle L_{\text{con}(\mathbf{c})}=L_{\text{dist}}(\mathbf{c},\mathbf{c} _{\text{pos}})+L_{\text{sim}}(\mathbf{c},\mathbf{c}_{\text{neg}})$ (3) $$

To implement $L_{\text{dist}}$ and $L_{\text{sim}}$, we use mean absolute cosine distance and mean absolute cosine similarity, respectively.

### 3.5 Creating Contrastive Pair with respect to Transformation Latent Factor

Imposing contrastive loss on the content latent factor satisfies condition 1 of the content-transformation disentanglement (Section 3.1). However, the content-transformation disentanglement remains only partial since it does not satisfy condition 2. To solve this issue, we design another contrastive loss w.r.t. the transformation latent factor. Designing a contrastive loss that explicitly enforces condition 2 is not possible, given the transformation that may equalize the transformation factor of two samples can not be approximated. To this end, we design a contrastive loss that implicitly encourages condition 2. We validated the design experimentally in Section 4.

The negative pair of a sample $\mathbf{x}\in\mathcal{X}$ w.r.t. transformation latent factor is generated while creating a positive pair of it w.r.t content factor. This is because a transformation $T$ on $\mathbf{x}$ alters its transformation, but not its contents. So, we regard $\mathbf{x}_{\text{pos}(\mathbf{c})}$ as $\mathbf{x}_{\text{neg}(\mathbf{z})}$.
However, creating positive pairs of samples w.r.t. transformation is very challenging as they can not be created in the input space, unlike others. To this end, we adopted a strategy of creating pairs of samples from the same distribution in the transformation factor embedding space. We randomly sample $\mathbf{z}^{(1)}$ and $\mathbf{z}^{(2)}$ from $\mathcal{N}(0,1)$ and use the decoder $p_{\theta}$ to generate samples $\mathbf{x}^{(1)}_{\text{pos}(\mathbf{z})}$ and $\mathbf{x}^{(2)}_{\text{pos}(\mathbf{z})}$ from them (Fig. 1 (c)). In the decoder, as content input, we use content factor $\mathbf{c}^{(1)}$ and $\mathbf{c}^{(2)}$ obtained by encoding two random samples $\mathbf{x}^{(1)}$ and $\mathbf{x}^{(2)}$. Consequently, $\mathbf{x}^{(1)}_{\text{pos}(\mathbf{z})}=p_{\theta}(\mathbf{c}^{(1)},\mathbf{
z}^{(1)})$ and $\mathbf{x}^{(2)}_{\text{pos}(\mathbf{z})}=p_{\theta}(\mathbf{c}^{(2)},\mathbf{
z}^{(2)})$ serve as the positive pairs w.r.t. transformation latent factor. We then use the encoder $q_{\psi}$ to generate $\mathbf{z}$, $\mathbf{z}_{\text{neg}}$, $\mathbf{z}^{(1)}_{\text{pos}}$, and $\mathbf{z}^{(2)}_{\text{pos}}$ from $\mathbf{x}$, $\mathbf{x}_{\text{neg}(\mathbf{z})}$, $\mathbf{x}^{(1)}_{\text{pos}(\mathbf{z})}$, and $\mathbf{x}^{(2)}_{\text{pos}(\mathbf{z})}$ respectively. We use a contrastive loss to penalize the distance between $\mathbf{z}^{(1)}_{\text{pos}}$ and $\mathbf{z}^{(2)}_{\text{pos}}$ and similarity between $\mathbf{z}$ and $\mathbf{z}_{\text{neg}}$. Our contrastive loss for transformation latent factor can be written as:

$$ $\displaystyle L_{\text{con(z)}}=L_{\text{dist}}(\mathbf{z}^{(1)}_{\text{pos}}, \mathbf{z}^{(2)}_{\text{pos}})+L_{\text{sim}}(\mathbf{z},\mathbf{z}_{\text{neg }})$ (4) $$

Similar to $L_{\text{con}(\mathbf{c})}$, we use mean absolute cosine distance and mean absolute cosine similarity to implement $L_{\text{dist}}$ and $L_{\text{sim}}$, respectively.

Objective Function:
In summary, we train the encoder and decoder networks by simultaneously minimizing the loss components. Our overall objective function to minimize is:

$$ $\displaystyle L=L_{\text{vae}}+L_{\text{con(c)}}+L_{\text{con(z)}}$ (5) $$

## 4 Experiments & Results

**Table 1: Quantitative Results of unsupervised Content-Transformation Disentangling Methods on MNIST and LineMod datasets. This evaluation only shows the disentanglement performance w.r.t ground truth content factors. We report mean $\pm$ std. dev. over model training by setting 3 different random seeds. Here, for Harmony and SpatialVAE, the transformation latent factor is restricted to dimension 3. For the other methods, there is no such restriction on the transformation latent factor. For SimCLR, (Disc) means discriminative model (without any reconstruction), and (Gen) means Generative model (with reconstruction).**
|  | MNIST | LineMod |  |  |
| --- | --- | --- | --- | --- |
| Method | $D_{\text{score}}(c|c)(\uparrow)$ | $D_{\text{score}}(c|s)(\downarrow)$ | $D_{\text{score}}(c|c)(\uparrow)$ | $D_{\text{score}}(c|s)(\downarrow)$ |
| SpatialVAE [] | 0.81 $\pm$ 0.01 | 0.28 $\pm$ 0.02 | 0.95 $\pm$ 0.02 | 0.32 $\pm$ 0.01 |
| Harmony [] | 0.82 $\pm$ 0.01 | 0.31 $\pm$ 0.02 | 0.90 $\pm$ 0.01 | 0.56 $\pm$ 0.02 |
| SimCLR (Disc) with rotation [] | 0.58 $\pm$ 0.01 | 0.60 $\pm$ 0.01 | 0.62 $\pm$ 0.01 | 0.40 $\pm$ 0.01 |
| SimCLR (Gen) with rotation [] | 0.53 $\pm$ 0.01 | 0.67 $\pm$ 0.02 | 0.61 $\pm$ 0.02 | 0.79 $\pm$ 0.01 |
| VAE with two latent space | 0.63 $\pm$ 0.00 | 0.63 $\pm$ 0.00 | 0.87 $\pm$ 0.00 | 0.87 $\pm$ 0.00 |
| VITAE [] | 0.77 $\pm$ 0.00 | 0.32 $\pm$ 0.00 | 0.92 $\pm$ 0.00 | 0.90 $\pm$ 0.00 |
| DualContrast (w/o $L_{\text{con(c)}}$) | 0.87 $\pm$ 0.02 | 0.21 $\pm$ 0.02 | 0.86 $\pm$ 0.01 | 0.31 $\pm$ 0.02 |
| DualContrast (w/o $L_{\text{con(z)}}$) | 0.79 $\pm$ 0.04 | 0.85 $\pm$ 0.02 | 0.79 $\pm$ 0.03 | 0.86 $\pm$ 0.03 |
| DualContrast | 0.89 $\pm$ 0.01 | 0.31 $\pm$ 0.02 | 0.95 $\pm$ 0.01 | 0.48 $\pm$ 0.02 |

Benchmark Datasets: We report results on MNIST [], LineMod [], Starmen [], and a realistically simulated protein subtomogram dataset. Only LineMod features RGB images among these datasets, while the other three consist of grayscale images. The subtomogram is a 3D volumetric dataset, whereas the others are 2D images.

Baselines: We used the unsupervised content-transformation disentanglement methods, SpatialVAE [], Harmony [], VITAE [] with explicit parameterization and InfoNCE loss based standard self-supervised learning method with rotation augmentation [] as baseline approaches. Several generic disentangled representation learning methods, such as beta-VAE, Factor-VAE, beta-TCVAE, DIP-VAE, etc., are available that do not aim to disentangle specific factors such as content or transformation in the data but rather aim to disentangle all factors of variation. However, previous studies [, ] have shown that such methods perform worse than methods specifically aiming to disentangle content and transformation. Therefore, we did not use these methods as baselines in our experiments.

Implementation Details: We implemented the encoder using a Convolutional Neural Network (CNN) and the decoder using a Fully Connected Network (FCN). Our experiments used the same latent dimension for content and transformation factors, except for Harmony [] and SpatialVAE [], where a specific dimension needs to be used for transformation latent factors. We train our models for $200$ epochs with a learning rate of $0.0001$ in NVIDIA RTX A500 GPUs and AMD Raedon GPUs. We optimize the model parameters with Adam optimizer.

Evalutation Metrics: We assess DualContrast based on two criteria: (1) the quantitative disentanglement of content and transformation factors and (2) the qualitative capability of image generation through content-transformation transfer.

Figure: (a)
Refer to caption: x2.png

For (1), there exists several metrics, *e.g*., MIG score, $D_{\text{score}}$, etc. Locatello et al [] demonstrated that these metrics are highly correlated. To this end, similar to [], we only use $D_{\text{score}}$ to measure the disentanglement. $D_{\text{score}}$ is simply a measurement of how well the ground truth factors can be predicted from the corresponding latent factors. In our scenario, there are four such quantities - (a) predictivity of content labels given content factor $D_{\text{score}}(c|c)$ or transformation factor $D_{\text{score}}(c|s)$ (b) predictivity of transformation labels given content factor $D_{\text{score}}(s|c)$ or transformation factor $D_{\text{score}}(s|s)$. For MNIST and LineMod, we do not have any transformation labels, so we only reported values for (a). Similar to [], we use linear Logistic Regression Classifiers to measure predictivity ($D_{\text{score}}$) across all our experiments.

Figure: Figure 3: tSNE embedding plots of content latent factor learned by the unsupervised content-style disentanglement methods. (a), (b), (c), and (d) shows the results for Harmony [], Spatial-VAE [], C-VITAE [], and DualContrast respectively. Overall, DualContrast shows superior performance.
Refer to caption: x6.png

The quantitative measurement, particularly when transformation labels are not available, might be misleading. To this end, we perform qualitative evaluation by generating images while transferring content and transformation from original images from the dataset. We take a set of $K$ images $\mathbf{x}$ and create a $K\times K$ grid of $K^{2}$ images using the different combinations of content and transformation factors from those $K$ images. In the grid, we keep the content factor and transformation factor consistent across the columns and the rows respectively. In other words, the image at position $(i,j)$ in the grid is generated using the content of image $x_{j}$ and transformation of image $x_{i}$. For the protein subtomogram dataset, we also provide qualitative results through embedding visualization (Figure 8).

### 4.1 DualContrast disentangles writing style from MNIST handwritten digit images

Similar to [, , ], we start our experiments with the MNIST dataset of handwritten digits. We performed a quantitative comparison between DualContrast and the baseline approaches, as shown in Table 1. We observed that DualContrast achieves higher $D_{score}(c|c)$ compared to other explicitly parameterized approaches [, , ].

Figure: (a)
Refer to caption: x7.png

However, since transformation labels are not present for MNIST, relying entirely on the content prediction performance for disentanglement might be misleading. For example, the degenerate solution achieves even lower $D_{score}(c|s)$ as transformation does not represent anything. To this end, we report qualitative results for baseline methods and our approach DualContrast (Fig. 2). In generating images with varying content and styles, we observe that Harmony, SpatialVAE, and C-VITAE generate many images with erroneous content and styles. On the other hand, DualContrast produces a few minor errors and no extreme mistakes. We also include tSNE embedding of the content factors inferred by the models on the MNIST test dataset associated with class labels (Figure 3). On the embedding space, DualContrast clearly shows superior clustering performance. All of these suggest better content-transformation disentanglement with our approach than existing explicit parameterization approaches.

### 4.2 DualContrast Disentangles ViewPoint from Linemod Object Dataset

Figure: Figure 5: Qualitative Results of Unsupervised Content-Transformation Disentanglement on Starmen (a) and (b) shows the results obtained with VITAE [] and DualContrast, respectively.
Refer to caption: x9.png

LineMod is an RGB object recognition and 3D pose estimation dataset [] with 15 unique object types. It contains 15 unique objects: ‘ape’, ‘bench vise’, ‘bowl’, ‘cam’, ‘can’, ‘cat’, ‘cup’, ‘driller’, ‘duck’, ‘eggbox’, ‘glue’, ‘hole puncher’, ‘iron’, ‘lamp’ and ‘phone’, photographed in a highly cluttered environment. We use a synthetic version of the dataset [], which has the same objects rendered under different viewpoints. The dataset is publicly available at [this url](https://bop.felk.cvut.cz/media/data/bop_datasets/lm_train.zip). The dataset has $1,313$ images per object category. We used $1,000$ images per category for training and the remaining for testing.

We evaluated the disentanglement scores $D_{\text{score}}$ (Table 1) for DualContrast and baseline approaches, where DualContrast clearly shows the best performance. We used the object identities as content labels during the evaluation.

We also performed qualitative measurement with image transfer and depicted the result of DualContrast and its most relevant (in terms of the generality of transformations) VITAE in Fig 4. We observe that VITAE distorts the images significantly while performing image generation with content-style transfer. On the other hand, DualContrast does not face such issues. Moreover, DualContrast can disentangle different viewpoints of the objects as the transformation factor using only 2D rotation to create contrastive pairs in its model.

### 4.3 DualContrast disentangles Human Deformation from starmen shapes dataset

Starmen is a 2D synthetic shape dataset of the deformable human body [] consisting of $10,000$ shape-focused images, which is publicly available at [https://zenodo.org/records/5081988](https://zenodo.org/records/5081988) under CC-BY license. Each starmen image is a grayscale image of size $64\times 64$. They are generated by deforming five control points (two on arms, two on legs, one on head) on template human shapes.

We randomly took $6,000$ starmen images to train our models and kept the remaining $4,000$ for testing. We report the qualitative results by VITAE [] and DualContrast in Figure 5. From the image generation results with the content-transformation transfer, we can observe that DualContrast disentangles the deformation of starmen shapes as transformation and the identity (based on leg thickness) as content, providing a smooth image generation result. On the other hand, VITAE [] does not perform such meaningful content-transformation disentanglement, generating implausible results.

### 4.4 DualContrast Disentangles Protein Identity from states in Cryo-ET Subtomograms

Figure: Figure 6: (a) Isosurface visualization of Nucleosome Density Maps. Each density map slightly varies in terms of conformation. They are associated with the PDB IDs in the figure. (b) 3D Slice by slice visualization of several subtomograms of the corresponding nucleosome density maps.
Refer to caption: x10.png

We created a realistic simulated cryo-ET subtomogram dataset of 18,000 subtomograms of size $32^{3}$. The dataset consists of 3 protein classes of similar sizes- Nucleosomes, Sars-Cov-2 spike protein, and Fatty Acid Synthase Unit. These proteins are significantly different in terms of their composition, which determines their different identities. Moreover, structures for all of these three types of proteins have been resolved in cellular cryo-ET [, , ], which makes it feasible to use them for our study. Moreover, cellular cryo-ET is the primary method to capture these proteins inside the cells in their native state [].

For each protein class, we collected 6 different protein structures from the RCSB PDB website []. RCSB PDB is a web server containing the structure of millions of proteins. For nucleosomes, we collected PDB IDs ‘2pyo’, ‘7kbe’, ‘7pex’, ‘7pey’, ‘7xzy’, and ‘7y00’. All of these are different conformations of nucleosomes that slightly vary in composition (Figure LABEL:fig:supp:nuc-dmaps). For sars-cov-2 spike proteins, we collected PDB IDs ‘6vxx’, ‘6vyb’, ‘6xr8’, ‘6xra’, ‘6zox’, and ‘6zp0’. Among them, only ‘6xra’ shows much variation in structure from the other ones, and the rest of the PDB IDs are very similar in structure. For Fatty Acid Synthase (FAS) Unit, we used PDB IDs ‘8prv’, ‘8ps1’, ‘8ps9’, ‘8psj’, ‘8psm’, ‘8psp’. They also vary very slightly in terms of the structure.

After collecting these 18 PDB structures as PDB files, we used EMAN PDB2MRC [] to create density maps (as MRC file extension) from the PDB files. We create density maps of size $32^{3}$ with $1$ nm resolution. We then randomly rotate and translate each density map and create $1000$ such copies. We then convolve the density maps with CTF with CTF parameters common in experimental datasets (Defocus -5 nm, Spherical Abberation 1.7, Voltage 300 kV). Afterward, we add noise to the convolved density maps so that the SNR is close to $0.1$. Thus, we prepare $18,000$ realistic subtomograms with 3 different protein identities, each with 6 different conformations. We uploaded the entire dataset anonymously at [https://zenodo.org/records/11244440](https://zenodo.org/records/11244440) under CC-BY-SA license. 3D slice-by-slice visualizations for sample subtomograms for nucleosomes, spike proteins, and FAS units are provided in Figure 6.

We trained DualContrast and baseline approaches against the protein subtomogram dataset. We could not train VITAE [] on subtomogram datasets since it did not define any transformation for 3D data. Designing CPAB transformation for 3D data by ourselves was challenging. Among the other baseline approaches, spatialVAE could not distinguish the protein identities with high heterogeneity at all, which is evident by its embedding space UMAP (Figure 7). Only Harmony provided somewhat plausible results apart from DualContrast.

Figure: Figure 7: UMAP of content factor learned by SpatialVAE [] on our subtomogram dataset. The protein identities are indistinguishable from the embedding space.
Refer to caption: extracted/5622084/figs/SpatialVAE_subtomo_content_umap.png

We observed that the DualContrast content factor perfectly clustered the three protein identities from the dataset (GMM clustering accuracy $0.99$), profoundly disentangling the transformations (Figure 8(b)). We do not see such perfect clustering for Harmony (Figure 8(a)); rather, protein identities with very different compositions get mixed up due to their conformational variation (GMM clustering accuracy $0.70$). We further observe that DualContrast transformation embedding clusters have similar nucleosome conformations, showing disentanglement of conformation from protein identity. On the other hand, in Harmony, the transformation factor only represents rotation and translation by design and can not capture nucleosome conformations. Moreover, even its content embedding for manually selected (not through automatic clustering) nucleosome subtomograms can not show good clustering of conformations (Figure 8(b)).

Figure: Figure 8: Disentanglement of composition and conformations in cellular subtomogram dataset. (a) UMAP embedding of the content factor of all subtomograms (top) and only nucleosome subtomograms (bottom) with Harmony [] (b) UMAP embedding of the content factor (top) and transformation factor (bottom) obtained by DualContrast.
Refer to caption: x14.png

Ablation Study: To evaluate the individual contribution of the contrastive losses, we conduct both quantitative and qualitative ablation analyses of DualContrast. We trained (1) DualContrast without any contrastive loss, which is basically a VAE with two latent spaces, (2) DualContrast with only $L=L_{\text{VAE}}+L_{\text{con (z)}}$, and (3) $L=L_{\text{VAE}}+L_{\text{con (c)}}$. We qualitatively and quantitatively evaluated each model. We show the quantitative results in Table 1 and qualitative results of MNIST in Figure 9.

Figure: (a)
Refer to caption: x15.png

For (1), the $D_{\text{score}}$ is almost the same for both content and transformation factors, indicating equal predictivity of the digit classes by both factors. This is obvious given that the model has no inductive bias to make different factors capture different information. In model (2), using contrastive loss w.r.t. only transformation factor $z$ makes it uninformative of the data. Thus, it provides a small $D(c|z)$ score as desired, but the changing $z$ does not affect the image generation (Figure 9). On the other hand, in the model (3), using contrastive loss w.r.t. only content factor $c$ gives a less informative content factor, a lower $D(c|c)$ score, and a higher $D(c|z)$ score, contrary to what is desired. These results indicate contrastive loss w.r.t to both factors is crucial to get the desired disentanglement.

## 5 Limitation

Similar to existing content-transformation disentanglement methods [, , , , ], DualContrast is applicable to data samples with single objects. This is reasonable, given content and transformation can not be defined for data with multiple objects.
Moreover, we limited our study only to shape-focused image datasets, motivated by the problem of conformation disentanglement of protein subtomograms. Consequently, we only used rotation to create contrastive pairs, which worked pretty well in our cases. For image datasets where appearance or illumination plays a crucial role, using only rotation to create contrastive pairs may not work well. However, disentanglement in natural image datasets is a different scenario that is outside of the scope of this paper.

## 6 Conclusion

This work focuses on a challenging setting of the unsupervised content-transformation disentanglement problem in shape-focused image datasets, where the transformation latent factor is not explicitly parameterized. To tackle this problem, we propose a novel method termed DualContrast. DualContrast employs generative modeling with a novel contrastive learning strategy that creates positive and negative pairs for content and transformation latent factors. Our extensive experiments show DualContrast’s effective disentanglement of challenging transformation factors across various shape-based image datasets; including simulated cellular subtomograms, where it solved the unexplored problem of isolation of protein conformations from protein identities, as proof of principle.

## 7 Acknowledgement

This work was supported in part by U.S. NIH grants R01GM134020 and P41GM103712, NSF grants DBI-1949629, DBI-2238093, IIS-2007595, IIS-2211597, and MCB-2205148. This work was supported in part by Oracle Cloud credits and related resources provided by Oracle for Research, and the computational resources support from AMD HPC Fund. MRU were supported in part by a fellowship from CMU CMLH.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 1
presentation: 2
questions: Please refer to the Weaknesses section.
soundness: 2
strengths: 1. This paper proposes well-defined conditions for the disentanglement of content and transformation. 
2. The experiments are conducted on four datasets, and comprehensive qualitative results are shown.
summary: This paper proposes an unsupervised disentangling method to disentangle content and transformation of the input. Specifically, this paper first proposes two conditions that disentanglement of content and transformation should satisfy. Then, this paper proposes a method to construct positive and negative samples with respect to both content and transformation. The key idea is to utilize a variational autoencoder to construct these samples. The experiments are conducted on four datasets, i.e., three of them (mnist, linemod, and starmen) are pure images, and one is protein subtomogram. One quantitative result and several qualitative results are shown to prove the effectiveness of the proposed method.
weaknesses: The main concern of this paper is evaluation, which is insufficient and less significant.
1. The first three datasets (mnist, linemod, starmen) are somehow toy datasets, which is less significant in real-world applications.
2. I agree that protein conformation is one meaningful real-world application, but other than map visualizations, it fails to produce convincing evaluation results.
3. There lacks some widely used evaluating metrics in Table 1 to demonstrate the application of the disentanglement.
4. This paper also does not provide comparisons with other baseline methods or state-of-the-art methods.

--------------------------------------

Post-rebuttal:
Thanks for the authors' response. The revision includes more quantitative results and qualitative comparisons, which is appreciated. These additions partially resolved my concerns. However, I noticed that these compared methods are still up to date in 2022, which cannot be considered state-of-the-art to some extent. 

Another issue to mention is that the metrics D(c|c) and D(c|z) were used in the initial submission, and the protein experiments are the most important one (in my personal view); it is usually discouraged not to report the main quantitative results of the critical experiments in the initial submission, but added during rebuttal.

Considering the rebuttal and revision, I raised my rating to 5. Please use it sparsely.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: In the method section, condition 1 is very confusing. I think it was meant to be "for all $T \in T$ and $x \in X, h_c(T(x)) = h_c(x)$", but please correct me if I misunderstood this.

In the method section, many terms are used seemingly interchangeably, such as "latent space", "factor", "representation", "transformation", "content". Please clarify these terms. For example, at line 206, "transformation" is used, however, I think it was meant to be "transformation representation" or "transformation factor", unless I misunderstood.

There are a few grammatical and syntactical mistakes, such as inconsistencies in the use of uppercase and lowercase, sometimes writing "shape focused" while other times "shape-focused".
soundness: 3
strengths: Originality.
In this work, a novel method to address the problem of creating positive pairs of transformations under content change has been proposed. The core of this work is original.

Quality.
The method proposed was evaluated on a sufficient number of datasets. Although not very complex, they could suffice in showing the potential for this approach. The baselines chosen are also relevant to the method proposed.

Clarity.
The figures in the experiment section allow for a quick qualitative assessment of the performance of the methods. The method explanation is quite clear.

Significance.
This work and the proposed method have shown some potential for successful applications.
summary: Unsupervised disentanglement of transformations and content is a challenging task that was previously approached primarily through using separate ad-hoc transformation methods, or by self-supervised contrastive-based methods. Ad-hoc transformations suffer from being limited to the given parameterization chosen, while self-supervised methods do not tackle this disentanglement problem directly. In this work, DualContrast is proposed, which consists of a VAE with additional contrastive losses designed to disentangle content and transformation. The hardest challenge is obtaining positive pairs of samples with respect to transformations: changing the content while keeping the transformation constant. In this work, this has been done by decoding two random samples from the prior of the transformation latent space while feeding different permutations of the content latent representation to obtain similar transformations with different content. The method is applied to MNIST, LineMod, Starmen Shapes, and Cryo-ET subtomograms with positive results.
weaknesses: The explanation of the method in the abstract and introduction is especially unclear. This is also a problem because Figure 2 fails to properly and intuitively show the method. Reading the method section explains this more. To improve, I would suggest clearly highlighting  the role of the latent space in the creation of the positive pair that would otherwise be impossible. This could be done similarly to how it was done in Figure 3. Figure 2 would then become useful. Additionally, Figure 2 lacks proper annotations such as labeling of all elements present, and proper caption explaining what happens in the figure in a more complete way. There is some inconsistency in how things are called. In the figure, style, and content are mentioned, however, in the text it is clear that "style" is supposed to be "transformation", please pick one and stick with it in the whole manuscript, either one would suffice, however, transformation is likely to be more accurate.

The contributions are a bit bold. The first contribution, especially, is more context for the work than a contribution and could be removed entirely. Please consider reworking the contributions to be more reflective of the actual content.

The related work section should expand a bit more on the protein part, which is currently very unclear for somebody who is not a practitioner. Please provide more examples, even referring to the appendix to understand the data and the context better.

The method section has a few mistakes and the explanation is very wordy, which makes it hard to follow. I wrote a few observations in the questions section of this review.

The manuscript should include the limitations of this method, especially regarding the latent space-based approach to creating positive transformation pairs. For example, the limitations should address whether this approach could be extended to real-world datasets or whether this approach should be limited to specific types of datasets. 

The experiments lack in quantitative results. Although disentanglement is very hard to measure, given the ability to choose datasets, it would be much more convincing to have datasets where a quantitative assessment is possible either in the form of direct supervision (similar to what the disentanglement metric is currently doing), or through some downstream tasks where the disentanglement would be useful. Such tasks could be segmentation, or visual question answer.
The choice of the "human deformation" as a dataset is very confusing, and the results reported are also very underwhelming. Although the generated shapes are better, the data appears to be very trivial, so some more information on the training and the difficulty of fitting such samples would be more convincing. The number of parameters, ablation performed, and results from more baselines would be a step in the right direction. It is especially important to keep including all baselines for all datasets used, the results on human deformation appear to be unfinished. If so, it would have been better to simply exclude the dataset from the manuscript.
Additionally, plots of the latent space are only available for the cellular dataset

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - On L147, the authors say that they model "is highly effective in disentangling wide range of transformations from the content in various shape-focused image datasets by only using simple rotation for creating contrastive pairs since the representation for disentangled rotation generalizes over other shape transformations." Could they elaborate on this? Why is it the case? Where is it shown on the paper? How can we be sure it will work to other modalities besides the toy tasks tested?
- Is the VAE trained at the smae time as the contrastive losses? Since the VAE is used to generated samples for the CLs, how do training jointly vs training in two stages (VAE followed by CLs) change the performance?
soundness: 2
strengths: + The paper is well written and easy to follow
+ The idea of disentangling features is an important problem in many applications of machine learning
+ The proposed approach is simple and well motivated
summary: In this paper, the authors propose a model that learn unsupervised representation for  "shape-focused images". In particular, their method, DualContrast, learn to disentangle "content" and "transformation" in an unsupervised fashion. The model is trained with an a combination of 2 contrastive losses (one for context, one for transformations) and a VAE loss (the VAE is used to sample positive samples for the CL of transformations). The authors show results of the proposed model on multiple small/toysh datasets.
weaknesses: -  The author often mention that the work focus on "shape-focused real-world images", but they only applied in very simplified, toysh settings, very far from "real-world images". Even the CryoEM task is a very simplified task.
- The choice of positive/negative samples for each factor is very ad-hoc. The paper lacks explanation and empirical validation on why this choice makes sense versus others. 
- I found very strange the choice of using VAE generated samples as data to train the contrastive loss. This idea of using generated smples to train a model is not well understood. This approach might have worked in the very simplified tasks tested on the paper, but It is very unlikely that the proposed model would work on any real-world dataset.
- I also find the experimental results a bit weak. First, the datasets utilized in this work are very simple and results on them probably wont guarantee their utility on real-world problems. Second, the metrics utilized on Table 1 are not particularly significant. Third, most of the results are qualitative and based on one or two images that can potentially be cherry picked.

Meta-review: This paper utilizes a VAE with dual latent spaces for disentangling of content and transformations of cellular 3D protein images. Apart from standard VAE loss, the novel part is contrastive learning loss based on positive and negative pairs in terms of content and transformation. The method here is unsupervised, and those positive and negative pairs are generated without supervision or labels. Clustering of content and transformation codes clearly shows the effectiveness of the method as shown in Fig. 6 in the revised version.
However, reviewers were concerned about real-world impact on downstream tasks. Clustering of latents is another way of visualization for the latents, but not really an application with a profound impact. (There might be a knowledge gap for the reviewers and AC who do not work on biological science).
Another weakness is the lack of comparison to more recent methods on disentangled representation learning as mentioned by the reviewer. Another metric was added in the rebuttal, but overall the amount of quantitative results seems insufficient to be convincing.
The way to construct positive pairs and negative pairs is highly unreliable. Randomly selecting pairs of samples can lead to positive pairs with the same content. Also, it is not clear how to bootstrap the joint network to produce synthetic images with the same content but different transformations, in particular at the beginning of training when the network cannot produce images with high fidelity.

The authors are encouraged to improve the empirical significance of the method and resubmit it to the next conference or to a more domain-specific venue where the impact can be better appreciated.

Decision: Reject