Power law attention biases for molecular transformers

TL;DR

This work targets how molecular transformers encode interatomic structure, arguing that standard positional encodings are suboptimal for graphs or Euclidean point clouds. It introduces a power-law attention bias $b_{ij} = p \log || r_i - r_j ||$ to weight attention by interatomic distance, with $p$ theory-derived, fixed, or per-head learnable; the bias is integrated into the attention logits as $A_{ij} \propto || r_i - r_j ||^{p} \exp((q_i k_j^T)/\sqrt{d_k})$. On QM9 and SPICE, these biases outperform positional encodings and graph attention baselines, with Gaussian kernel biases often delivering the best accuracy and power-law biases with negative $p$ matching that performance at lower compute. The study also shows that fixed attention using the bias can substitute for dynamic dot-product attention, offering a cheaper, more interpretable path toward scalable molecular transformers, though the experiments are limited to a modest model size and dataset scale. Future work should scale data, model size, and compute to validate generality and explore broader deployment.

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.

Figures (2)

• Figure 1: Histogram of learned power law exponents. Exponents are collected from an ensemble of 7 independently trained models.
• Figure 2: Histogram of learned power law exponents when dot-product attention is ablated. Exponents are collected from an ensemble of 7 independently trained models.