[Paper Review] Explicitly disentangling image content from translation and rotation with spatial-VAE

Tristan Bepler, Ellen D. Zhong+3 more|arXiv (Cornell University)|Sep 25, 2019

Advanced Electron Microscopy Techniques and Applications34 citations

TL;DR

Spatial-VAE disentangles rotation/translation from content in a VAE by making the generator depend on pixel coordinates, enabling explicit pose inference and improved reconstruction on transformed images, EM, and galaxy data.

ABSTRACT

Given an image dataset, we are often interested in finding data generative factors that encode semantic content independently from pose variables such as rotation and translation. However, current disentanglement approaches do not impose any specific structure on the learned latent representations. We propose a method for explicitly disentangling image rotation and translation from other unstructured latent factors in a variational autoencoder (VAE) framework. By formulating the generative model as a function of the spatial coordinate, we make the reconstruction error differentiable with respect to latent translation and rotation parameters. This formulation allows us to train a neural network to perform approximate inference on these latent variables while explicitly constraining them to only represent rotation and translation. We demonstrate that this framework, termed spatial-VAE, effectively learns latent representations that disentangle image rotation and translation from content and improves reconstruction over standard VAEs on several benchmark datasets, including applications to modeling continuous 2-D views of proteins from single particle electron microscopy and galaxies in astronomical images.

Motivation & Objective

• Motivate learning latent representations of images where semantic content is separated from pose transformations like rotation and translation.
• Propose a spatially conditioned VAE to infer rotation and translation jointly with content in an unsupervised manner.
• Demonstrate improved reconstruction and true-content disentanglement on benchmarks and domain datasets (MNIST variants, cryo-EM simulations, Galaxy Zoo).

Proposed method

• Define a spatial generator network where the image likelihood p_g(y^i|x^i,z) is computed by an MLP taking coordinates x^i (and rotated/transformed versions) and latent z.
• Transform coordinates to model rotation and translation directly in the coordinate space, so p_g(y|z,θ,Δx) depends on x^i R(θ) + Δx.
• Infer rotation θ, translation Δx, and unstructured latent z via an approximate posterior q(z,θ,Δx|y) = N(μ(y),σ(y)^2 I) within a VAE objective.
• Use a variational lower bound that combines expected log-likelihood with KL terms for z, θ, and Δx; provide a KL alternative for θ when a uniform prior is desired (Eq. 3).
• Provide a training scheme where θ and Δx are allowed to vary during inference, enabling end-to-end learning of transformed and content factors.
• Experimentally compare spatial-VAE to a vanilla VAE and a variant with fixed pose to demonstrate disentanglement and reconstruction improvements.

Experimental results

Research questions

• RQ1 Can a VAE be trained to disentangle continuous pose factors (rotation and translation) from semantic content without supervision?
• RQ2 Does explicitly modeling coordinate-based generation and pose inference improve reconstruction when images are randomly rotated and translated?
• RQ3 Can the framework recover true content variation in domain datasets such as simulated cryo-EM projections and galaxy images while remaining invariant to pose?
• RQ4 Is the learned latent space able to separate content from rotation/translation and reflect meaningful interpolations across domains?

Key findings

• Spatial-VAE outperforms vanilla VAEs in ELBO on transformed MNIST variants, especially when the unstructured latent dimension is small.
• The model learns to separate content (z) from rotation θ and translation Δx, demonstrated by learned data manifolds where pose is removed from content representations.
• In simulated cryo-EM data, spatial-VAE with one unconstrained latent variable plus pose inference captures conformational variation with orientation factored out, outperforming vanilla VAEs.
• On Galaxy Zoo and EM datasets, spatial-VAE learns transformation-invariant representations that reflect galaxy morphology and protein conformations independently of rotation/translation.
• Interpolations in the latent space show smooth transitions in content while pose remains disentangled, indicating a robust separation of factors.