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

TL;DR

This paper introduces Mo3ENet, an end-to-end equivariant autoencoder for multi-type point clouds that provides a complete, rotatable molecular representation useful for various 3D molecular tasks.

Contribution

The paper presents Mo3ENet, a novel equivariant autoencoder with a Gaussian mixture loss for comprehensive molecular point cloud embedding.

Findings

Mo3ENet achieves rotatable, complete molecular representations.

The latent space supports scalar and vector property predictions.

Demonstrated effectiveness on multiple molecular tasks.

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.