[Paper Review] Molecule Attention Transformer

Łukasz Maziarka, Tomasz Danel+4 more|arXiv (Cornell University)|Feb 19, 2020

Computational Drug Discovery MethodsComputer Science44 references108 citations

TL;DR

Molecule Attention Transformer (MAT) augments Transformer self-attention with inter-atomic distances and molecular graph structure, enabling competitive molecule property prediction and interpretable attention; pretrained MAT reduces hyperparameter tuning and improves performance.

ABSTRACT

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.

Motivation & Objective

• Develop a deep model for wide-ranging molecule property prediction tasks.
• Incorporate chemical structure information into Transformer attention via adjacency and distance data.
• Demonstrate that simple self-supervised pretraining reduces hyperparameter tuning and boosts performance.
• Show that attention weights learned by MAT are chemically interpretable.

Proposed method

• Replace Transformer self-attention with Molecule Self-Attention that blends attention, distance, and graph information via a weighted sum controlled by λa, λd, λg.
• Represent atoms with a 26-dimensional embedding incorporating atomic identity and local features.
• Compute distance matrix D from 3D conformers using RDKit and incorporate it into attention via g(D) (e.g., softmax over rows or exp(-d)).
• Augment input with a dummy node to allow the model to skip patterns if none are present.
• Pretrain MAT using node-level masked feature prediction to stabilize learning and reduce hyperparameter search.
• Evaluate on a diverse set of molecular datasets under scaffold and random splits and compare to baseline graph and fingerprint-based models.

Experimental results

Research questions

• RQ1 Can an attention mechanism augmented with molecular graph structure and inter-atomic distances perform competitively across diverse molecule property prediction tasks?
• RQ2 Does self-supervised pretraining enable strong performance with minimal hyperparameter tuning in MAT?
• RQ3 Are attention weights in MAT interpretable from a chemical perspective?

Key findings

• MAT achieves competitive or superior performance across a wide range of datasets compared to shallow models and several graph-based baselines.
• With pretraining, MAT outperforms non-pretrained variants and requires tuning of only a few hyperparameters (learning rate), reducing search effort significantly.
• Pretrained MAT attains the best average rank across datasets, using only 7 hyperparameter evaluations versus hundreds for other methods.
• Ablation studies show using the dummy node and incorporating graph and distance information improves stability and performance across tasks.
• Attention heads in MAT learn interpretable chemical patterns, with some heads focusing on specific atom types or local substructures.