Explicitly disentangling image content from translation and rotation with spatial-VAE

TL;DR

This paper introduces spatial-VAE, a novel variational autoencoder framework that explicitly disentangles image content from pose variables like rotation and translation, leading to improved reconstruction and interpretability.

Contribution

The paper proposes a new method for explicitly disentangling pose from content in VAEs by formulating the generative model as a function of spatial coordinates, enabling better control over latent factors.

Findings

Spatial-VAE effectively disentangles rotation and translation from content.

The method improves reconstruction quality over standard VAEs.

Applications include modeling 2-D views of proteins and galaxies.

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.