Unified and explainable molecular representation learning for imperfectly annotated data from the hypergraph view

Abstract

Molecular representation learning (MRL) has shown promise in accelerating drug development by predicting chemical properties. However, imperfectly annotation among datasets pose challenges in model design and explainability. In this work, we formulate molecules and corresponding properties as a hypergraph, extracting three key relationships: among properties, molecule-to-property, and among molecules, and developed a unified and explainable multi-task MRL framework, OmniMol. It integrates a task-related meta-information encoder and a task-routed mixture of experts (t-MoE) backbone to capture correlations among properties and produce task-adaptive outputs. To capture underlying physical principles among molecules, we implement an innovative SE(3)-encoder for physical symmetry, applying equilibrium conformation supervision, recursive geometry updates, and scale-invariant message passing to facilitate learning-based conformational relaxation. OmniMol achieves state-of-the-art performance in properties prediction, reaches top performance in chirality-aware tasks, demonstrates explainability for all three relations, and shows effective performance in practical applications. Our code is available in our https://github.com/bowenwang77/OmniMol public repository. AI models for drug discovery often struggle with real-world, incomplete data. Here, the authors present OmniMol, a framework using hypergraphs to improve predictions of molecular properties, addressing challenges of imperfect data annotation and enhancing model explainability.