Contextual Representation Anchor Network to Alleviate Selection Bias in Few-Shot Drug Discovery

TL;DR

A novel method named contextual representation anchor Network (CRA), where an anchor refers to a cluster center of the representations of molecules and serves as a bridge to transfer enriched contextual knowledge into molecular representations and enhance their expressiveness, is presented.

Abstract

In the drug discovery process, the low success rate of drug candidate screening often leads to insufficient labeled data, causing the few-shot learning problem in molecular property prediction. Existing methods for few-shot molecular property prediction overlook the sample selection bias, which arises from non-random sample selection in chemical experiments. This bias in data representativeness leads to suboptimal performance. To overcome this challenge, we present a novel method named contextual representation anchor Network (CRA), where an anchor refers to a cluster center of the representations of molecules and serves as a bridge to transfer enriched contextual knowledge into molecular representations and enhance their expressiveness. CRA introduces a dual-augmentation mechanism that includes context augmentation, which dynamically retrieves analogous unlabeled molecules and captures their task-specific contextual knowledge to enhance the anchors, and anchor augmentation, which leverages the anchors to augment the molecular representations. We evaluate our approach on the MoleculeNet and FS-Mol benchmarks, as well as in domain transfer experiments. The results demonstrate that CRA outperforms the state-of-the-art by 2.60% and 3.28% in AUC and $Δ$AUC-PR metrics, respectively, and exhibits superior generalization capabilities.

Figures (7)

• Figure 1: The sample selection bias in the NR-AR and NR-AhR tasks from the Tox21 dataset. The top row shows the distribution of positive and negative samples, highlighting the imbalance in data sample. The bottom row displays 2D PCA projections of the molecular MACCS fingerprints, illustrating the clustering patterns influenced by sample selection bias.
• Figure 2: The augmentation of molecular embeddings and Contextual representation anchors in CRA. The left part shows the initial molecular embeddings and anchors. The middle part shows the change of anchors after augment with reference data. The right part shows the change of molecular embeddings by using Contextual representation anchors for augmentation.
• Figure 3: Overview of CRA. The architecture comprises three main components: the context augmentation module (CAM), the anchor augmentation module (AAM), and the matching module (MM). The procedure begins with a shared encoder to obtain molecular embeddings and initial anchors. Next, The CAM augments these initial anchors with unlabeled reference molecules. Then, the augmented anchors are used to form augmented molecular embeddings, which are enhanced by the AAM. The MM leverages the similarities between the query samples and support samples to derive the final predictions. The attention matrix highlights the similarity between molecular structures (e.g., a similarity score of 0.63).
• Figure 4: The mean performance with standard errors of the ablation study on FS-Mol benchmark. (a) The results (AUC and $\Delta$AUC-PR) of the ablation study for the components of CRA. (b) The performance of CRA with varying size of reference sets.
• Figure 5: Performances of CRA and MHNfs for various support set sizes during inference.