# So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems

arXiv: https://arxiv.org/abs/2205.14276 · HTML: https://mlsanity.com/paper/2205.14276/

## Key facts

- Cited by 39 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 64 references that are in the map.
- First posted on arXiv 28 May 2022.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [Machine Learning Force Fields](https://mlsanity.com/paper/2010.07067/) (2020)
- [TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials](https://mlsanity.com/paper/2202.02541/) (2022)
- [Point Convolutional Neural Networks by Extension Operators](https://mlsanity.com/paper/1803.10091/) (2018)
- [Machine learning for molecular simulation](https://mlsanity.com/paper/1911.02792/) (2019)
- [SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects](https://mlsanity.com/paper/2105.00304/) (2021)
- [XAI for Transformers: Better Explanations through Conservative Propagation](https://mlsanity.com/paper/2202.07304/) (2022)
- [SE(3)-equivariant prediction of molecular wavefunctions and electronic densities](https://mlsanity.com/paper/2106.02347/) (2021)
- [Hierarchical modeling of molecular energies using a deep neural network](https://mlsanity.com/paper/1710.00017/) (2017)
- [Machine Learning Molecular Dynamics for the Simulation of Infrared Spectra](https://mlsanity.com/paper/1705.05907/) (2017)

## Built on by (most cited first, 30 of 39)

- [FAENet: Frame Averaging Equivariant GNN for Materials Modeling](https://mlsanity.com/paper/2305.05577/) (2023)
- [A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems](https://mlsanity.com/paper/2312.07511/) (2023)
- [Geometric Algebra Transformer](https://mlsanity.com/paper/2305.18415/) (2023)
- [Ewald-based Long-Range Message Passing for Molecular Graphs](https://mlsanity.com/paper/2303.04791/) (2023)
- [ET-Flow: Equivariant Flow-Matching for Molecular Conformer Generation](https://mlsanity.com/paper/2410.22388/) (2024)
- [The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains](https://mlsanity.com/paper/2410.24169/) (2024)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [Navigating protein landscapes with a machine-learned transferable coarse-grained model](https://mlsanity.com/paper/2310.18278/) (2023)
- [Are High-Degree Representations Really Unnecessary in Equivariant Graph Neural Networks?](https://mlsanity.com/paper/2410.11443/) (2024)
- [Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing](https://mlsanity.com/paper/2405.15389/) (2024)
- [Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems](https://mlsanity.com/paper/2502.19227/) (2025)
- [Euclidean, Projective, Conformal: Choosing a Geometric Algebra for Equivariant Transformers](https://mlsanity.com/paper/2311.04744/) (2023)
- [Higher-Rank Irreducible Cartesian Tensors for Equivariant Message Passing](https://mlsanity.com/paper/2405.14253/) (2024)
- [Probing the effects of broken symmetries in machine learning](https://mlsanity.com/paper/2406.17747/) (2024)
- [Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Cost](https://mlsanity.com/paper/2412.08541/) (2024)
- [Pushing the limits of unconstrained machine-learned interatomic potentials](https://mlsanity.com/paper/2601.16195/) (2026)
- [Molecular relaxation by reverse diffusion with time step prediction](https://mlsanity.com/paper/2404.10935/) (2024)
- [Accurate Ab-initio Neural-network Solutions to Large-Scale Electronic Structure Problems](https://mlsanity.com/paper/2504.06087/) (2025)
- [Towards Symbolic XAI -- Explanation Through Human Understandable Logical Relationships Between Features](https://mlsanity.com/paper/2408.17198/) (2024)
- [Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics](https://mlsanity.com/paper/2601.22123/) (2026)
- [GeoMFormer: A General Architecture for Geometric Molecular Representation Learning](https://mlsanity.com/paper/2406.16853/) (2024)
- [Learning Equivariant Non-Local Electron Density Functionals](https://mlsanity.com/paper/2410.07972/) (2024)
- [Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration](https://mlsanity.com/paper/2501.19179/) (2025)
- [Equivariant Matrix Function Neural Networks](https://mlsanity.com/paper/2310.10434/) (2023)
- [Reconstructing Kernel-based Machine Learning Force Fields with Super-linear Convergence](https://mlsanity.com/paper/2212.12737/) (2022)
- [Representing spherical tensors with scalar-based machine-learning models](https://mlsanity.com/paper/2505.05404/) (2025)
- [PRISM: Periodic Representation with multIscale and Similarity graph Modelling for enhanced crystal structure property prediction](https://mlsanity.com/paper/2511.20362/) (2025)
- [Stress and heat flux via automatic differentiation](https://mlsanity.com/paper/2305.01401/) (2023)
- [From Peptides to Nanostructures: A Euclidean Transformer for Fast and Stable Machine Learned Force Fields](https://mlsanity.com/paper/2309.15126/) (2023)
- [FreeCG: Free the Design Space of Clebsch-Gordan Transform for Machine Learning Force Fields](https://mlsanity.com/paper/2407.02263/) (2024)

## Abstract (arXiv)

The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates - a self-attention based message passing neural network - uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism. We show that in contrast to other published methods, So3krates is able to describe non-local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. So3krates matches or exceeds state-of-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25 - 0.4x) while at the same time giving a substantial speedup (6 - 14x for training and 2 - 11x for inference) compared to other models.