Geometric Hyena Networks for Large-scale Equivariant Learning
This addresses a computational bottleneck for researchers in biology, chemistry, and physics by enabling efficient global equivariant learning, though it is an incremental improvement over existing long-convolutional models adapted for geometric systems.
The paper tackles the challenge of processing global geometric context with equivariance in large-scale biological and physical systems by introducing Geometric Hyena, an equivariant long-convolutional model that outperforms existing methods in tasks like RNA property prediction and protein dynamics, achieving 20x faster processing and 72x longer context within the same budget.
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.