Scaling Molecular Representation Learning with Hierarchical Mixture-of-Experts

Abstract

Recent advancements in large-scale self-supervised pretraining have significantly improved molecular representation learning, yet challenges persist, particularly when addressing distributional shifts (e.g., under scaffold-split). Drawing inspiration from the success of Mixture-of-Experts (MoE) networks in NLP, we introduce H-MoE, a hierarchical MoE model tailored for molecular representation learning. Since conventional routing strategies struggle to capture global molecular information--such as scaffold structures, which are crucial for enhancing generalization--we propose a hierarchical routing mechanism. This mechanism first utilizes scaffold-level structural guidance before refining molecular characteristics at the atomic level. To optimize expert assignment, we incorporate scaffold routing contrastive loss, ensuring scaffold-consistent routing while preserving discriminability across molecular categories. Furthermore, a curriculum learning approach and dynamic expert allocation strategy are employed to enhance adaptability. Extensive experiments on molecular property prediction tasks demonstrate the effectiveness of our method in capturing molecular diversity and improving generalization across different tasks.