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

2024-12-11

arXiv.org

AI Summary

To address the high computational cost of modeling long-range interactions in machine learning force fields (MLFFs), this work introduces the first linear-complexity Euclidean-space attention mechanism. The method integrates two key innovations: (1) Euclidean Rotation Position Encoding (ERoPE), the first position encoding that rigorously preserves translational, rotational, and reflectional equivariance; and (2) symmetry-constrained embedding coupled with a linear-scaling self-attention approximation, enabling equivariant global representation learning for 3D atomic systems. The approach significantly improves prediction accuracy for critical chemical interactions—including van der Waals forces, charge transfer, and long-range polarization—outperforming classical force fields across multiple benchmark datasets. With an inference complexity of O(N), it achieves both physical consistency and computational scalability, marking a step toward efficient, physically principled MLFFs for large-scale molecular simulation.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Computer Vision: 3D Computer Vision

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Capturing long-range correlations in Euclidean data efficiently
• Overcoming quadratic complexity of self-attention mechanisms
• Modeling chemical interactions with linear-scaling attention

Innovation

• Linear-scaling attention mechanism for Euclidean data
• Novel Euclidean rotary positional encodings for spatial information
• Efficiently captures long-range interactions in chemical systems