Adapting Differential Molecular Representation with Hierarchical Prompts for Multi-label Property Prediction

TL;DR

Multi-label molecular property prediction faces an exponential output space $2^m$ and gradient conflicts among tasks. The authors introduce HiPM, a hierarchical prompted molecular representation learning framework with a Molecular Representation Encoder (MRE) and a Task-Aware Prompter (TAP) to model multi-granular task correlations. HiPM achieves state-of-the-art performance on six MoleculeNet datasets, particularly when label correlations are strong, and provides interpretability through affinity-guided task clustering and motif-weight analysis. This work mitigates negative transfer in multi-label settings and offers a scalable tool for drug discovery applications.

Abstract

Accurate prediction of molecular properties is crucial in drug discovery. Traditional methods often overlook that real-world molecules typically exhibit multiple property labels with complex correlations. To this end, we propose a novel framework, HiPM, which stands for hierarchical prompted molecular representation learning framework. HiPM leverages task-aware prompts to enhance the differential expression of tasks in molecular representations and mitigate negative transfer caused by conflicts in individual task information. Our framework comprises two core components: the Molecular Representation Encoder (MRE) and the Task-Aware Prompter (TAP). MRE employs a hierarchical message-passing network architecture to capture molecular features at both the atom and motif levels. Meanwhile, TAP utilizes agglomerative hierarchical clustering algorithm to construct a prompt tree that reflects task affinity and distinctiveness, enabling the model to consider multi-granular correlation information among tasks, thereby effectively handling the complexity of multi-label property prediction. Extensive experiments demonstrate that HiPM achieves state-of-the-art performance across various multi-label datasets, offering a novel perspective on multi-label molecular representation learning.

Figures (5)

• Figure 1: The salicylic acid structures in aspirin, salsalate, and sodium salicylate are specially emphasized by dashed circles, which collectively imply that these compounds are anti-inflammatory, acidic, and hydrophilic.
• Figure 2: Overview of the HiPM Framework: (A) illustrates the architecture of our model, where the prompt matrix is directly fused with the molecular representations generated using Eq \ref{['eqIns2']}. (B) details the process of calculating task affinities using cosine similarities. (C) describes the methodology for constructing the hierarchical prompt tree structure, utilizing the agglomerative hierarchical clustering algorithm. (D) outlines the process of computing soft prompts for non-leaf nodes in a bottom-up manner. (E) explains how the prompt matrix is derived from the nodes corresponding to the prefix paths of tasks on the prompt tree. We use Eq. \ref{['eqPro']} to perform fusion on all prompts along the path.
• Figure 3: Visualization results of the affinity matrix (A) and hierarchical clustering tree (B) for the 12 tasks of Tox21. In the affinity matrix, darker colors represent higher task affinities. In the hierarchical clustering tree, the leaf nodes are labeled, and potential reasons for the prioritized clustering of certain tasks are provided. Specifically, CS-DM stands for Cellular Stress and Defense Mechanisms, CSR stands for Cellular Stress Response, NR-HM stands for Nuclear Receptor and Hormone Metabolism, NR-HR stands for Nuclear Receptor and Hormone Regulation, and NRA stands for Nuclear Receptor Activity. The clustering of these leaf nodes reflects their functional similarities and associations within their respective toxicological mechanisms.
• Figure 4: Results of ablation experiments. Each variant was run with three random seeds. We report the average ROC-AUC (classification) or MAE (regression) scores along with their standard deviations. Higher ROC-AUC values indicate superior performance, while lower MAE values are preferable. The error bars denote standard deviations.
• Figure 5: Visualization analysis results that highlight key motifs and the changes in the corresponding attention weights before and after integrating the TAP module. Bar charts illustrate the variations in predicted probabilities for each label generated by HiPM before and after the inclusion of TAP. The value of each bar is calculated as the predicted probability (after) minus the predicted probability (before). An increase (or decrease) in predicted probability after adding TAP, where the true label is 1 (or 0), indicates an increase in the model’s confidence in the correct answer.