Multi-Type Point Cloud Autoencoder: A Complete Equivariant Embedding for Molecule Conformation and Pose

TL;DR

This work introduces Mo3ENet, a complete O(3)-equivariant autoencoder for multi-type point clouds that encodes full molecular conformation and pose into a fixed-size latent representation. It uses a Gaussian-mixture reconstruction loss and an end-to-end EGNN-based encoder–decoder architecture to achieve lossless-like reconstruction and powerful embeddings. The approach yields strong molecule reconstruction, chemically meaningful embeddings, accurate QM9 property predictions, and valuable crystal-embedding capabilities for lattice energy estimation. Overall, Mo3ENet provides a universal, rotatable embedding suitable for downstream tasks in molecular design, clustering, and crystal structure prediction, with potential as a generative conditioning backbone for inverse problems.

Abstract

Representations are a foundational component of any modelling protocol, including on molecules and molecular solids. For tasks that depend on knowledge of both molecular conformation and 3D orientation, such as the modelling of molecular dimers, clusters, or condensed phases, we desire a rotatable representation that is provably complete in the types and positions of atomic nuclei and roto-inversion equivariant with respect to the input point cloud. In this paper, we develop, train, and evaluate a new type of autoencoder, molecular O(3) encoding net (Mo3ENet), for multi-type point clouds, for which we propose a new reconstruction loss, capitalizing on a Gaussian mixture representation of the input and output point clouds. Mo3ENet is end-to-end equivariant, meaning the learned representation can be manipulated on O(3), a practical bonus. An appropriately trained Mo3ENet latent space comprises a universal embedding for scalar and vector molecule property prediction tasks, as well as other downstream tasks incorporating the 3D molecular pose, and we demonstrate its fitness on several such tasks.

Figures (15)

• Figure 1: (a) Overall summary of autoencoder architecture. Molecule point cloud coordinates, $\vec{r}$ and atom types $z$ are embedded and bottlenecked to dimension $f_{emb}$ by the encoder. Then, the decoder (an equivariant MLP or GNN with fixed graph size) transforms the embedding to $n_j$ independently weighted points in the $3+n_c+1$ dimensional output cloud, where $n_c$ is the number of classes or types in the dataset and one extra dimension for the learned weights. (b) Outline of the architecture of the encoder graph model, outputting the equivariant embedding $\bm{\vec{g}}$ and its scalarized counterpart $\bm{\tilde{g}}$.
• Figure 2: Illustrative diagram demonstrating the correspondence between point clouds and Gaussian Mixtures. We show the explicit GMs projecting in different colors each type dimension as well as cartoons of point-centered Gaussians, projected in the 'carbon' dimension. These GMs are compared pairwise over input and output points to determine the overall overlap with each input node.
• Figure 3: (a), (b), and (c) Convergence of the output Gaussian mixture on a flat molecule (point cloud shown), while reducing the Gaussian width, projecting each of the typewise dimensions in individual colors. (d) Input GM.
• Figure 4: Mean distance distribution for Mo3ENet models on the test QM9 dataset (10k samples). The models omit or include hydrogen atoms, respectively. In panel (a) we show the mean deviations averaged over the whole molecule, for molecules where every atom was matched by our clustering procedure. We also give the overall average and matched molecule fractions. In panel (b) we show the same statistics on a per-atom basis.
• Figure 5: The UMAP decomposition for the full QM9 dataset according to our autoencoder trained with hydrogens, with the x and y axes the UMAP reduced dimensions of the 64 dimensional encoder embedding. Points are color coded according to, (a) compositional (atom fraction) and, (b) geometric molecular (principal inertial ratios) factors, with legends for each on the right hand side.