SE(3)-Hyena Operator for Scalable Equivariant Learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Moskalev, Artem, Prakash, Mangal, Liao, Rui, Mansi, Tommaso
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929456888676352
author Moskalev, Artem
Prakash, Mangal
Liao, Rui
Mansi, Tommaso
author_facet Moskalev, Artem
Prakash, Mangal
Liao, Rui
Mansi, Tommaso
contents Modeling global geometric context while maintaining equivariance is crucial for accurate predictions in many fields such as biology, chemistry, or vision. 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 SE(3)-Hyena operator, an equivariant long-convolutional model based on the Hyena operator. The SE(3)-Hyena captures global geometric context at sub-quadratic complexity while maintaining equivariance to rotations and translations. Evaluated on equivariant associative recall and n-body modeling, SE(3)-Hyena matches or outperforms equivariant self-attention while requiring significantly less memory and computational resources for long sequences. Our model processes the geometric context of 20k tokens x3.5 times faster than the equivariant transformer and allows x175 longer a context within the same memory budget.
format Preprint
id arxiv_https___arxiv_org_abs_2407_01049
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle SE(3)-Hyena Operator for Scalable Equivariant Learning
Moskalev, Artem
Prakash, Mangal
Liao, Rui
Mansi, Tommaso
Machine Learning
Modeling global geometric context while maintaining equivariance is crucial for accurate predictions in many fields such as biology, chemistry, or vision. 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 SE(3)-Hyena operator, an equivariant long-convolutional model based on the Hyena operator. The SE(3)-Hyena captures global geometric context at sub-quadratic complexity while maintaining equivariance to rotations and translations. Evaluated on equivariant associative recall and n-body modeling, SE(3)-Hyena matches or outperforms equivariant self-attention while requiring significantly less memory and computational resources for long sequences. Our model processes the geometric context of 20k tokens x3.5 times faster than the equivariant transformer and allows x175 longer a context within the same memory budget.
title SE(3)-Hyena Operator for Scalable Equivariant Learning
topic Machine Learning
url https://arxiv.org/abs/2407.01049