# Learning to Group Auxiliary Datasets for Molecule

arXiv: https://arxiv.org/abs/2307.04052 · HTML: https://mlsanity.com/paper/2307.04052/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 65 references that are in the map.
- First posted on arXiv 8 Jul 2023.

## Builds on (most cited first, 30 of 43)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [Generalized Out-of-Distribution Detection: A Survey](https://mlsanity.com/paper/2110.11334/) (2021)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [GCC: Graph Contrastive Coding for Graph Neural Network Pre-Training](https://mlsanity.com/paper/2006.09963/) (2020)
- [GraphMAE: Self-Supervised Masked Graph Autoencoders](https://mlsanity.com/paper/2205.10803/) (2022)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Tutorial on Variational Autoencoders](https://mlsanity.com/paper/1606.05908/) (2016)
- [Analyzing Learned Molecular Representations for Property Prediction](https://mlsanity.com/paper/1904.01561/) (2019)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [Meta Pseudo Labels](https://mlsanity.com/paper/2003.10580/) (2020)
- [OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs](https://mlsanity.com/paper/2103.09430/) (2021)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [Order Matters: Sequence to sequence for sets](https://mlsanity.com/paper/1511.06391/) (2015)
- [Beam Search Strategies for Neural Machine Translation](https://mlsanity.com/paper/1702.01806/) (2017)
- [Pareto Multi-Task Learning](https://mlsanity.com/paper/1912.12854/) (2019)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)

## Built on by (most cited first, 3 of 3)

- [Semi-overlapping Multi-bandit Best Arm Identification for Sequential Support Network Learning](https://mlsanity.com/paper/2512.24959/) (2025)
- [Diversity Curves for Graph Representation Learning](https://mlsanity.com/paper/2605.06466/) (2026)
- [CalibHyper: Chance-Corrected Relational Hypergraphs for Few-Shot Molecular Property Prediction](https://mlsanity.com/paper/2609.33342/) (2026)

## Abstract (arXiv)

The limited availability of annotations in small molecule datasets presents a challenge to machine learning models. To address this, one common strategy is to collaborate with additional auxiliary datasets. However, having more data does not always guarantee improvements. Negative transfer can occur when the knowledge in the target dataset differs or contradicts that of the auxiliary molecule datasets. In light of this, identifying the auxiliary molecule datasets that can benefit the target dataset when jointly trained remains a critical and unresolved problem. Through an empirical analysis, we observe that combining graph structure similarity and task similarity can serve as a more reliable indicator for identifying high-affinity auxiliary datasets. Motivated by this insight, we propose MolGroup, which separates the dataset affinity into task and structure affinity to predict the potential benefits of each auxiliary molecule dataset. MolGroup achieves this by utilizing a routing mechanism optimized through a bi-level optimization framework. Empowered by the meta gradient, the routing mechanism is optimized toward maximizing the target dataset's performance and quantifies the affinity as the gating score. As a result, MolGroup is capable of predicting the optimal combination of auxiliary datasets for each target dataset. Our extensive experiments demonstrate the efficiency and effectiveness of MolGroup, showing an average improvement of 4.41%/3.47% for GIN/Graphormer trained with the group of molecule datasets selected by MolGroup on 11 target molecule datasets.