SE(3)-Hyena Operator for Scalable Equivariant Learning
Keywords: architecrture, equivariance, global context, long convolution, scalability, mechanistic interpretability
TL;DR: We introduce the SE(3)-Hyena, an equivariant long-convolutional model that efficiently captures global geometric context at sub-quadratic complexity
Abstract: Modeling global geometric context while maintaining equivariance is crucial for accurate predictions in many fields such as biology, chemistry, or when modeling physical systems. Yet, this is challenging due to the computational demands of processing high-dimensional data at scale. Existing approaches such as equivariant self-attention or distance-based message passing, suffer from quadratic complexity with respect to sequence length, while localized methods sacrifice global information. Inspired by the recent success of state-space and long-convolutional models, in this work, we introduce the SE(3)-Hyena operator, the first equivariant network that adopts a long-convolutional framework for geometric systems. SE(3)-Hyena captures global geometric context at sub-quadratic complexity while maintaining equivariance to rotations and translations. Evaluated on the task of all-atom property prediction of large RNA molecules, SE(3)-Hyena matches or outperforms equivariant self-attention while requiring significantly less memory and compute for long geometric sequences. Additionally, we propose equivariant associative recall as a new mechanistic interpretability task for studying the contextual learning capabilities of equivariant models. Notably, our model processes the geometric context of $30k$ tokens $20 \times$ faster than the equivariant transformer and allows $72 \times$ longer context within the same memory budget. The code will be released upon the acceptance.
Primary Area: learning on graphs and other geometries & topologies
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 3951
Loading