$\texttt{MiniMol}$: A Parameter-Efficient Foundation Model for Molecular Learning

TL;DR

MiniMol presents a parameter-efficient, 10M-parameter graph neural network pre-trained on a diverse LargeMix of ~6M molecules and ~526M labels to produce transferable molecular fingerprints. By leveraging a fingerprinting pipeline with a GINE backbone and graph-level pooling, it achieves state-of-the-art or near state-of-the-art results on the Therapeutic Data Commons ADMET benchmarks, outperforming the previous foundation model MolE on 17 of 22 tasks while using only a fraction of the parameters. The work analyzes how pre-training task selection and backbone choice influence downstream transfer, revealing that certain graph-level tasks can negatively correlate with downstream performance and that careful loss balancing is essential in multi-task pre-training. Overall, MiniMol demonstrates that strong molecular representations can be learned with modest parameter budgets and efficiently adapted to low-data downstream tasks, with an open-source release to foster further research.

Abstract

In biological tasks, data is rarely plentiful as it is generated from hard-to-gather measurements. Therefore, pre-training foundation models on large quantities of available data and then transfer to low-data downstream tasks is a promising direction. However, how to design effective foundation models for molecular learning remains an open question, with existing approaches typically focusing on models with large parameter capacities. In this work, we propose $\texttt{MiniMol}$, a foundational model for molecular learning with 10 million parameters. $\texttt{MiniMol}$ is pre-trained on a mix of roughly 3300 sparsely defined graph- and node-level tasks of both quantum and biological nature. The pre-training dataset includes approximately 6 million molecules and 500 million labels. To demonstrate the generalizability of $\texttt{MiniMol}$ across tasks, we evaluate it on downstream tasks from the Therapeutic Data Commons (TDC) ADMET group showing significant improvements over the prior state-of-the-art foundation model across 17 tasks. $\texttt{MiniMol}$ will be a public and open-sourced model for future research.

Figures (3)

• Figure 1: Workflow overview of the MiniMol pre-training and downstream task evaluation. MiniMol is pre-trained on the LargeMix datasets. Then MiniMol embeddings are used as molecular fingerprints for downstasks such as those in TDC.
• Figure 2: Schematic of the architecture of MiniMol. An example molecule is featurized in the first block. Node feature vectors are created by combining chemical features with positional and structural encodings, edge features are generated using RDKit and a random initial global vector is generated. Each initial vector is processed with a separate embedding MLP. The backbone of the model is a stack of MPNN layers, which output the molecular fingerprint $\psi$ after pooling. The pre-pooling output is used for pre-training on node-level tasks, in our case, PCQM4M_N4. The fingerprint $\psi$ is used either for pre-training multiple graph-level task heads or as an input to the downstream tasks, including the full set of ADMET tasks from the TDC benchmarks.
• Figure 3: Example of the MPNN block architecture given in Eq.\ref{['eq.mpnn']}. The edge update in Eq.\ref{['eq:edge_update']} gathers the nodes and edges before passing through the MLP first, then this output is used for the node update in Eq.\ref{['eq:node_update']}, gathering all connected node features and updated edge features. The global update in Eq.\ref{['eq:global_update']} connects all nodes and edges, before finally the skip connections in Eq.\ref{['eq:skip']}.