Molecular Prior Networks

AI summary

Overview

Research area: Machine learning for drug discovery, specifically protein–ligand affinity prediction, few-shot/meta-learning, and Prior-data Fitted Networks (a class of tabular foundation models).

Technical level: Advanced. The paper assumes familiarity with transformer attention, meta-learning, and kernel regression, though its headline results are stated in accessible terms.

Scope: A workshop paper (AI for Drug Discovery – Bridging the Translation Gap) introducing a transformer that predicts molecular affinity labels in a single forward pass from a labelled support set, trained across thousands of existing affinity datasets and evaluated on the regression variant of the FS-Mol benchmark.

What This Paper Is About

Drug-target affinity models are usually trained on very small datasets, ranging from tens to hundreds of molecules, because wet-lab experiments are expensive. This makes it hard to generalize across targets and chemical series, so existing affinity datasets tend to be discarded for new projects. The authors ask whether a transformer pre-trained over thousands of existing affinity datasets can learn general priors about molecular activity and then predict labels for an unseen target's molecules in one forward pass, without any task-specific training.

Key Contributions

1. A Molecular Prior Network (MPN): a transformer that takes a labelled support set and an unlabelled query set of molecules as input and predicts all query labels jointly in a single forward pass, with no per-task fitting.
2. A domain-specific prior instead of a synthetic one: rather than generating synthetic datasets as TabPFN and TabICL do, the model is meta-trained on thousands of existing affinity datasets, each corresponding to a single protein target.
3. A reduced "kernel MPN" diagnostic: an ablation that replaces the transformer stack with a similarity-weighted average of support labels in a meta-learned feature space, used to test how much of the full model's performance is explained by simple molecular similarity.
4. An empirical comparison on FS-Mol regression against ADKF-IFT, TabPFNv2, TabICLv2, a tuned random forest, and a Gaussian process with a Tanimoto kernel, plus timing measurements and a public release of code, experiment code, and trained weights.

Main Findings

• Parity at the smallest support size: at 16 support molecules, MPN (16 layers) reaches Kendall's τ of 0.298 and MPN (4 layers) 0.299, matching ADKF-IFT at 0.298. Beyond that, ADKF-IFT remains the strongest method, scoring 0.349, 0.391, 0.434 and 0.456 at 32, 64, 128 and 256 support molecules.
• Depth helps at larger support sets: moving from 4 to 16 layers changes τ from 0.299 to 0.298 at 16 molecules, but improves it from 0.326 to 0.335 at 32, 0.352 to 0.368 at 64, 0.382 to 0.398 at 128, and 0.395 to 0.411 at 256.
• The kernel baseline is unexpectedly strong: kernel MPN scores 0.233, 0.282, 0.326, 0.369 and 0.379 across the five support sizes, beating both the tuned random forest (0.138 to 0.354) and GP-ST (0.163 to 0.358) at every support size, while remaining behind the full transformer from 32 molecules upward.
• Similarity alone is not the whole story: because the full transformer outperforms the kernel variant, the authors conclude that MPN performance is only partially explained by kernel regression over molecular similarity in the learned feature space.
• Generic tabular priors underperform a domain-specific one: TabPFNv2 ranges from 0.146 to 0.322 and TabICLv2 from 0.150/0.156 to 0.350/0.358 (500-bit and 2048-bit fingerprints respectively), both below the MPNs.
• Speed: MPNs perform inference over thousands of molecules in seconds (attention) or hundredths of a second (kernel); both MPN and ADKF-IFT are roughly two times faster on a GPU, but GPU memory becomes a bottleneck for ADKF-IFT beyond 1000 samples.
• Benchmark caveat: the 256-support-set column is computed only on the 43 of 111 test tasks that have enough samples to form a 256-molecule support set.

Methodology in Plain English

Each molecule is converted into a fixed-length binary fingerprint (Morgan, radius 2, length 2048) and treated as a single token, so individual fingerprint bits keep consistent meaning across datasets. A task is presented to the model as a labelled "support" set and an unlabelled "query" set.

Labels are normalized using the mean and standard deviation of the support set and embedded by a separate linear map into a smaller space (one quarter of the hidden dimension). In every transformer block, the support tokens are concatenated with their label embeddings to form attention values, and attention is asymmetric: queries and supports both attend only to support molecules, so information flows only from labelled data. No positional encoding is used, making predictions invariant to the order of the support set. The label embeddings are re-injected at every layer, which the authors found improves performance. A final linear head reads predictions off the query tokens.

Training is meta-learning: at each step a dataset is sampled from a corpus of existing affinity datasets, randomly split into support and query sets with sizes drawn from a distribution, and the model is trained with mean squared error on the query labels. In the experiments, training used the 4938 FS-Mol training tasks with pIC50 targets; datasets larger than 500 molecules were randomly subsampled to 500; support sizes were drawn uniformly from {16, 32, 64, 128, 256} and capped at 80% of the dataset so at least 20% remained as queries. The hidden dimension was 1024, the feed-forward inner dimension 4d, and optimization used Adam with learning rates of 10⁻⁴ (kernel) and 10⁻⁵ (attention), batches of 10 tasks, and gradient norm clipping at 100. The kernel variant trained for 100 epochs and the attention models for 400; a 16-layer MPN took approximately 4 days on a single NVIDIA A10G GPU. Model selection used mean absolute error on 40 validation tasks, and results are averaged over 111 test tasks and 10 random splits.

Why This Matters

The work reframes affinity prediction from "fit a model per project" to "amortize learning across many past projects," which matters when each new drug program starts with only a handful of measured compounds.

• Virtual screening and hit triage: ranking large candidate libraries against a target using only a small set of measured actives.
• Lead optimization: re-ranking analogs after each round of synthesis, where assay data arrive in small batches.
• Cross-target knowledge transfer: reusing patterns learned from historical affinity datasets for targets with no dedicated training data.
• Cheap deployment: a single forward pass means no hyperparameter tuning or per-task training, reducing compute per project.

Industry relevance: for pharmaceutical and biotech teams, and contract research organizations, the appeal is turnaround time and reproducibility. A pre-trained model that predicts in seconds, with published weights, is easier to slot into existing screening pipelines than a model that must be re-fit and tuned per target. The paper also provides a caution about generic tabular foundation models: when a domain-specific pre-training corpus exists, it can outperform generic priors.

Future Directions

• Uncertainty quantification: the authors propose explicit modelling of prediction uncertainty to enable use in active learning and Bayesian optimization.
• Broader label space: training on a wider range of experimentally and computationally derived molecular properties, turning MPNs into a general foundation model for molecular property prediction.
• Better molecular representations: augmenting Morgan fingerprints with additional descriptors, or replacing them with a graph neural network encoder trained end-to-end with the transformer layers.
• Scaling to larger support sets: the reported timings suggest inference over thousands of molecules is feasible, and the authors frame applying the architecture across a wider range of support set sizes as future work.

Target Audience

Machine learning researchers working on drug discovery and molecular property prediction; cheminformaticians interested in few-shot and meta-learning methods; practitioners evaluating whether pre-trained foundation models can replace per-target model fitting; and methodologists interested in Prior-data Fitted Networks and the boundary between attention-based and kernel-based inference.

Authors’ abstract

Predicting affinity between small molecules and their protein targets remains one of the central problems in drug discovery. Due to the substantial cost of wet-lab experiments, machine learning models of drug-target affinity are usually trained on relatively small datasets, ranging from tens to hundreds of molecules. Given the inherent difficulty of generalization between different targets and chemical series, existing datasets of protein-ligand affinities are typically not used for new projects. At the same time, tabular foundation models, including Prior-data Fitted Networks, have recently shown how transformers can be used to solve machine learning problems in a single forward pass of the network, after being trained over a large corpus of, typically, synthetic datasets. Inspired by that line of work, we propose a novel affinity modelling approach, in which the labels of the query molecules are predicted from the labelled support molecules in a single forward pass of a transformer model. The model is trained on thousands of existing affinity datasets, enabling it to capture the general patterns of molecular activity across different targets, chemical series and dataset sizes. For the smallest support sets, our model matches the state of the art on the regression variant of the FS-Mol benchmark, while being faster than existing methods, thanks to the lack of task-specific training. A comparison against a reduced variant of our architecture, which amounts to kernel regression in a meta-learned feature space, shows that the predictive performance of the transformer model is only partially explained by simple molecular similarity.