Published March 8, 2024 | Version v7
Preprint Open

Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets

  • 1. Mila, Valence, University of Montreal
  • 2. Mila, Valence, McGill University
  • 3. Graphcore
  • 4. Mila, Valence, New Jersey Institute of Technology
  • 5. RWTH Aachen University
  • 6. Mila
  • 7. Mila, University of Montreal
  • 8. Valence
  • 9. Mila, University of Montreal, CIFAR AI Chair
  • 10. Mila, McGill University, CIFAR AI Chair
  • 11. Mila, HEC Montreal, CIFAR AI Chair


Recently, pre-trained foundation models have shown significant advancements in multiple fields. However, the lack of datasets with labeled features and codebases has hindered the development of a supervised foundation model for molecular tasks. Here, we have carefully curated seven datasets specifically tailored for node- and graph-level prediction tasks to facilitate supervised learning on molecules. Moreover, to support the development of multi-task learning on our proposed datasets, we created the Graphium graph machine learning library. Our dataset collection encompasses two distinct categories. Firstly, the TOYMIX category modifies three small existing datasets with additional data for multi-task learning. Secondly, the LARGEMIX category includes four large-scale datasets with 344M graph-level data points and 409M node-level data points from ∼5M unique molecules. Finally, the ultra-large dataset contains 2,210M graph-level data points and 2,031M node-level data points coming from 86M molecules. Hence our datasets represent an order of magnitude increase in data volume compared to other 2D-GNN datasets. In addition, recognizing that molecule-related tasks often span multiple levels, we have designed our library to explicitly support multi-tasking, offering a diverse range of multi-level representations, i.e., representations at the graph, node, edge, and node-pair level. We equipped the library with an extensive collection of models and features to cover different levels of molecule analysis. By combining our curated datasets with this versatile library, we aim to accelerate the development of molecule foundation models. Datasets and code are available at


For the ultra large datasets due to its large size, it is uploaded at


Files (2.6 GB)

Name Size Download all
2.1 MB Preview Download
2.7 MB Preview Download
152.8 MB Preview Download
2.5 GB Preview Download
12.8 MB Preview Download
96.5 kB Preview Download
395.3 kB Preview Download