Aligning Rotational and Hierarchical Geometry in Molecular Representation Learning with Product-Manifold Latent Spaces

Urvi Awasthi, Alexander Arjun Lobo, Leonid Zhukov

Keywords: molecular representation learning, molecular machine learning, equivariant graph neural networks, rotational symmetry, SO(3)-equivariance, hyperbolic geometry, product manifold, hierarchical chemical scaffolds, geometric deep learning, message passing, manifold Bayesian optimization, scaffold split generalization, molecular property prediction, generative molecular design, curvature-aware aggregation, symmetry-preserving learning

Abstract: Learning effective molecular representations requires capturing two fundamental but largely disjoint aspects of the structure of molecules: rotational symmetries in 3D conformations and the hierarchical organization of chemical scaffolds. We introduce a new paradigm of product-manifold representation learning with product-manifold message passing on $\mathrm{SO}(3) \times \mathbb{H}^d$, which couples equivariant geometric features with hyperbolic embeddings of chemical hierarchy. Our construction preserves $\mathrm{SO}(3)$-equivariance in the geometric channel and uses an $\mathrm{E}(3)$‑invariant readout for scalar properties while enabling curvature-aware aggregation in the hyperbolic channel, with cross-coupling restricted to scalar invariants to maintain symmetry. Unlike prior approaches that fuse equivariant and hierarchical encoders via concatenation or stacking, our method defines message passing directly on the product manifold, yielding a unified representation. We outline how such models could be evaluated on molecular property prediction, scaffold-split generalization, and generative design, and discuss how embeddings in $\mathrm{SO}(3) \times \mathbb{H}^d$ provide a natural surrogate space for manifold Bayesian optimization, enabling more sample-efficient discovery of high-value molecules compared to Euclidean BO. Together, these results suggest a principled path toward unifying physical symmetries and chemical hierarchies within a single geometric learning framework.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 24738