SE(3)-Hyena Operator for Scalable Equivariant Learning

要約

等分散性を維持しながらグローバルな幾何学的コンテキストをモデル化することは、生物学、化学、視覚などの多くの分野で正確な予測を行うために非常に重要です。
しかし、高次元データを大規模に処理するには計算量が必要となるため、これは困難です。
等変セルフアテンションや距離ベースのメッセージパッシングなどの既存のアプローチは、シーケンスの長さに関して二次的な複雑さの影響を受ける一方、局所的な方法ではグローバルな情報が犠牲になります。
状態空間モデルと長期畳み込みモデルの最近の成功に触発されて、この研究では、Hyena オペレーターに基づく等変長期畳み込みモデルである SE(3)-Hyena オペレーターを導入します。
SE(3)-Hyena は、回転と平行移動に対する等変性を維持しながら、二次二次複雑さでグローバルな幾何学的コンテキストを捕捉します。
等変連想想起と n-body モデリングで評価すると、SE(3)-Hyena は等変自己注意と同等またはそれを上回り、長いシーケンスに必要なメモリと計算リソースが大幅に少なくなります。
私たちのモデルは、等変変換器よりも 3.5 倍速く 20,000 トークンの幾何学的コンテキストを処理し、同じメモリ バジェット内で 175 倍長いコンテキストを許可します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google