Power law attention biases for molecular transformers

TL;DR

This paper introduces power law-based attention biases inspired by physical laws to improve molecular transformers, outperforming traditional positional encodings and enabling more interpretable models.

Contribution

It proposes a novel family of attention biases based on power laws that effectively encode molecular structure and enhance transformer performance.

Findings

Outperforms positional encodings and graph attention on QM9 and SPICE datasets.

Can compensate for the removal of scaled dot-product attention.

Enables more interpretable molecular transformer models.

Abstract

Transformers are the go-to architecture for most data modalities due to their scalability. While they have been applied extensively to molecular property prediction, they do not dominate the field as they do elsewhere. One cause may be the lack of structural biases that effectively capture the relationships between atoms. Here, we investigate attention biases as a simple and natural way to encode structure. Motivated by physical power laws, we propose a family of low-complexity attention biases $b_{ij} = p \log|| \mathbf{r}_i - \mathbf{r}_j||$ which weigh attention probabilities according to interatomic distances. On the QM9 and SPICE datasets, this approach outperforms positional encodings and graph attention while remaining competitive with more complex Gaussian kernel biases. We also show that good attention biases can compensate for a complete ablation of scaled dot-product attention, suggesting a low-cost path toward interpretable molecular transformers.