PaperlayerView exact revision on arXiv

Source-linked AI summary

UniMatch: Universal Matching from Atom to Task for Few-Shot Drug Discovery

Ruifeng Li, Mingqian Li, Wei Liu, Yuhua Zhou, Xiangxin Zhou, Yuan Yao, Qiang Zhang, Hongyang Chen

arXiv:2502.12453v1cs.LGcs.AIq-bio.BM

TL;DR

Scarce experimental annotations and methods that overlook hierarchical molecular structures motivate few-shot drug discovery models that capture multiple molecular scales. UniMatch combines hierarchical molecular matching with task-level meta-learning. It outperforms state-of-the-art methods on MoleculeNet and FS-Mol and demonstrates generalization on Meta-MolNet.

  • Problem

    Few-shot drug discovery faces scarce experimental annotations, and existing methods often overlook molecular structures across atom, substructure, and whole-molecule scales.

    Source: §1 · pp. 1–2

  • Method

    UniMatch combines explicit hierarchical molecular matching across molecular levels with meta-learning for implicit task-level matching.

    Source: §1 · p. 2

  • Results

    UniMatch outperforms state-of-the-art methods on MoleculeNet and FS-Mol, improving AUROC by 2.87% and ∆AUPRC by 6.52%, respectively, and demonstrates generalization on Meta-MolNet.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    UniMatch combines multi-scale structural features and inter-task relationships for few-shot molecular property prediction.

    Source: §6 · p. 10, H DISCUSSION, LIMITATION AND FUTURE WORK · p. 26

    UniMatch underfits regression tasks when its final fusion module uses a linear layer; weighted-average fusion performs better and converges properly.

    Source: H DISCUSSION, LIMITATION AND FUTURE WORK · p. 26

Abstract

from arXiv · show

Drug discovery is crucial for identifying candidate drugs for various diseases.However, its low success rate often results in a scarcity of annotations, posing a few-shot learning problem. Existing methods primarily focus on single-scale features, overlooking the hierarchical molecular structures that determine different molecular properties. To address these issues, we introduce Universal Matching Networks (UniMatch), a dual matching framework that integrates explicit hierarchical molecular matching with implicit task-level matching via meta-learning, bridging multi-level molecular representations and task-level generalization. Specifically, our approach explicitly captures structural features across multiple levels, such as atoms, substructures, and molecules, via hierarchical pooling and matching, facilitating precise molecular representation and comparison. Additionally, we employ a meta-learning strategy for implicit task-level matching, allowing the model to capture shared patterns across tasks and quickly adapt to new ones. This unified matching framework ensures effective molecular alignment while leveraging shared meta-knowledge for fast adaptation. Our experimental results demonstrate that UniMatch outperforms state-of-the-art methods on the MoleculeNet and FS-Mol benchmarks, achieving improvements of 2.87% in AUROC and 6.52% in delta AUPRC. UniMatch also shows excellent generalization ability on the Meta-MolNet benchmark.

1 INTRODUCTION

Drug discovery faces scarce experimental annotations, while existing approaches can miss molecular information at one or more structural scales. UniMatch combines hierarchical molecular matching with task-level meta-learning and reports improvements on MoleculeNet and FS-Mol.

1 INTRODUCTION

  • Limited availability of labeled experimental data is linked to the lengthy duration, high cost, and low success rate of chemical wet experiments.The paper describes these constraints as limiting experimental annotations for molecular prediction.Source: §1 · p. 1
  • Existing methods may miss distinct structural scales: graph methods can over-smooth substructures, while fingerprint methods use fragmented local features and may overlook overall molecular structure.The paper argues that capturing structural levels from atoms through whole molecules matters because different properties depend on different levels.Source: §1 · p. 2, Figure 1 · p. 2
  • UniMatch combines explicit hierarchical molecular matching with implicit task-level matching, aligning molecular structures with tasks.The paper describes the two matching mechanisms as complementary and says their combination enhances adaptability and generalization across tasks.Source: §1 · p. 2
  • Attention-based matching integrates atom-to-higher-level structural information and selects relevant features across multiple molecular levels.The mechanism is presented as a way to capture complex molecular features for prediction.Source: §1 · p. 2
  • Meta-learning provides implicit task-level matching by learning shared parameters, capturing task similarities, and supporting rapid adaptation and improved generalization.The task similarities are captured through optimization.Source: §1 · pp. 2–3
  • UniMatch outperforms state-of-the-art methods on MoleculeNet and FS-Mol, improving AUROC by 2.87% and ∆AUPRC by 6.52%, respectively, and shows outstanding generalization on Meta-Mol.The reported generalization evaluation is on the Meta-Mol benchmark.Source: §1 · p. 3

2 BACKGROUND

The background defines few-shot molecular property prediction through training tasks with small labeled support sets and query sets for evaluation. It also introduces GNNs as graph-representation methods that update node features using neighboring nodes.

2.1 PROBLEM DEFINITION

  • Few-shot molecular property prediction trains across sampled tasks to generalize to new tasks, with each task containing a support set and a query set.The support set provides a few labeled examples for task adaptation; the query set evaluates unseen examples.Source: §2.1 · p. 3

2.2 PRELIMINARIES

  • In a graph, nodes and edges represent graph elements, and GNNs learn representations by aggregating information from neighboring nodes.Node features are updated at each layer using an aggregation step followed by an update step.Source: §2.2 · p. 3

3 METHOD

UniMatch combines explicit hierarchical molecular matching with meta-learning for implicit task-level matching. Its method pools and aligns molecular representations at multiple scales, then uses task relationships to support adaptation across tasks.

3.1 ARCHITECTURE: EXPLICIT HIERARCHICAL MOLECULAR MATCHING

  • Mean pooling at each GNN layer builds representations from nodes toward substructures and whole molecules, enabling comparisons across molecular scales.The method uses GIN as its graph encoder.Source: §3.1 · pp. 3–4
  • Attention matches query representations against support representations and labels at each level, and concatenated layerwise results are fused into final predictions.The architecture integrates the predictions from different GNN layers through a linear layer.Source: §3.1 · p. 4

3.2 META-LEARNING

  • Meta-learning provides implicit task-level matching by learning task relationships that support adaptation and generalization across tasks.The method frames task relationships as a higher-level abstraction.Source: §3.2 · p. 5

3.2.1 TRAINING AND INFERENCE

  • Training adapts task-specific matching and fusion parameters in an inner loop, then updates shared meta-parameters across tasks; inference fine-tunes task-specific parameters before prediction.The outer loop minimizes aggregated loss across training tasks.Source: §3.2.1 · p. 5

3.2.2 IMPLICIT TASK-LEVEL MATCHING

  • A task relationship matrix quantifies relationships between training tasks; implicit matching uses these relationships to shape task-specific parameters and guide updates to the shared encoder.The matrix is defined using a relation function based on the meta-parameters.Source: §3.2.2 · pp. 5–6

4 EXPERIMENT

The experiments evaluate UniMatch on MoleculeNet and FS-Mol, test generalization on Meta-MolNet, and visualize hierarchical representations for toxicity prediction.

4 EXPERIMENT

  • The evaluation spans MoleculeNet and FS-Mol benchmarks, seven Meta-MolNet datasets in single-task and multi-task scenarios, and visualization experiments; experiments ran on an NVIDIA RTX A6000 GPU.Further experiments and analyses are reported in Appendix E.Source: §4 · p. 6

4.1 FEW-SHOT MOLECULAR PROPERTY PREDICTION ON THE MOLECULENET BENCHMARK

  • 2.87% average improvement on MoleculeNet; UniMatch and Pre-UniMatch outperform state-of-the-art methods on Tox21, SIDER (pre-training stage only), and ToxCast, while UniMatch ranks second on MUV.Table 1 reports mean AUROC% with standard deviations for methods evaluated with a support set of 20.Source: Table 1 · p. 6, §4.1 · p. 7

4.2 FEW-SHOT MOLECULAR PROPERTY PREDICTION ON THE FS-MOL BENCHMARK

  • 4.27%, 8.53%, 4.40%, 6.07%, and 4.26% gains at support-set sizes 16, 32, 64, 128, and 256, respectively, place UniMatch ahead of all compared methods on FS-Mol.The evaluation averages ten random support/query splits for each support-set size.Source: Figure 3 · p. 7, §4.2 · p. 8
  • Ablations report advantages from dual matching across GIN, GCN, GAT, and Graph-SAGE backbones, indicating transferability across these GNN architectures.The ablation examines hierarchical pooling and matching alongside meta-learning.Source: Figure 3 · p. 7, §4.2 · p. 8

4.3 CROSS-DOMAIN DRUG DISCOVERY ON THE META-MOLNET BENCHMARK

  • On Meta-MolNet, UniMatch performs exceptionally well on GSK3, JNK3, Tox21, and ToxCast, less well on HIV and PCBA, and faces significant challenges on MUV.Results use a support set of size 2 across seven classification datasets.Source: Figure 4 · p. 8, §4.3 · p. 9

4.4 VISUALIZATION

  • The visualization examines PCA projections of 10 active and inactive molecules and layer-wise attention shifts for one molecule in NR-AhR toxicity prediction.The selected molecule illustrates layers capturing atoms, substructures, and the whole molecule.Source: §4.4 · p. 9, Figure 5 · p. 9

5 RELATED WORK

Prior molecular prediction methods use graph encoders or fingerprints, while matching methods face challenges representing molecular graphs’ hierarchical and global structure.

5.1 GRAPH-BASED MOLECULAR PROPERTY PREDICTION

  • Graph-based few-shot methods commonly use GNN encoders, but typically focus on single-scale features and overlook molecular hierarchy; LLM-based approaches can incur high computational costs.UniMatch instead uses hierarchical pooling and matching to align complex molecular structures.Source: §5.1 · p. 10

5.2 MATCHING LEARNING

  • Matching methods successful in NLP and computer vision struggle with molecular graphs; prior hierarchical methods can miss global structure or generalize poorly across varied molecular graphs.UniMatch combines attention-based matching at atomic, substructural, and molecular levels with task-level meta-learning.Source: §5.2 · p. 10

6 CONCLUSION

UniMatch combines hierarchical molecular matching with implicit task-level matching, while its meta-training procedure updates shared and task-specific parameters across sampled tasks. The appendix also describes task relationships and the datasets used in MoleculeNet.

6 CONCLUSION

  • UniMatch improves AUROC by 2.87% on MoleculeNet and ∆AUPRC by 6.52% on FS-Mol, and demonstrates excellent generalization on Meta-MolNet.The conclusion identifies future work on improving fusion with attention fusion or multi-scale feature aggregation.Source: §6 · p. 10

A.1 ALGORITHM

  • Meta-training samples task batches, computes node and molecular representations at each GNN layer, and updates meta-parameters using query-set performance.The procedure alternates inner task-specific fine-tuning with outer updates to shared parameters.Source: §A.1 · p. 17

A.2 DETAILS OF IMPLICIT TASK-LEVEL MATCHING

  • Task vectors define a relationship matrix used to share information during inner- and outer-loop optimization; new-task parameters can be generated from training-task matches without explicit gradient updates.The relationship matrix can be based on task-vector similarity or distance, and the resulting parameters are used for direct inference.Source: §A.2 · p. 18

B DETAILS OF MOLECULENET BENCHMARK

B.1 DETAILS OF DATASETS

  • The datasets span toxicity, drug side effects, biological activities, and toxicology, with ToxCast comprising 8,615 compounds and 617 tasks.Tox21 has 8,014 compounds and 12 tasks; SIDER has 1,427 compounds and 27 categories; MUV has 93,127 compounds and 17 tasks.Source: §B.1 · p. 19

B.2 DETAILS OF EXPERIMENTAL SETUP

The experimental setup describes UniMatch’s implementation and optimization settings. The following FS-Mol section covers benchmark details, baselines, the ∆AUPRC metric, and setup.

B.2 DETAILS OF EXPERIMENTAL SETUP

  • HieMatch and Pre-HirMatch use a GIN with 5 layers and hidden size 300; UniMatch trains for a maximum of 5000 epochs using Adam with a 0.001 meta-learning rate.The matching and fusion modules use a 0.05 fine-tuning learning rate; dropout is 0.1 for all components except the graph-based molecular encoder.Source: §B.2 · p. 19
  • Table 3 is identified as the location of UniMatch’s hyperparameters.The supplied passages describe the table’s topic but do not provide its entries.Source: Table 3 · p. 19

C DETAILS OF FS-MOL BENCHMARK

  • The FS-Mol appendix section outlines benchmark, baseline, ∆AUPRC metric, and experimental-setup details.The supplied section overview does not provide the actual baseline comparisons or metric calculation.Source: C DETAILS OF FS-MOL BENCHMARK · p. 19

C.1 DETAILS OF BENCHMARKS

FS-Mol comprises many molecular assays, each with relatively few compounds, and is evaluated against several families of baselines. The appendix also describes ∆AUPRC and the benchmark’s training setup.

C.1 DETAILS OF BENCHMARKS

  • FS-Mol contains 5,120 assays and 233,786 compounds, with 157 test tasks and an average of 94 compounds per task.The dataset is partitioned into 4,938 training tasks, 40 validation tasks, and 157 test tasks.Source: §C.1 · p. 19

C.2 DETAILS OF BASELINES.

  • FS-Mol comparisons include single-task, multi-task pre-training, self-supervised pre-training, and meta-learning baselines.Examples described include GP-ST, GNN-ST, CHEF, MixHop, GNN-MT, MAT, PAR, ProtoNet, GNN-MAML, and ADKF-IFT.Source: §C.2 · p. 20

C.3 EVALUATION METRICS OF FS-MOL BENCHMARK

  • The appendix uses ∆AUPRC to compare baseline models, following Stanley et al. (2021).The supplied passages do not include the metric’s calculation formula.Source: §C.3 · p. 21

C.4 DETAILS OF EXPERIMENTAL SETUP

  • For FS-Mol, UniMatch is trained for 10,000 epochs with batch task size 21 and weight decay 5e-5.Source: §C.4 · p. 21

D DETAILS OF META-MOLNET BENCHMARK

D.1 DETAILS OF BENCHMARKS

On FS-Mol, the reported HieMatch comparison covers 157 test tasks across support-set sizes. The appendix describes Meta-MolNet’s task range, comparison baselines, evaluation metrics, and setup.

D.1 DETAILS OF BENCHMARKS

  • The Meta-MolNet evaluation uses seven classification tasks: GSK3, JNK3, HIV, Tox21, ToxCast, PCBA, and MUV.The benchmark integrates diverse datasets through multitask and transfer learning.Source: §D.1 · p. 21

D.2 DETAILS OF BASELINES

  • Meta-MolNet comparisons include classical machine learning, graph-based, message-passing neural-network, and self-supervised pre-training methods.The described candidates include SVM, XGBoost, RF, GCN, DMPNN, CMPNN, Attentive FP, TrimNet, CDDD, Mol2Context-vec, MolBERT, N-gram, and Pre-GNN.Source: §D.2 · pp. 21–22

D.3 EVALUATION METRICS OF META-MOLNET

  • Meta-MolNet classification is evaluated with AUROC and AUPRC, with AUPRC identified as particularly suitable for severely imbalanced PCBA and MUV tasks.Source: §D.3 · p. 23

D.4 DETAILS OF EXPERIMENTAL SETUP

  • The Meta-MolNet setup uses query sets of size 8 and support sets of size 2, with 100 training epochs.AdamW is used with a 0.001 meta-learning rate and a 0.001 inner learning rate.Source: §D.4 · p. 23

E.1 OVERALL PERFORMANCE

  • HieMatch demonstrates superior performance to the state-of-the-art method across all metrics on 157 FS-Mol test tasks and varying support-set sizes.Figure 6 presents the results as box plots of classification-accuracy distributions.Source: §E.1 · p. 23

E.2 SUB-BENCHMARK PERFORMANCE

Across FS-Mol sub-benchmarks, HieMatch leads in half of classification sub-benchmarks at support size 16. UniMatch takes slightly longer to infer than ProtoNet and GNN-MAML, but remains relatively fast.

E.2 SUB-BENCHMARK PERFORMANCE

  • At support size 16, HieMatch is the top performer in half of FS-Mol classification sub-benchmarks.Table E.2 reports results across seven EC-number sub-benchmarks.Source: §E.2 · p. 23, Table 5 · p. 23
  • Figure 6 compares methods across 157 FS-Mol test tasks and varying support set sizes.The caption identifies the box plots' scope but does not state their outcomes.Source: Figure 6 · p. 24

E.3 META-TESTING COSTS

  • UniMatch inference takes slightly longer than ProtoNet and GNN-MAML, while ADKF-IFT has the longest reference time.The authors describe UniMatch's inference time as relatively fast and viable for meta-learning tasks.Source: §E.3 · p. 25, Figure 7 · p. 25

F VISUALIZATION EXPERIMENTS

The visualization section identifies ten Tox21 molecules for NR-AhR toxicity prediction and their activity labels. The related-work discussion describes hierarchical graph methods and their stated limitations.

F VISUALIZATION EXPERIMENTS

  • The visualization experiment uses ten molecules sampled from Tox21 for NR-AhR toxicity prediction; labels 1 and 0 indicate active and inactive molecules.Table 6 lists the molecules used in Section 4.4.Source: Table 6 · p. 25, F VISUALIZATION EXPERIMENTS · p. 25

G.1 HIERARCHICAL REPRESENTATION LEARNING ON GRAPHS

  • Hierarchical graph representations can capture multi-scale structure and help models discern local and global patterns.The passage frames this as motivation for hierarchical representation learning on graphs.Source: §G.1 · p. 25
  • Layer-wise methods struggle to balance local and global structures, while architecture-level methods may lack task-specific adaptation and generalization.The discussion names DiffPool, JK-Nets, Top-K Pooling, ASAP, HGP-SL, MixHop, FraGAT, MGSSL, and UniCorn as examples.Source: §G.1 · p. 26
  • Supergraph methods are described as computationally complex and insufficiently sensitive to structural diversity; HiPM still requires substantial data.The passage names HiMol as a supergraph method and HiPM as a prompt-tree approach.Source: §G.1 · p. 26

H DISCUSSION, LIMITATION AND FUTURE WORK

UniMatch outperforms state-of-the-art methods on MoleculeNet and FS-Mol and shows strong generalization on Meta-MolNet. The discussion identifies fusion design, particularly for regression, as a limitation and proposes broader improvements.

H DISCUSSION, LIMITATION AND FUTURE WORK

  • UniMatch outperforms state-of-the-art methods on MoleculeNet and FS-Mol and shows exceptional generalization on Meta-MolNet.The conclusion reports significant improvements in AUROC and ∆AUPRC without specifying values in this passage.Source: H DISCUSSION, LIMITATION AND FUTURE WORK · p. 26
  • The fusion mechanism is relatively simplistic and may limit integration of information across hierarchical levels and capture of cross-scale dependencies.The authors suggest attention-based fusion or multi-scale feature aggregation, while noting these may increase computational complexity.Source: H DISCUSSION, LIMITATION AND FUTURE WORK · p. 26
  • On regression tasks, the linear final fusion layer is associated with underfitting; weighted-average fusion instead achieves better performance and proper convergence.The authors attribute the issue to the linear layer's ability to capture relationships among features from different layers.Source: H DISCUSSION, LIMITATION AND FUTURE WORK · p. 26
  • Future work includes testing advanced fusion strategies, broader datasets and tasks, and domain-specific knowledge and features.The authors also plan to investigate interpretability and improve computational efficiency and scalability.Source: H DISCUSSION, LIMITATION AND FUTURE WORK · pp. 26–27

Related papers

Loading 2502.12453v1…