Geometric Transformer for End-to-End Molecule Properties Prediction

TL;DR

This paper introduces a geometry-aware Transformer architecture for molecule property prediction that effectively captures molecular structure, improves performance over state-of-the-art methods, and is designed to work well with small datasets without domain-specific knowledge.

Contribution

The authors propose a novel Transformer-based model with geometric encoding and data augmentation for molecule property prediction, outperforming existing methods without relying on domain knowledge.

Findings

Outperforms state-of-the-art methods in molecule property prediction

Effective in small dataset regimes due to augmentation scheme

Captures molecular geometry using modified positional encoding

Abstract

Transformers have become methods of choice in many applications thanks to their ability to represent complex interactions between elements. However, extending the Transformer architecture to non-sequential data such as molecules and enabling its training on small datasets remains a challenge. In this work, we introduce a Transformer-based architecture for molecule property prediction, which is able to capture the geometry of the molecule. We modify the classical positional encoder by an initial encoding of the molecule geometry, as well as a learned gated self-attention mechanism. We further suggest an augmentation scheme for molecular data capable of avoiding the overfitting induced by the overparameterized architecture. The proposed framework outperforms the state-of-the-art methods while being based on pure machine learning solely, i.e. the method does not incorporate domain knowledge from quantum chemistry and does not use extended geometric inputs besides the pairwise atomic distances.