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

2024-05-22

arXiv.org

AI Summary

To address the challenge of jointly modeling molecular conformations and 3D orientations, this paper introduces Mo3ENet—the first strictly O(3)-equivariant multi-type point cloud autoencoder. Mo3ENet uniformly encodes atomic types, Cartesian coordinates, and global rotations/reflections, achieving exact SE(3)-equivariance for coordinates and SE(3)-invariance for scalars—fully decoupling equivariance from invariance. It innovatively employs a Gaussian mixture density reconstruction loss to enhance 3D structural fidelity. Its latent space natively supports prediction of both scalar properties (e.g., energy) and vectorial properties (e.g., dipole moment), and directly enables controllable 3D molecular generation. On joint conformation–pose reconstruction, Mo3ENet significantly outperforms existing baselines. It establishes the first general-purpose, rigorously group-equivariant representation framework for molecular clusters, dimers, and condensed-phase systems.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: 3D Computer Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

The point cloud is a flexible representation for a wide variety of data types, and is a particularly natural fit for the 3D conformations of molecules. Extant molecule embedding/representation schemes typically focus on internal degrees of freedom, ignoring the global 3D orientation. For tasks that depend on knowledge of both molecular conformation and 3D orientation, such as the generation of molecular dimers, clusters, or condensed phases, we require a representation which is provably complete in the types and positions of atomic nuclei and roto-inversion equivariant with respect to the input point cloud. 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 for downstream learning tasks. 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.

Problem

• Develops a rotatable, complete representation for molecular conformation and 3D orientation
• Proposes an equivariant autoencoder for multi-type point clouds with O(3) symmetry
• Creates a universal embedding for 3D molecular property prediction tasks

Innovation

• Equivariant autoencoder for molecular point clouds
• Gaussian mixture loss for point cloud reconstruction
• O(3)-equivariant latent space for 3D tasks