PaperlayerView exact revision on arXiv

Source-linked AI summary

Molecule Attention Transformer

Łukasz Maziarka, Tomasz Danel, Sławomir Mucha, Krzysztof Rataj, Jacek Tabor, Stanisław Jastrzębski

arXiv:2002.08264v1cs.LGphysics.comp-phstat.ML

TL;DR

Molecular property prediction lacks a single neural architecture that is competitive across a broad range of tasks, despite its importance to drug discovery. MAT adapts Transformer attention to molecular graphs and inter-atomic distances. The paper reports strong cross-task performance, improved results with simple pretraining, reduced tuning requirements, and chemically interpretable attention weights.

  • Problem

    Molecular property prediction lacks a single neural architecture that is competitive across a broad range of tasks, despite its importance to drug discovery.

  • Method

    MAT adapts Transformer self-attention to molecules by incorporating inter-atomic distances and molecular graph structure.

  • Results

    MAT performs well across a wide range of molecule property prediction tasks, while self-supervised pretraining further improves performance and drastically reduces hyperparameter tuning.

  • Takeaways & Limitations

    MAT provides a versatile molecular-property architecture with simpler tuning requirements and attention weights that are interpretable from a chemical perspective.

  • Takeaways & Limitations

    The paper leaves exploration of more sophisticated compound 3D-structure minimization and better pretraining tasks for future work.

Abstract

from arXiv · show

Designing a single neural network architecture that performs competitively across a range of molecule property prediction tasks remains largely an open challenge, and its solution may unlock a widespread use of deep learning in the drug discovery industry. To move towards this goal, we propose Molecule Attention Transformer (MAT). Our key innovation is to augment the attention mechanism in Transformer using inter-atomic distances and the molecular graph structure. Experiments show that MAT performs competitively on a diverse set of molecular prediction tasks. Most importantly, with a simple self-supervised pretraining, MAT requires tuning of only a few hyperparameter values to achieve state-of-the-art performance on downstream tasks. Finally, we show that attention weights learned by MAT are interpretable from the chemical point of view.

1. Introduction

Molecular property prediction is important for drug discovery and material design, but deep networks can be difficult to train and may underperform shallow models. MAT is proposed as a broadly effective, easier-to-use deep model for this setting.

  • Motivation: Molecular property prediction supports applications including drug discovery and material design.The paper notes that correctly predicting properties such as toxicity or bioactivity could potentially avoid some costly clinical-trial failures.
  • Motivation: 85% of United States drug candidates are estimated to fail clinical trials after a long and costly development process.
  • Background: Deep neural networks are expected to transform drug discovery and material design, following breakthroughs in image and text classification.
  • Challenge: On many molecular property prediction tasks, deep neural networks are outperformed by support vector machines or random forests.Deep networks can also be difficult to train and require tuning many hyperparameters.
  • Aim: The paper aims to develop a deep model that is simple to use out-of-the-box and performs strongly across many molecule property prediction tasks.
  • Contribution: MAT augments Transformer self-attention with inter-atomic distances and molecular graph structure, achieving strong performance across a wide range of tasks.Self-supervised pretraining further improves performance and drastically reduces hyperparameter-tuning time; the experiments tune only the learning rate across 7 values.

2. Related work

Prior molecular-property models use physics, data-driven learning, molecular graphs, three-dimensional structure, or molecular strings, each addressing different modeling needs. MAT unifies graph- and distance-aware molecular modeling with Transformer self-attention in one architecture.

  • Molecular property prediction: Physics-based molecular-property prediction remains prohibitively costly for accurately predicting many properties, while data-driven methods face a lack of large datasets.
  • Deep learning approaches: Molecule property prediction methods have progressed from handcrafted representations to molecular strings and currently popular molecular-graph approaches.
  • Graph models: Graph convolutional networks aggregate information from adjacent nodes, and graph structure improves performance across molecule-modeling tasks.
  • Graph models: Attention mechanisms and learned neighborhood functions have been added to graph convolutional networks to improve neighbor-data aggregation.
  • Three-dimensional structure: 3DGCN and MPNN integrate graph and distance information, achieving strong performance on tasks such as solubility prediction.
  • Transformer models: Prior Transformer applications represent compounds as SMILES sentences, whereas MAT represents compounds as atom lists and augments self-attention with molecular structure.The paper reports that ablation studies identify this structural augmentation as a critical component.
  • Contribution: MAT unifies related ideas in a Transformer-based model intended to preserve strong performance across many chemical tasks.
  • Motivation: Deep neural networks are not always competitive with support vector machines or random forests, and graph neural networks can be difficult to train.The paper describes this as making deep learning less applicable to some molecule property prediction settings.

3. Molecule Attention Transformer

MAT adapts Transformer self-attention to molecular inputs by incorporating graph adjacency and inter-atomic distances, avoiding linearized molecular representations. The model replaces standard self-attention layers with Molecule Self-Attention and uses masked-node pretraining to stabilize learning and reduce hyperparameter search.

  • Molecule Self-Attention: MAT augments Transformer self-attention with molecular graph adjacency and inter-atomic distances instead of encoding molecules as sentences.The adjacency matrix represents molecular connectivity, while the distance matrix represents inter-atomic distances.
  • Molecule Self-Attention: Molecule Self-Attention combines self-attention, distance, and adjacency matrices using weighting scalars λa, λd, and λg.The distance function can use row-normalized softmax or element-wise exp(−d).
  • Molecule Attention Transformer: The architecture can incorporate additional molecular information, such as forces between atoms, beyond adjacency and distance matrices.This extension is presented as straightforward within the augmented attention mechanism.
  • Molecule Attention Transformer: MAT replaces every self-attention layer in the Transformer encoder with Molecule Self-Attention layers and embeds each atom as a 26 dimensional vector.The weighting scalars are treated as hyperparameters and kept frozen during training.
  • Pretraining: Masked-node self-supervised pretraining stabilizes learning and reduces the need for extensive hyperparameter search.The paper reports that MAT already performs well with this simple pretraining task and leaves other pretraining tasks for future work.
  • Implementation and limitations: Each compound is represented using one pre-computed 3D conformer, with distance matrices calculated from RDKit-generated conformers.Sampling additional conformations did not produce a consistent performance boost, though more sophisticated structure minimization could improve results.

4. Experiments

Experiments evaluate MAT across seven molecular property datasets, showing strong average-rank performance, benefits from pretraining, and chemically interpretable attention patterns. Ablations indicate that graph structure, inter-atomic distances, and the dummy node contribute to performance, while extra edge features do not.

  • Molecule Attention Transformer: Graph neural networks do not outperform other models on average; the best graph model ranks 3.28 compared with RF’s 3.14.Deep-model performance improves with larger hyperparameter-search budgets, but common deep-learning methods can require substantial computation and still perform poorly.
  • Molecule Attention Transformer: MAT achieves the best average rank across seven datasets, scoring 2.71 with 150 evaluations and 2.42 with 500, versus 3.14 for RF.The results attribute this broad performance to MAT’s flexible architecture and molecular inductive bias.
  • Molecule Attention Transformer: RF and SVM outperform MAT on Estrogen-β, MetStablow, and MetStabhigh because extended-connectivity fingerprints encode substructures relevant to metabolic stability.The passage connects metabolic stability to particular molecular moieties recognized by enzymes.
  • Ablation studies: Removing the dummy node reduces performance, while retaining graph and distance information produces the most stable performance across three tasks.MAT can also use three-dimensional distance information in a toy task involving functional groups.
  • Ablation studies: Adding additional edge features does not improve MAT performance, although richer features or better integration methods might help in future work.The authors leave comprehensive edge-feature design and integration for future work.
  • Attention analysis: Many MAT attention heads focus on the dummy node or a few atoms, with early layers learning simple chemical patterns and later layers capturing more complex arrangements.On BBBP, selected heads assign higher attention weights to atoms matching chemically specified SMARTS patterns.
  • Pretrained Molecule Attention Transformer: Pretrained MAT converges faster and has lower training-error variance than MAT on ESOL and FreeSolv.The experiments use the learning rate giving the best generalization for each dataset and repeat training six times.

5. Conclusions.

MAT is presented as a versatile architecture for molecular property prediction that performs well across many tasks. Self-supervised pretraining further improves performance and reduces hyperparameter-tuning needs, while better pretraining tasks remain future work.

  • MAT performs well across a wide range of molecule property prediction tasks.
  • Self-supervised pretraining further improves MAT’s performance and drastically reduces the need for hyperparameter tuning.
  • Exploring better pretraining tasks for MAT is identified as a promising direction for future work.

A. Dataset details.

The benchmark includes regression and binary-classification datasets covering solubility, barrier permeability, metabolic stability, and estrogen-receptor activity. These tasks represent molecular properties relevant to drug design and bioavailability.

  • FreeSolv and ESOL are regression tasks for predicting hydration free energy and logS, respectively.Molecular solubility influences drug bioavailability.
  • BBBP is a binary-classification task for predicting blood-brain-barrier penetration.The task is relevant when drug design seeks either central-nervous-system access or avoidance of brain effects.
  • MetStabhigh and MetStablow classify compounds with high or low metabolic stability against the other stability categories.The underlying stability measure is compound half-life within an organism.
  • Estrogen Alpha and Estrogen Beta are binary-classification tasks based on active and inactive compounds for the estrogen receptors.The passage links these tasks to drug-discovery concerns about receptor potency.

B. Other experimental details

The experimental details define MAT’s tunable architecture, optimization settings, pooling and prediction pathway, and baseline implementations. Baselines use DeepChem configurations or extended-connectivity fingerprints as specified.

  • MAT hyperparameters include model dimension, encoder layers, attention heads, feed-forward layers, distance-matrix kernel, dropout, and weight decay.The distance-matrix kernel transforms the distance matrix D.
  • The reported experimental hyperparameter tables cover MAT, SVM, RF, GCN, Weave, and EAGCN.
  • The model is trained with Adam and a learning-rate scheduler using optimizer factor 100 · LEARNING RATE and warmup steps based on WARMUP FACTOR.The scheduler increases learning rate linearly during warmup and then decreases it proportionally to the inverse square root of the step number.
  • After N encoder layers, MAT mean-pools atom representations and passes the resulting molecule embedding through a single linear prediction layer.
  • SVM, RF, GCN, and Weave baselines use DeepChem implementations with the hyperparameter-tuning settings and proposed defaults described in the experiments.
  • RF and SVM operate on 2048-dimensional extended-connectivity fingerprints computed with radius 2.

C. Additional results for Sec. 4.2

An additional internal-energy experiment tests MAT on a quantum-property regression task involving atom-level interactions. Pretrained MAT slightly outperforms GCN, while fingerprint-based RF performs much worse; training loss eventually plateaus for both graph models.

  • Additional internal-energy experiment: The additional experiment predicts internal energy at 298.15 K from the Alchemy dataset of 200K molecules with 12 quantum properties.The experiment tests whether molecule self-attention performs well on atom-level interaction tasks.
  • Additional internal-energy experiment: Three methods—RF, GCN, and pretrained MAT—are compared using mean absolute error on the original dataset splits.RF and GCN receive random searches with 500 hyperparameter sets, while pretrained MAT tunes only learning rate.
  • Additional internal-energy experiment: MAT achieves a slightly lower error than GCN, whereas fingerprint-based RF has an MAE almost two orders of magnitude higher than the other methods.The authors attribute graph-method performance to locality preserved in graph structure.
  • Training error for graph-based neural networks: Training loss for both GCN and MAT eventually flattens as layer count and model dimensionality increase.Figure 7 reports training loss as a function of these two architectural variables.

D. Additional details for Sec. 4.3

The pretraining details describe masked-node prediction on a large molecular corpus, followed by standardized fine-tuning and selected dataset-specific handling.

  • Pretraining: Masked-node prediction trains the model to infer atom attributes from neighboring nodes, encouraging domain-specific relationships between atoms.The task replaces input nodes with mask tokens and predicts their attributes from corresponding node embeddings.
  • Pretraining setting: 2 mln ZINC15 molecules formed the pretraining dataset, with models trained for 8 epochs using binary cross entropy.MAT used Noam optimization, while EAGCN used Adam; both used learning rate 0.001 and batch size 256.
  • Fine-tuning setting: Fine-tuning used 100 epochs, batch size 32, and learning-rate selection from seven candidate values.The candidate learning rates ranged from 1e−3 to 1e−6.
  • Dataset handling: Three largest molecules were excluded from Estrogen Alpha experiments because of memory issues.
  • Reference settings: The section provides hyperparameter tables for pretrained MAT, pretrained EAGCN, and SMILES-Transformer comparison settings.

E. Additional results for Sec. 4.4

The additional results describe edge-feature encoding for MAT, producing an edge-feature matrix used in molecule self-attention.

  • Edge features: Each molecular bond is represented by an embedded vector of edge features that is transformed into a scalar by a linear layer and ReLU.
  • Molecule Self-Attention: The resulting E ∈ R^N_atoms×N_atoms matrix replaces the adjacency matrix in the Molecule Self-Attention layer.Nonexistent bonds are represented by passing a zero vector through the same layer.
  • Experimental configuration: Table 16 lists the edge features used for the experiments associated with Table 6.

F. Toy task

The toy task tests whether MAT can use inter-atomic distances to classify spatial relationships between molecular fragments. Performance improves with distance information and reaches 0.98 AUC, exceeding the reported GCN result of 0.93 AUC.

  • Motivation: The task is motivated by molecular geometry-dependent effects that graph information alone may not represent well.The authors note that emphasizing distance information may require lower values for λg or λa.
  • Task description: The task predicts whether an -NH2 fragment and tert-butyl group are closer than a predefined spatial threshold.The dataset contains 2677 PubChem molecules, using a 20 Å threshold selected for balanced positive and negative examples.
  • Experimental setting: MAT performance is evaluated while varying λd under three settings of λa and λg, with all three λ parameters constrained to sum to 1.
  • Results: 0.98 AUC is achieved by MAT with λd = 1.0, compared with 0.93 AUC for GCN on the toy task.
  • Results: Using distance information significantly improves performance across all three tested λ settings.The results motivate tuning the weighting parameters in the main experiments.
  • Baseline comparison: A grid search for GCN hyperparameters selected the best validation AUC configuration, whose test AUC was 0.925 ± 0.006.

G. Interpretability analysis

The interpretability analysis examines first-layer MAT attention heads for chemically meaningful patterns and tests those patterns quantitatively using SMARTS-defined atom groups.

  • Results: Table 19 reports means and standard deviations for matching atoms and other atoms across the six attention-head patterns.
  • Pattern discovery: Six first-layer attention-head patterns correspond to chemical structures identified through qualitative inspection.The analysis then tests hypotheses about what those heads represent quantitatively.
  • Quantitative test: SMARTS expressions select atoms matching each hypothesized pattern, whose mean attention weights are compared with nonmatching atoms.The comparison uses attention-matrix column averages on BBBP molecules.
Loading 2002.08264v1…