Topographic VAEs learn Equivariant Capsules

T. Anderson Keller, Max Welling

Keywords: Topographic Organization, Unsupervised, Equivariance, Variational Inference, Deep Generative Model, Disentanglement

TL;DR: We show how deep topographic generative models can be efficiently trained and subsequently leverged to learn equivariant sets of features from unsupervised sequences.

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.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

Code: https://github.com/akandykeller/TopographicVAE

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2109.01394/code)