Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Cost

TL;DR

This paper introduces Euclidean fast attention (EFA), a linear-scaling attention mechanism with Euclidean rotary positional encodings (ERoPE), enabling efficient modeling of long-range correlations in Euclidean data, particularly for computational chemistry applications.

Contribution

The paper presents EFA and ERoPE, novel methods that enable linear-scaling attention for Euclidean data, improving long-range interaction modeling in machine learning models.

Findings

EFA captures diverse long-range effects effectively.

EFA-equipped MLFFs outperform conventional models on chemical interactions.

EFA reduces computational complexity from quadratic to linear.

Abstract

Long-range correlations are essential across numerous machine learning tasks, especially for data embedded in Euclidean space, where the relative positions and orientations of distant components are often critical for accurate predictions. Self-attention offers a compelling mechanism for capturing these global effects, but its quadratic complexity presents a significant practical limitation. This problem is particularly pronounced in computational chemistry, where the stringent efficiency requirements of machine learning force fields (MLFFs) often preclude accurately modeling long-range interactions. To address this, we introduce Euclidean fast attention (EFA), a linear-scaling attention-like mechanism designed for Euclidean data, which can be easily incorporated into existing model architectures. A core component of EFA are novel Euclidean rotary positional encodings (ERoPE), which enable efficient encoding of spatial information while respecting essential physical symmetries. We empirically demonstrate that EFA effectively captures diverse long-range effects, enabling EFA-equipped MLFFs to describe challenging chemical interactions for which conventional MLFFs yield incorrect results.