Fig. 1: O2-VAE learns orientation-invariant representations of cells and organelles.

a The O2-VAE model. An input image to the orientation-invariant encoder produces the same output vector for any input orientation. In the box labelled `encoder' each layer of the convolutional encoder is constrained to be orientation-equivariant with a final spatial pooling layer to produce an orientation-invariant learned representation vector of cell phenotype. During training, the representation is decoded using a separate neural-network decoder, which is trained to reconstruct the input. By design, the learned representation is orientation invariant and thus the reconstruction orientation may differ from the input. In the box labelled `re-alignment', we efficiently estimate and correct the misalignment using Fourier transform-based methods. We use a loss function (not shown, for simplicity of the figure) that promotes accurate reconstruction while constraining the distribution of representations. b O2-VAE can be trained on any image including: binary, greyscale, and multi-channel images. The learned representation vector or phenotypic profile can be used for downstream analysis; three representative tasks are shown with each dot corresponding to an image in embedding space. For discretely-varying shapes, objects form separated clusters; for continuously-varying shapes, data can be visualised with dimensionality reduction; outlier data will be far from most other data.