# Affinity-VAE: incorporating prior knowledge in representation learning from scientific images

arXiv: https://arxiv.org/abs/2209.04517 · HTML: https://mlsanity.com/paper/2209.04517/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 49 references that are in the map.
- First posted on arXiv 9 Sep 2022.

## Builds on (most cited first, 12 of 12)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [Learning SO(3) Equivariant Representations with Spherical CNNs](https://mlsanity.com/paper/1711.06721/) (2017)
- [Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing](https://mlsanity.com/paper/1903.10145/) (2019)
- [Explicitly disentangling image content from translation and rotation with spatial-VAE](https://mlsanity.com/paper/1909.11663/) (2019)
- [Positive-unlabeled convolutional neural networks for particle picking in cryo-electron micrographs](https://mlsanity.com/paper/1803.08207/) (2018)
- [Weakly Supervised Disentanglement by Pairwise Similarities](https://mlsanity.com/paper/1906.01044/) (2019)
- [AtomAI: A Deep Learning Framework for Analysis of Image and Spectroscopy Data in (Scanning) Transmission Electron Microscopy and Beyond](https://mlsanity.com/paper/2105.07485/) (2021)
- [SHREC 2021: Classification in cryo-electron tomograms](https://mlsanity.com/paper/2203.10035/) (2022)
- [A convolutional autoencoder approach for mining features in cellular electron cryo-tomograms and weakly supervised coarse segmentation](https://mlsanity.com/paper/1706.04970/) (2017)

## Built on by (most cited first, 1 of 1)

- [The microscope is the mask: privileged views and labels from a cryo-ET forward model](https://mlsanity.com/paper/2609.04325/) (2026)

## Abstract (arXiv)

Learning compact and interpretable representations of data is a critical challenge in scientific image analysis. Here, we introduce Affinity-VAE, a generative model that enables us to impose our scientific intuition about the similarity of instances in the dataset on the learned representation during training. We demonstrate the utility of the approach in the scientific domain of cryo-electron tomography (cryo-ET) where a significant current challenge is to identify similar molecules within a noisy and low contrast tomographic image volume. This task is distinct from classification in that, at inference time, it is unknown whether an instance is part of the training set or not. We trained affinity-VAE using prior knowledge of protein structure to inform the latent space. Our model is able to create rotationally-invariant, morphologically homogeneous clusters in the latent representation, with improved cluster separation compared to other approaches. It achieves competitive performance on protein classification with the added benefit of disentangling object pose, structural similarity and an interpretable latent representation. In the context of cryo-ET data, affinity-VAE captures the orientation of identified proteins in 3D which can be used as a prior for subsequent scientific experiments. Extracting physical principles from a trained network is of significant importance in scientific imaging where a ground truth training set is not always feasible.