Geometric Hyena Networks for Large-scale Equivariant Learning

TL;DR

Geometric Hyena introduces an efficient equivariant long-convolutional model that captures global geometric context at scale, outperforming existing methods in biological molecular modeling tasks with reduced computational costs.

Contribution

It is the first equivariant long-convolutional model for geometric systems, enabling scalable global context processing with sub-quadratic complexity and high efficiency.

Findings

Outperforms existing equivariant models in biological tasks

Processes 30k tokens 20x faster than equivariant transformer

Allows 72x longer context within the same computational budget

Abstract

Processing global geometric context while preserving equivariance is crucial when modeling biological, chemical, and physical systems. Yet, this is challenging due to the computational demands of equivariance and global context at scale. Standard methods such as equivariant self-attention suffer from quadratic complexity, while local methods such as distance-based message passing sacrifice global information. Inspired by the recent success of state-space and long-convolutional models, we introduce Geometric Hyena, the first equivariant long-convolutional model for geometric systems. Geometric Hyena captures global geometric context at sub-quadratic complexity while maintaining equivariance to rotations and translations. Evaluated on all-atom property prediction of large RNA molecules and full protein molecular dynamics, Geometric Hyena outperforms existing equivariant models while requiring significantly less memory and compute that equivariant self-attention. Notably, our model processes the geometric context of 30k tokens 20x faster than the equivariant transformer and allows 72x longer context within the same budget.