[Paper Review] Semi-supervised learning of images with strong rotational disorder: assembling nanoparticle libraries

Maxim Ziatdinov, Muammer Y. Yaman+3 more|arXiv (Cornell University)|May 24, 2021

Genetics, Bioinformatics, and Biomedical Research4 citations

TL;DR

This paper proposes a semi-supervised rotationally invariant variational autoencoder (ss-rVAE) that enables classification of images with strong rotational and translational disorder by leveraging a small amount of labeled data to generalize to a large unlabeled dataset. The model learns a rotationally invariant latent space and transfers this invariance to a downstream classifier, enabling robust nanoparticle library construction and disentanglement of physical factors of variation in experimental microscopy data.

ABSTRACT

The proliferation of optical, electron, and scanning probe microscopies gives rise to large volumes of imaging data of objects as diversified as cells, bacteria, pollen, to nanoparticles and atoms and molecules. In most cases, the experimental data streams contain images having arbitrary rotations and translations within the image. At the same time, for many cases, small amounts of labeled data are available in the form of prior published results, image collections, and catalogs, or even theoretical models. Here we develop an approach that allows generalizing from a small subset of labeled data with a weak orientational disorder to a large unlabeled dataset with a much stronger orientational (and positional) disorder, i.e., it performs a classification of image data given a small number of examples even in the presence of a distribution shift between the labeled and unlabeled parts. This approach is based on the semi-supervised rotationally invariant variational autoencoder (ss-rVAE) model consisting of the encoder-decoder "block" that learns a rotationally (and translationally) invariant continuous latent representation of data and a classifier that encodes data into a finite number of discrete classes. The classifier part of the trained ss-rVAE inherits the rotational (and translational) invariances and can be deployed independently of the other parts of the model. The performance of the ss-rVAE is illustrated using the synthetic data sets with known factors of variation. We further demonstrate its application for experimental data sets of nanoparticles, creating nanoparticle libraries and disentangling the representations defining the physical factors of variation in the data. The code reproducing the results is available at https://github.com/ziatdinovmax/Semi-Supervised-VAE-nanoparticles.

Motivation & Objective

• To address the challenge of classifying images with strong rotational and translational disorder when only a small amount of labeled data is available.
• To develop a method that generalizes from labeled data with weak disorder to unlabeled data with strong distribution shift in orientation and position.
• To create a scalable, interpretable nanoparticle library from experimental microscopy data using disentangled representations.
• To enable transfer of rotational and translational invariance from the autoencoder to a downstream classifier for deployment on new data.

Proposed method

• The ss-rVAE model combines an encoder-decoder architecture with rotational and translational invariance in the latent space through data augmentation and equivariant design.
• The model uses a variational autoencoder (VAE) framework with a latent space that is invariant to rotations and translations via augmentation and group-equivariant constraints.
• A classifier head is trained on the invariant latent representations, inheriting the rotational and translational invariance from the encoder.
• The method leverages both labeled data for supervised learning and unlabeled data for self-supervised pre-training, enabling semi-supervised generalization.
• The model is trained end-to-end using a combination of reconstruction loss, KL divergence, and cross-entropy loss on labeled data.
• The framework is applied to synthetic data with known factors of variation and then to real experimental nanoparticle datasets from microscopy.

Experimental results

Research questions

• RQ1 Can a semi-supervised model generalize from a small set of labeled images with weak disorder to a large set of unlabeled images with strong rotational and translational disorder?
• RQ2 To what extent can the learned latent space disentangle physical factors of variation such as particle shape and size in nanoparticle imaging data?
• RQ3 How well does the ss-rVAE transfer rotational and translational invariance from the autoencoder to a downstream classifier?
• RQ4 Can the model create a reliable, interpretable nanoparticle library from experimental microscopy data with minimal human annotation?

Key findings

• The ss-rVAE achieves high classification accuracy on synthetic datasets with known factors of variation, demonstrating effective disentanglement of rotation and translation from other factors.
• The model successfully generalizes from labeled data with weak disorder to unlabeled data with strong disorder, achieving robust performance despite distribution shift.
• The trained classifier, which inherits invariance from the encoder, achieves high accuracy on independent test sets, confirming transferability of learned invariances.
• Application to experimental nanoparticle datasets enables the creation of a structured, interpretable nanoparticle library from raw microscopy images.
• The disentangled representations reveal physical factors of variation such as particle morphology and size, enabling meaningful clustering and analysis.