[Paper Review] Topographic VAEs learn Equivariant Capsules

Timothy A. Keller, Max Welling|arXiv (Cornell University)|Sep 3, 2021

Image Processing and 3D Reconstruction45 references6 citations

TL;DR

This paper introduces the Topographic Variational Autoencoder (TVAE), a deep generative model that learns topographically organized, approximately equivariant capsule representations from unlabeled sequential data by enforcing temporal coherence. By training on transformed input sequences—such as rotated or color-shifted MNIST digits—the model discovers disentangled, equivariant features that generalize to unseen transformations, achieving higher likelihood and quantitatively verified equivariance through commutativity of inference and transformation operators.

ABSTRACT

In this work we seek to bridge the concepts of topographic organization and equivariance in neural networks. To accomplish this, we introduce the Topographic VAE: a novel method for efficiently training deep generative models with topographically organized latent variables. We show that such a model indeed learns to organize its activations according to salient characteristics such as digit class, width, and style on MNIST. Furthermore, through topographic organization over time (i.e. temporal coherence), we demonstrate how predefined latent space transformation operators can be encouraged for observed transformed input sequences -- a primitive form of unsupervised learned equivariance. We demonstrate that this model successfully learns sets of approximately equivariant features (i.e. "capsules") directly from sequences and achieves higher likelihood on correspondingly transforming test sequences. Equivariance is verified quantitatively by measuring the approximate commutativity of the inference network and the sequence transformations. Finally, we demonstrate approximate equivariance to complex transformations, expanding upon the capabilities of existing group equivariant neural networks.

Motivation & Objective

• To bridge topographic organization and equivariance in deep generative models by introducing a novel inductive bias.
• To enable unsupervised learning of approximately equivariant capsule representations from sequential data without explicit supervision or hard-coded symmetry.
• To demonstrate that temporal coherence in sequences can induce topographic organization and equivariance in latent space.
• To verify that the learned representations maintain relative pose structure under transformations, enabling generalization to unseen transformation combinations.
• To extend the capabilities of group equivariant networks by learning complex, non-linear transformations such as combined color and rotation or perspective shifts.

Proposed method

• The model uses a Variational Autoencoder (VAE) with a topographically organized latent space where feature activations are spatially structured according to input transformations.
• It enforces temporal coherence by modeling sequences of transformed inputs as time-series, encouraging smooth, slow changes in latent representations across time steps.
• The key innovation is the use of a 'Roll' operation within capsule dimensions to simulate transformation operators (e.g., rotation), allowing the model to decode sequences by rolling activations across capsules.
• The model learns a shared transformation operator via a differentiable routing mechanism that aligns latent activations with observed input transformations.
• The inference network is trained to produce latent codes that approximately commute with transformation operators, quantifying equivariance via a commutativity loss.
• The architecture supports both circular and causal convolutions, enabling application to online or streaming settings with temporal dependencies.

Experimental results

Research questions

• RQ1 Can temporal coherence in sequential data induce topographic organization and equivariance in a deep generative model's latent space?
• RQ2 Can a topographic VAE learn approximately equivariant capsule representations directly from unlabeled sequences without explicit supervision?
• RQ3 Does the model generalize to unseen combinations of transformations, such as combined color and rotation, without architectural changes?
• RQ4 To what extent can the model achieve quantitatively verified equivariance through approximate commutativity of inference and transformation operators?
• RQ5 Can the model learn complex, non-linear transformations (e.g., perspective shifts) in a completely unsupervised manner?

Key findings

• The TVAE successfully learns topographically organized latent representations that reflect salient characteristics such as digit class, width, and style on MNIST.
• The model achieves higher likelihood on test sequences with complex transformations, including combined color and rotation, compared to baseline VAEs.
• Quantitative evaluation confirms approximate equivariance through a commutativity loss, demonstrating that the inference network and transformation operators nearly commute.
• The model generalizes to unseen transformation combinations, such as color and perspective shifts, without architectural changes, indicating robustness and disentanglement.
• When $ K $ equals the capsule size, the model learns invariant representations, as rolling produces no perceptible change in the observation space.
• Capsule traversals show that the model can reconstruct full transformation sequences by rolling initial activations, with accurate reconstructions even for complex transformations like combined color and rotation.