SE(3)-Hyena Operator for Scalable Equivariant Learning

## Abstract

Abstract 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 20 ⁢ k 20 𝑘 20k 20 italic_k tokens × 3.5 absent 3.5 \times 3.5 × 3.5 faster than the equivariant transformer and allows × 175 absent 175 \times 175 × 175 longer a context within the same memory budget.

## 1 Introduction

Modeling global geometric context while preserving equivariance is crucial in many real-world tasks. The properties of a protein depend on the global interaction of its residues (Baker & Sali, ). Similarly, the global geometry of DNA and RNA dictates their functional properties (Leontis & Westhof, ; Sato et al., ). In computer vision, modeling global geometric context is crucial when working with point clouds or meshes (Thomas et al., ; De Haan et al., ). In all these tasks, maintaining equivariance while capturing global context is essential for robust modeling and prediction.

Processing global geometric context with equivariance is challenging due to the computational demands of processing high-dimensional data at scale. Existing methods either rely on global all-to-all operators such as self-attention (Liao & Smidt, ; de Haan et al., ; Brehmer et al., ), which do not scale well due to their quadratic complexity, or they restrict processing to local neighborhoods (Thomas et al., ; Köhler et al., ; Fuchs et al., ), losing valuable global information. This limitation is a significant practical bottleneck, necessitating more efficient solutions for scalable equivariant modeling with a global geometric context.

Figure: Figure 1: SE(3)-Hyena operator marries global context, equivariance and scalability towards long sequences. SE(3)-Hyena operator can process global geometric context in sub-quadratic time while preserving equivariance to rotations and translations.
Refer to caption: x1.png

An efficient algorithm for modeling global context should support parallelization during training while maintaining bounded computational costs relative to sequence length during inference. One approach involves recurrent operators (Orvieto et al., ; De et al., ), which provide bounded compute but lack easy parallelization. Another family of methods relies on self-attention (Vaswani et al., ) allowing parallel processing at the cost of quadratic computational complexity. The most recent advances leverage state-space (Gu et al., ; Fu et al., ; Gu & Dao, ) and long-convolutional (Romero et al., ; Poli et al., ) frameworks, enabling global context reasoning in sub-quadratic time with easy parallelization. Extending these models to accommodate equivariance remains an unexplored direction.

Inspired by the success of state-space and long-convolutional methods, in this work we propose a SE(3)-equivariant long-convolutional model based on the recent Hyena operator (Poli et al., ). The SE(3)-Hyena efficiently models global geometric context in sub-quadratic time while preserving equivariance to rotations and translations (Figure 1). Central to our method is the equivariant vector long convolution that leverages cross products between equivariant queries and keys. This vector convolution can be implemented in the Fourier domain with a computational complexity of $O(Nlog_{2}N)$, enabling scaling to much longer sequences than self-attention. We evaluate the proposed SE(3)-Hyena against its self-attention counterpart on a novel equivariant associative recall benchmark and the n-body dynamical system modeling task. Our results suggest that SE(3)-Hyena matches or surpasses equivariant self-attention in performance while requiring less memory and compute for long sequences. In particular, for a sequence of $20k$ tokens, the equivariant Hyena runs $\times 3.5$ faster than the equivariant transformer approach. Notably, when the equivariant transformer runs out of memory on sequences over $20k$ tokens, our model can handle up to $3.5M$ million tokens on a single GPU, providing up to $175$ times longer context length within the same computational budget.

To sum up, we make the following contributions:

- •
We propose SE(3)-equivariant Hyena operator which enables modeling global geometric context in sub-quadratic time.
- •
We propose an equivariant counterpart for the mechanistic interpretability associative recall task.
- •
We demonstrate that the equivariant Hyena matches or outperforms the equivariant transformer, while requiring significantly less memory and compute for long-context modeling.

## 2 Related work

##### Equivariance

Equivariance to group transformations, particularly rotations and translations in 3D, is crucial for modeling physical systems  (Zhang et al., ). Schütt et al. () condition continuous convolutional filters on relative distances to build model invariant to rotations and translations. Thomas et al. (); Fuchs et al. (); Brandstetter et al. (); Liao & Smidt (); Bekkers et al. () utilize spherical harmonics as a steerable basis which enables equivariance between higher-order representations. Since computing spherical harmonics can be expensive, Jing et al. (, ); Satorras et al. (); Deng et al. () focus on directly updating vector-valued features to maintain equivariance, while Zhdanov et al. () employ another equivariant network to implicitly parameterize steerable kernels. Another recent line of work  (Ruhe et al., , ; Brehmer et al., ; Zhdanov et al., ) employs geometric algebra representation which natively provides a flexible framework for processing symmetries in the data (Dorst et al., ).

While these works focus on how to build equivariance into a neural network, in this paper we focus on efficient equivariance to model global geometric contexts at scale.

##### Modeling geometric context

Various strategies are employed to process context information in geometric data. Convolutional methods aggregate context linearly within a local neighborhood, guided by either a graph topology (Kipf & Welling, ) or spatial relations in geometric graphs (Schütt et al., ; Wu et al., ; Thomas et al., ). Message-passing framework (Gilmer et al., ) generalizes convolutions, facilitating the exchange of nonlinear messages between nodes with learnable message functions. These approaches are favored for their simplicity, balanced computational demands, and expressiveness (Wu et al., ). However, they are limited to local interactions and are known to suffer from oversmoothing (Rusch et al., ). This hinders building deep message-passing networks capable of encompassing a global geometric context in a receptive field. To address these limitations, recent methods have turned to self-attention mechanisms for graph (Yang et al., ; Kreuzer et al., ; Kim et al., ; Rampášek et al., ) and geometric graph (Fuchs et al., ; Liao & Smidt, ; Brehmer et al., ) data, outperforming convolutional and message-passing approaches. Yet, the quadratic computational cost of self-attention poses significant challenges when modeling large-scale physical systems. In this work, we aim to develop a method for global geometric context processing with sub-quadratic computational complexity.

##### State-space and long-convolutional models

The quadratic computational complexity of self-attention has driven the exploration of alternatives for modeling long context. Structured state-space models (Gu et al., ) have emerged as a promising alternative, integrating recurrent and convolutional mechanisms within a single framework. These models enable parallelized training in a convolutional mode and maintain linear complexity with respect to sequence length in a recurrent mode. Models like S4 (Gu et al., ), H3 (Fu et al., ), and Mamba (Gu & Dao, ; Li et al., ) have consistently matched or exceeded transformer performance in diverse tasks such as genomics (Schiff et al., ), long-range language (Wang et al., ), and vision tasks (Zhu et al., ). Concurrently, another line of work integrates long-convolutional framework with implicit filters (Sitzmann et al., ; Romero et al., ; Zhdanov et al., ) to capture global sequence context. The implicit filter formulation allows for data-controlled filtering similar to transformers, while FFT-based long convolution enables global context aggregation in sub-quadratic time (Poli et al., ). Such models have shown competitive performance comparable to state-space and transformer architectures in time-series modeling (Romero et al., ), genomics (Nguyen et al., ), and vision tasks (Poli et al., ).

Although state-space and long-convolutional methods dramatically reduced the computational costs associated with processing long sequences, their application to geometric data requiring equivariance remains unexplored. In this work, we adapt the recently proposed Hyena operator (Poli et al., ) to incorporate SE(3) equivariance. To the best of our knowledge, this is the first equivariant long-convolutional model that can process global geometric contexts with sub-quadratic memory and time requirements.

## 3 Method

We consider tasks that require modeling invariant and equivariant features in geometric graphs. A geometric graph of $N$ nodes is represented by a set of features $\{\mathbf{x}_{i},\mathbf{f}_{i}\}_{i=1}^{N}$ where $\mathbf{x}_{i}\in\mathbb{R}^{3}$ represents vector features (e.g. coordinates or velocities), and $\mathbf{f}_{i}\in\mathbb{R}^{S}$ represents scalar features (e.g. atom types, charges or fingerprints). We call $\mathbf{x}_{i}$ geometric or vector tokens, and $\mathbf{f}_{i}$ are scalar tokens. When working with geometric graphs, a neural network must respect symmetries of the input space such as rotation or translation. That means, a model must be equivariant with respect to geometric tokens and invariant with respect to scalar tokens.

Figure: Figure 2: SE(3)-Hyena building blocks. (a) Schematic of existing Hyena architecture (Poli et al., ). (b) The proposed architecture consists of the SE(3)-Hyena operator, residual connections, and an equivariant MLP. (c) The block architecture of SE(3)-Hyena operator consists of two streams processing invariant and equivariant features. The key components are scalar and vector long convolution responsible for global context aggregation.
Refer to caption: x2.png

### 3.1 SE(3)-Hyena operator

SE(3)-Hyena operator consists of invariant and equivariant streams which are responsible for processing scalar and vector features respectively, as illustrated in Figure 2. This way, the model takes scalar and vector inputs, and outputs processed scalar and vector features.

Formally, $\Psi:\mathbb{R}^{3}\times\mathbb{R}^{S}\rightarrow\mathbb{R}^{3}\times\mathbb{
R}^{D}$ satisfies equivariance property:

$$ $\{L_{g}(\hat{\mathbf{x}}_{i}),\hat{\mathbf{f}}_{i}\}_{i=1}^{N}=\Psi\left(\{L_{ g}(\mathbf{x}_{i}),\mathbf{f}_{i}\}_{i=1}^{N}\right)$ $$

where $L_{g}:\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$ is a representation of a group action $g\in SE(3)$.
Thus, geometric tokens $\mathbf{x}_{i}$ transform accordingly with the group action while scalar tokens $\mathbf{f}_{i}$ remain invariant.

For both invariant and equivariant streams, we make the overall information flow similar to the information flow of a transformer with the key difference in how the global context is aggregated.
Akin to the transformer architecture, the input is firstly projected into keys, queries, and values.
Next, the global context is aggregated via long convolution and gating.
Finally, the global context is combined with values, and then projected via a feed-forward network.

### 3.2 SE(3)-Hyena modules

We design each layer to be equivariant with respect to transformations of geometric tokens, and invariant for scalar tokens.
This way, a model consisting of a composition of equivariant layers is also equivariant by induction (Weiler & Cesa, ).

##### Input projection

Similar to transformers, we firstly project input tokens to queries, keys and values. Because our SE(3)-Hyena operator consists of equivariant and invariant streams, we need to obtain queries, keys and values for both invariant and equivariant features. We define the projection layer
$\phi:\mathbb{R}^{3}\times\mathbb{R}^{S}\rightarrow(\mathbb{R}^{3})^{3}\times(
\mathbb{R}^{D})^{3}$ as $\mathbf{z}^{eqv}_{i},\mathbf{z}^{inv}_{i}=\phi\left(\mathbf{x}_{i},\mathbf{f}_
{i}\right)$ where for $i$-th geometric token $\mathbf{z}^{eqv}_{i}=[\mathbf{q}^{eqv}_{i},\mathbf{k}^{eqv}_{i},\mathbf{v}^{
eqv}_{i}]$ is equivariant query, key, and value, each represented by a vector in $\mathbb{R}^{3}$, and $\mathbf{z}^{inv}_{i}=[\mathbf{q}^{inv}_{i},\mathbf{k}^{inv}_{i},\mathbf{v}^{
inv}_{i}]$ represents invariant query, key and value of dimensions $\mathbb{R}^{D}$ for $i$-th scalar token. This way, the projection layer emits both scalar and vector query, key, value triplets while also allowing interaction between equivariant and invariant subspaces.

To simultaneously preserve equivariance for geometric tokens and invariance for scalar tokens, while allowing interaction between them, we adopt E(n)-equivariant Clifford MLP (Ruhe et al., ) as the input projection function. The inputs are firstly embedded into $CL(\mathbb{R}^{3},q)$ Clifford algebra, and then processed via a geometric linear layer, geometric product, and normalization layers. Finally, the output is projected back to scalar and vector features by grade-one and grade-two projections respectively.

##### Scalar long convolution

To allow global context aggregation for invariant scalar features, we rely on long convolution (Romero et al., ; Poli et al., ) between query and key tokens. We treat the queries as input signal projection, and the keys constitute a data-controlled implicit filter. Similar to Romero et al. (); Poli et al. (), we employ circular FFT-convolution to reduce the computational complexity. Let $\mathbf{q}^{inv}$ and $\mathbf{k}^{inv}$ be two sequences of length $N$ composed of sets of one-dimensional invariant queries $\{\mathbf{q}^{inv}_{i}\}_{i=1}^{N}$ and keys $\{\mathbf{k}^{inv}_{i}\}_{i=1}^{N}$ respectively. Then, the global context can be aggregated by FFT-convolution as:

$$ $\displaystyle\mathbf{q}^{inv}\circledast\mathbf{k}^{inv}$ $\displaystyle=\mathbf{F}^{H}\mathbf{\Lambda_{k}}\mathbf{F}\mathbf{q}^{inv}$ $\displaystyle=\mathbf{F}^{H}\text{diag}(\mathbf{F}\mathbf{k}^{inv})\mathbf{F} \mathbf{q}^{inv}$ (1) $$

where F is a discrete Fourier transform matrix, and $\text{diag}(\mathbf{F}\mathbf{k}^{inv})$ is a diagonal matrix containing Fourier transform of the kernel $\mathbf{k}^{inv}$.

In the case when query’s and key’s dimension $D>1$, the scalar FFT-convolution runs separately for each dimension, rendering computational complexity of $O(DNlog_{2}N)$ that is sub-quadratic in sequence length.

##### SE(3) vector long convolution

To allow global context aggregation for geometric tokens, we build equivariant vector long convolution.
While a scalar convolution relies on dot-products between scalar signals, vector convolution operates with vector cross products $\times$ between vector signals.
Formally, given a vector signal consisting of $N$ vector tokens $\mathbf{q}^{eqv}\in\mathbb{R}^{N\times 3}$ and a vector kernel $\mathbf{k}^{eqv}\in\mathbb{R}^{N\times 3}$, we define the vector long-convolution as:

$$ $\left(\mathbf{q}^{eqv}\circledast_{\times}\mathbf{k}^{eqv}\right)_{i}=\sum_{j= 1}^{N}\mathbf{q}^{eqv}_{i}\times\mathbf{k}^{eqv}_{j-i}$ (2) $$

The computational complexity of a naive implementation for the vector convolution is quadratic since both the signal and a kernel are of the full length. To reduce the computational complexity, we show how the vector convolution can be formulated as a series of scalar convolutions that can be efficiently carried out by the FFT. This is due to the fact that a cross product can be written element-wise through the series of scalar products as $\left(\mathbf{a}\times\mathbf{b}\right)[l]=\mathbf{\varepsilon}_{lhp}\textbf{a
}[h]\textbf{b}[p]$ where $\varepsilon$ is Levi-Civita symbol, and $\textbf{a}[h]$ denotes a projection onto $h$-th basis vector. Thus, the $l$-th component of the vector convolution in Equation 2 can be written element-wise as:

$$ $\displaystyle\left(\mathbf{q}^{eqv}\circledast_{\times}\mathbf{k}^{eqv}\right) _{i}[l]$ $\displaystyle=\varepsilon_{lhp}\sum_{j=1}^{N}\textbf{q}^{eqv}_{i}[h]\hskip 2.8 4526pt\textbf{k}^{eqv}_{j-i}[p]$ $\displaystyle=\varepsilon_{lhp}\left(\textbf{q}^{eqv}[h]\circledast\textbf{k}^ {eqv}[p]\right)_{i}$ (3) $$

Thus, we can obtain $l$-th component of a resulting vector signal via a scalar convolution over the $h$-th and $p$-th components of the sequences $\mathbf{q}^{eqv}$ and $\mathbf{k}^{eqv}$ respectively. Since the scalar convolution can be implemented with the FFT, decomposing the vector convolution to the series of scalar convolutions allows reducing its quadratic complexity to $O(Nlog_{2}N)$.

Since a cross product is already equivariant to rotations, the whole vector convolution is also equivariant to rotations provided the queries $\mathbf{q}^{eqv}$ and the keys $\mathbf{k}^{eqv}$ are rotated accordingly. The latter is guaranteed when the input projection function is equivariant. We can further obtain SE(3) equivariance by firstly centering (by subtracting the center of mass) and then uncentering with respect to translation.

##### Gating

Similar to transformers, we aim to enable data-controlled gating. Similar to the input projection layer, we employ Clifford MLP to obtain a gating mask while allowing interaction between invariant and equivariant subspaces. The output of the Clifford MLP $\gamma:\mathbb{R}^{3}\times\mathbb{R}^{D}\rightarrow(\mathbb{R})^{2}$ is two grade-one projected scalar features $m^{eqv}_{i},m^{inv}_{i}$ that are passed through a sigmoid function. The sigmoid outputs are multiplied element-wise with the scalar and vector tokens for the invariant and equivariant streams, respectively. In other words, given the input vector token $\mathbf{x}^{eqv}_{i}$, the gated token $\hat{\mathbf{x}}^{eqv}_{i}$ is computed as $\hat{\mathbf{x}}^{eqv}_{i}=\sigma(m^{eqv}_{i})\cdot\mathbf{x}^{eqv}_{i}$, with scalar tokens processed in a similar manner.

To align the information flow with that of transformers, the gating is applied on top of the long convolution between queries and keys. Finally, the gated tokens are integrated with value tokens $\mathbf{v}^{eqv}_{i}$ and $\mathbf{v}^{inv}_{i}$ using cross and element-wise products for scalar and vector tokens respectively.

By employing grade-one output projection in the Clifford MLP to obtain masking values, the gating mechanism remains E(3)-invariant, thus maintaining invariance for scalar tokens and preserving the equivariance of vector tokens.

##### Output projection

We add a residual connection between input tokens and the result from the gating layer, and pass it through the output equivariant Clifford MLP $\xi:\mathbb{R}^{3}\times\mathbb{R}^{D}\rightarrow\mathbb{R}^{3}\times\mathbb{R
}^{D}$ with grade-one and grade-two projections to extract processed scalar and vector tokens.

### 3.3 Algorithm

We detail the SE(3)-Hyena algorithm in Algorithm 1. The input projection, gating, and output projection operate in parallel for each token and thus of $O(N)$ complexity. The scalar and vector long convolutions are implemented with FFT and thus come with a complexity of $O(Nlog_{2}N)$.

Figure: Algorithm 1 SE(3)-Hyena forward pass

Figure: Figure 3: Equivariant associative recall task. An equivariant associative recall requires retrieving a vector token for a given vector query based on the context. The retrieval mechanism requires equivariance to rotation of tokens in a sequence. As standard associative recall serves to test the capability of models to learn global context, the equivariant associative recall task serves to test capability of models to learn global context with equivariance.
Refer to caption: x3.png

## 4 Experiments

##### Overview

We conduct experiments to assess the performance of the SE(3)-equivariant Hyena in modeling global geometric context. We start with the associative recall (induction heads) task from mechanistic interpretability (Elhage et al., ) that has become a standard for comparing efficiency of models in processing global context. Due to the equivariance requirement of geometric context, we have developed a vector-valued extension of the standard associative recall with equivariance to rotations. Next, we evaluate our model on the n-body dynamical system task, where the focus is on accurately predicting the positions of particles based on their initial velocities and coordinates while maintaining equivariance to rotations. Lastly, we compare the runtime and memory profiles of the equivariant Hyena against the equivariant transformer to highlight the computational efficiency of our method for long sequences.

##### Baselines

As a baseline, we evaluate the SE(3)-Hyena against the equivariant transformer, focusing specifically on global context modeling. To facilitate fair comparison, we aim to minimize architectural differences between the two models, altering only the global context mechanism. In the SE(3)-Hyena, we utilize scalar and vector long convolutions, whereas the transformer employs equivariant vector self-attention. The details on equivariant vector self-attention are in Appendix B. We also include comparisons with non-equivariant versions of both the Hyena and transformer models to underscore the benefits provided by equivariance.

Figure: Figure 4: Top row: The MSE between retrieved and target vectors for the fixed vocabulary associative recall task is plotted across various sequence lengths. Equivariant models effectively learn and generalize the underlying vocabulary across different orientations. Bottom row: The MSE for the random vocabulary associative recall task. The SE(3)-Hyena excels in learning the equivariant retrieval function, successfully associating target queries with their corresponding value vectors.
Refer to caption: x4.png

### 4.1 Equivariant associative recall

Associative recall (Olsson et al., ) is one of the standard mechanistic interpretability tasks used for quantifying the contextual learning capabilities of sequence models. In this task, a model is required to perform associative recall and copying; for instance, if a model previously encounters the bigram ”Harry Potter” in a sequence, it should accurately predict ”Potter” the next time ”Harry” appears, by referencing its past occurrence (Gu & Dao, ).

To adapt this for geometric contexts, we modify the standard associative recall to accommodate 3D vectors. In this version, the tokens within a bigram (key and value) relate to each other by a rotation matrix. The model processes a sequence of $N$ vector tokens concluding with a query token(^1^11key,value and query tokens here refer to associative recall task and different from keys,values and queries outputted by the input projection) and must predict a vector associated with this query seen earlier in the sequence, as illustrated in Figure 3. The task ensures that rotating the entire sequence affects only the orientation of the predicted vector, not the key-value relationship within a bigram, thereby making the task equivariant to rotations. The task complexity depends on the number of tokens in a sequence and on vocabulary size where the vocabulary items are unique key-value vector bigrams.

We propose two versions of the equivariant associative recall task: random and fixed vocabulary versions. In the random vocabulary version, sequences are sampled from a randomly generated vocabulary at each training iteration while validation and test sets are fixed. Thus, an only way for a model to solve this task is to learn an equivariant retrieval mechanism that can associate a given query with a corresponding value vector. In the fixed vocabulary version, the vocabulary is fixed and shared among training, test and validation sets, and the sequences are randomly rotated during the training. This way, the fixed version tests the model’s ability to learn underlying vocabulary and generalize it for various orientations.

#### 4.1.1 Implementation details

##### Dataset

We set the training data size to $2600$, and validation and test set size to $200$ sequences each. The vocabulary size is set to $4$ for both random and fixed versions. The vocabularies are generated as follows: one key-value bigram consists of two vectors with orientations sampled as random unit 3D vectors from an isotropic normal distribution, and with magnitudes randomly sampled from a uniform distribution in a range of $[1,\texttt{vocab_size}]$. We generate datasets with various sequence lengths from $2^{5}$ to $2^{8}$ tokens. When generating sequences from a vocabulary, the last token is a key that corresponds to a target value. We additionally constraint the generation so the target key-value pair is present in a sequence at least once.

##### Models

The vector tokens are used as the input to the equivariant branch, while for the invariant branch, we use positional encoding features (Vaswani et al., ) of dimension $16$ as the input. The SE(3)-Hyena model consists of $3$ SE(3)-Hyena operators followed by a mean pooling of equivariant features and one output equivariant MLP. In SE(3)-Hyena operator we set the hidden dimension to $16$ for invariant and to $128$ for equivariant streams. In the gating operator, we use a smaller dimension of $8$ for both equivariant and invariant features. This ends up in approximately $800k$ trainable parameters. Similarly, the equivariant transformer consists of $3$ SE(3)-equivariant vector self-attention blocks followed by a mean pooling and one equivariant MLP at the end. The hidden dimension is kept similar to SE(3)-Hyena which results in a nearly identical number of trainable parameters. Non-equivariant baselines consist of $3$ Hyena or standard self-attention blocks but with $3$ times larger hidden dimensions to balance the number of trainable parameters with equivariant models.

##### Training

We train all models for $300$ epochs with a batch size of $32$. We employ Adam optimizer (Kingma & Ba, ) with an initial learning rate of $0.001$ and cosine learning rate annealing (Loshchilov & Hutter, ). The weight decay is set to $0.00001$. Mean squared error is used as a loss function. For the fixed vocabulary experiment, we apply random rotation augmentation on sequences in the training batch, for the random vocabulary variant this is not necessary as sequences already appear in arbitrary orientations. Final models are selected based on the best validation loss.

#### 4.1.2 Results

The results for various sequence lengths are presented in Figure 4. We record the mean squared error between predicted and ground truth vectors as a performance measure.

For a fixed vocabulary variant, SE(3)-Hyena performs on par with the equivariant transformer across the whole range of sequence lengths. This demonstrates that both equivariant Hyena and transformer models can learn an underlying vocabulary and can generalize it for various orientations. We also observed that non-equivariant models were only able to learn the expectation across the training dataset.

For the random vocabulary variant, we observed that SE(3)-Hyena successfully learns the retrieval function to associate a target query with a corresponding target value vector. We observed better generalization for longer sequences which can be attributed to a higher frequency of target key-value bigram occurrence. Other models struggled to converge, with SE(3)-transformer only slightly outperforming non-equivariant models that only learn the expectation.

### 4.2 N-body problem

In the field of dynamical systems, comprehending the time evolution of point sets within geometric spaces is crucial for a variety of applications, such as control systems (Chua et al., ), model-based dynamics in reinforcement learning (Nagabandi et al., ), and simulations of physical systems (Watters et al., ). This task inherently requires equivariance, as any rotations and translations applied to the initial set of particles must be consistently represented throughout their entire trajectory.

We benchmark SE(3)-Hyena on forecasting trajectories in the dynamical system where the dynamics is governed by physical interaction between $5$ charged particles that carry a positive or negative charge. Particles have a position and a velocity in 3-dimensional space. The objective is to predict the future positions of particles given initial positions and velocities.

#### 4.2.1 Implementation details

##### Dataset

We utilize the standard n-body benchmark from Fuchs et al. (); Satorras et al. (). We sample $1000$ particles for training, $2000$ for validation, and $2000$ for testing. A training sample consists of equivariant features, i.e. initial positions $\mathbf{p}_{0}\in\mathbb{R}^{5\times 3}$ and velocities positions $\mathbf{v}_{0}\in\mathbb{R}^{5\times 3}$. Positions of the particles $\mathbf{p}_{1000}\in\mathbb{R}^{5\times 3}$ after $1.000$ timesteps are used as labels.

##### Models and training

The initial positions and velocities serve as the input to the equivariant branch, and one-hot encoded charges are used as the input to the invariant branch. The SE(3)-Hyena model consists of $2$ SE(3)-Hyena operators followed by one output equivariant MLP. The equivariant transformer baseline consists of $2$ vector self-attention blocks and one output equivariant MLP. For equivariant models, the invariant hidden dimension is set to $8$ and the equivariant hidden dimension to $16$. The gating in SE(3)-Hyena operator uses a hidden dimension of $8$ for both equivariant and invariant features. Non-equivariant baselines also consist of $2$ Hyena or standard self-attention blocks but with $3$ times larger invariant and equivariant hidden dimensions to match the number of trainable parameters with equivariant models. We also use the linear baseline model which is simply a linear motion equitation $\mathbf{p}_{t}=\mathbf{p}_{0}+t\cdot\mathbf{v}_{0}$. The models are trained for $300$ epochs with a batch size of $100$. We use Adam optimizer with a learning rate set to $0.0001$ and weight decay of $0.00001$. The mean squared error is used as the loss function.

#### 4.2.2 Results

The results are presented in Table 1. The SE(3)-Hyena model outperforms non-equivariant baselines and performs on par with the equivariant transformer. Interestingly, the non-equivariant Hyena also slightly outperformed the non-equivariant transformer. This may stem from Hyena’s structural similarities to SSMs (Poli et al., ) which are derived from the state-space representation of differential equation and hence are well-suited for modeling dynamical systems (Hinrichsen & Pritchard, ).

**Table 1: Mean squared error for the future position estimation in the N-body system experiment.**
| Method | MSE |  |
| --- | --- | --- |
| Linear | 0.0322 |  |
| Transformer | 0.0163 |  |
| Hyena | 0.0150 |  |
| SE(3)-Transformer | 0.0019 |  |
| SE(3)-Hyena | 0.0018 |  |

### 4.3 Runtime and memory benchmarks

We benchmark the runtime and memory consumption of a single layer SE(3)-Hyena against a single layer equivariant transformer when processing various sequence lengths. Similar to (Dao et al., ; Poli et al., ), we use random sequences for the runtime evaluation and we increase the sequence length until the SE(3)-Transformer runs out of memory. We record all runtimes on NVIDIA A10G GPU with CUDA version 12.2.

The comparison is reported in Figure 5, showing forward pass time in milliseconds and total GPU memory utilization in gigabytes. SE(3)-Hyena easily scales to longer sequences whereas SE(3)-Transformer is $3.5\times$ slower than our model for a sequence length of $20k$ tokens. Similarly to Nguyen et al. (), we observed Hyena to be slightly slower for shorter sequences. This can be attributed to expensive I/O between layers of the GPU memory hierarchy and can be addressed with the recent GPU-optimized versions of the FFT long convolution (Fu et al., ). Regarding memory usage, for the sequences length of $20k$ tokens the equivariant Hyena requires $18$ times less GPU memory than transformer for a forward pass. Moreover, we observed that when the equivariant transformer runs out of memory on $>20k$ tokens, our model supports up to $3.5M$ tokens on a single GPU allowing for $175$ times longer geometric context. This memory efficiency is attributed to the FFT long convolution that avoids materializing a quadratic self-attention matrix.

Figure: Figure 5: Top row: Forward runtime comparison. SE(3)-Hyena scales sub-quadratically and achieves a considerable speedup compared to SE(3)-Transformer when processing long sequences. Bottom row: Total GPU memory utilization for equivariant Hyena and transformer models.
Refer to caption: x6.png

## 5 Conclusions

We introduced the SE(3)-Hyena operator being, to the best of our knowledge, the first equivariant long-convolutional model with sub-quadratic complexity for global geometric context. Through experiments on the dynamical system and novel equivariant associative recall task, we demonstrated that the equivariant long-convolutional model can perform competitively to the equivariant self-attention while requiring a fraction of the computational and memory cost of transformers for long context modeling. Our scalable equivariant model efficiently captures the global context, highlighting its potential for a multitude of future applications in various domains.

##### Limitations and Future work

This work introduces a novel approach for equivariant modeling of global geometric context at scale, with initial experiments designed to validate the fundamental principles of our method. While these experiments confirm the key advantages of our approach, they represent the first step of a comprehensive experimental analysis that is necessary to uncover the model’s capabilities across a wider range of real-world tasks. On a technical side, an interesting direction for future improvement is to adapt the vector convolution to function across arbitrary dimensions as it currently relies on a cross product, which is only feasible in $3$ and $7$ dimensions (Massey, ). Also, our method relies on the FFT convolution that is not permutation equivariant. While this is not a significant limitation in domains such as molecular systems, where canonical ordering is typically available (Jochum & Gasteiger, ), it becomes critical in fields like point cloud processing, where establishing a canonical order is challenging. Enhancing the long-convolutional framework to incorporate permutation equivariance or learning this geometric constraint from the data (Moskalev et al., ) could offer substantial advantages in these areas.

## Appendix A Equivariant vector long convolution

We provide Pytorch implementation for the rotation-equivariant (without centering) vector long convolution in Code LABEL:lst:pytorch_conv.

Figure: Code 1: Pytorch implementation of the equivariant vector long convolution.

## Appendix B Equivariant vector self-attention

Since SE(3)-Hyena utilizes long convolutions based on cross products, we design cross product equivariant self-attention to align the global geometric context aggregation for Hyena and transformer models. Firstly, the self-attention tensor is built from all pairs of cross products between queries and keys. Then, the $L_{2}$ norm matrix is extracted and softmax is run row-wise on top of this matrix. The resulting softmax matrix serves as a selection mechanism to select vectors from the cross product tensor. Finally, the values are cross-multiplied with the self-attention tensor. This yields a set of processed values modulated with global self-attention information.

##### Cross product vector self-attention

Consider sequences of $N$ vector queries, keys, and values denoted as $\mathbf{q},\mathbf{k},\mathbf{v}\in\mathbb{R}^{N\times 3}$. We construct a query-key cross product tensor $\mathbf{C}\in\mathbb{R}^{(N\times N)\times 3}$ where each element $\mathbf{C}_{ij}=\mathbf{q}_{i}\times\mathbf{k}_{j}$, or using Levi-Civita notation as in Eq. 3.2, $\mathbf{C}_{ij}[l]=\varepsilon_{lhp}\mathbf{q}_{i}[h]\mathbf{k}_{j}[p]$. To integrate a softmax selection mechanism, as in standard self-attention, we first compute a matrix $\eta(\mathbf{C})\in\mathbb{R}^{N\times N}$ containing the $L_{2}$ norms of all cross products, specifically $\eta(\mathbf{C})_{ij}=\|\mathbf{q}_{i}\times\mathbf{k}_{j}\|_{2}$. Applying softmax to $\eta(\mathbf{C})$ then determines the vector pairs to select from the cross product tensor. Lastly, the values $\mathbf{v}$ are cross-multiplied with the softmax-filtered cross product tensor. Overall, the equivariant vector self-attention reads as:

$$ $\displaystyle\mathbf{S}=\texttt{softmax}(\frac{1}{\sqrt{N}}\eta(\mathbf{C})) \odot\mathbf{C}$ (4) $\displaystyle\mathbf{u}_{i}=\frac{1}{N}\sum_{j=1}^{N}\mathbf{S}_{ij}\times \mathbf{v}_{j}$ (5) $$

where the softmax is applied row-wise, and $\odot$ stands for element-wise product. Consequently, $\mathbf{S}\in\mathbb{R}^{(N\times N)\times 3}$ represents a tensor that encapsulates a soft selection of cross products between $\mathbf{q}_{i}$ and $\mathbf{k}_{j}$. Initially, we considered using just $\texttt{softmax}(\frac{1}{\sqrt{N}}\eta(\mathbf{C}))$ as a self-attention matrix, but early experiments indicated that the method outlined in Eq. 4 yields better results. Additionally, we found that normalizing the sum by $1/N$ in Eq. 5 further improves convergence.

Since the tensor $\mathbf{C}$ is constructed using cross products, it naturally maintains equivariance to rotations of queries and keys. Furthermore, the softmax is applied to the $L_{2}$ norms of these cross products, which makes it invariant to rotations. Consequently, the self-attention tensor $\mathbf{S}$ is a product of rotation-invariant scalar and rotation-equivariant vector quantities, rendering it rotation-equivariant. The Eq. 5 further preserves rotation-equivariance due to the inherent equivariance of the cross product. Equivariance to translations can be achieved by initially centering the data (subtracting the center of mass) and then re-centering the resulting tokens.

We provide Pytorch implementation for the rotation-equivariant (without centering) vector self-attention in Code LABEL:lst:pytorch_sa.

Figure: Code 2: Pytorch implementation of equivariant vector self-attention.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See Weakness.
soundness: 3
strengths: The idea of using global information to improve the model is very natural, and the convolution simplification of the cross product is very elegant.
summary: The paper introduces the SE(3)-Hyena operator, the first equivariant long-convolutional model with sub-quadratic complexity for global geometric context.  Importantly, authors claim their framework is flexible and can accommodate any equivariant network as the projection function.
weaknesses: > **W1. Lack of discussion on related work.**

There are many works that use global features to improve equivariance. Although this paper's work is obviously different from them, it is recommended to add a discussion on these works (e.g. FastEGNN [a], Neural P^3M [b]).

> **W2. The motivation for operator design is unclear.**

Why is the motivation for using the cross product not well explained? As we all know, the cross product is the Hodge star dual of the outer product. Can it be explained from this perspective? In addition, the introduction of the cross product actually produces pseudovectors, which also leads to the fact that this paper is only SE(3) equivariant rather than E(3) equivariant. Is such an introduction really reasonable?

> **W3. Results on N-body lack the latest baseline.**

Some of the latest results are not shown (e.g. 0.0043 of SEGNN [c], 0.0039 of CGENN [d]). Compared with these works, the results of this work seem to be insufficient.

> **W4. The significance of associative recall experiments is unclear.**

The current experiment cannot illustrate the "contextual learning capabilities of sequence models" that the authors want to claim. First, the model lacks more baselines (e.g. LEFTNet [e], MACE [f], EquiformerV2 [g], SO3krates [h]). It is more like explaining that in this setting, equivariant models are better than models without built-in symmetry priors. Secondly, this experiment lacks practical application significance. I can't seem to find a corresponding task in real life. I hope the authors can give further explanation.

> **W5. Lack of baseline in RNA dataset.** 

The baseline is also missing, and the baseline mentioned in W1 and W4 should be supplemented.

> **W6. Lack of expansion on other models.**

Authors claim their framework is flexible and can accommodate any equivariant network as the projection function. Is it possible to extend several common models (e.g. SchNet [i], EGNN [j], MACE [f], HEGNN [k])?

> **W7. Others (Some typos)**

- Line 144: where is the function $\Psi$, or the hat $\hat{\mathbf{x}_i}, \hat{\mathbf{f}_i}$ are the outputs?
- Line 231: in calculation of $\alpha_3$, it should be $\mathbf{r}_1^\top\mathbf{r}_2$
- Line 242: feature tuples $f_i$, LaTeX misses underscore
- Line 817: "hiddent dimension" should be "hidden dimension"
- Others: the logarithmic symbol should be $\log$ instead of $log$, and why is there a base sometimes without and sometimes with 2?

[a] Improving Equivariant Graph Neural Networks on Large Geometric Graphs via Virtual Nodes Learning

[b] Neural P3M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs

[c] Geometric and Physical Quantities Improve E(3) Equivariant Message Passing

[d] Clifford Group Equivariant Neural Networks

[e] A new perspective on building efficient and expressive 3D equivariant graph neural networks

[f] Mace: Higher order equivariant message passing neural networks for fast and accurate force fields

[g] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

[h] A Euclidean transformer for fast and stable machine learned force fields

[i] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions

[j] E(n) Equivariant Graph Neural Networks

[k] Are High-Degree Representations Really Unnecessary in Equivariant Graph Neural Networks?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Since the work focuses on "scalable" equivariant learning, could the authors provide results on larger-scale datasets to further demonstrate the model's effectiveness?
soundness: 2
strengths: Applying deep learning to problems that involve modeling geometric context, as done in this paper, is a valuable direction in the field. Performance improvements in this area often depend on architectural advances, which is also good to see.
summary: The paper introduces the SE(3)-Hyena operator to capture global geometric information while preserving equivariant constraints. It aims to address the computational limitations of existing methods, such as self-attention and local processing techniques. The proposed model is evaluated on dynamical system modeling and RNA property prediction tasks, and the authors introduce an "equivariant associative recall" task to assess contextual learning abilities.
weaknesses: - The experiments in the paper are not comprehensive enough to clearly demonstrate the advantages of the proposed method over existing ones. For example, baselines used in the paper -- SchNet, EGNN, and SE(3)-Transformer -- have been evaluated on the QM9 dataset in their original papers. It would be more convincing if the authors included results on QM9 as well.

- Several state-of-the-art baselines for dynamical system modeling are missing, making the performance of the proposed model not convincing enough. Examples include SEGNN [1], SAKE [2], SEGNO [3], and GeoMFormer [4].

[1] Geometric and Physical Quantities Improve E(3) Equivariant Message Passing. ICLR 2022.

[2] Spatial Attention Kinetic Networks with E(n)-Equivariance. ICLR 2023.

[3] Improving Generalization in Equivariant Graph Neural Networks with Physical Inductive Biases. ICLR 2024.

[4] GeoMFormer: A General Architecture for Geometric Molecular Representation Learning. ICML 2024.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: > Writing
1. Line 19: Give one or two sentences about long convolutions. Otherwise, it is hard to tell the difference from message passing.
2. Line 22: Give the name of datasets you tested so that people can better judge the scale of experiments at the beginning.
3. Line 32 -- 38: Give one or two sentences about equivariance.
4. Figure 1: I cannot tell any difference from typical self-attention, and thus the complexity O(N log N) is unclear.
5. Figure 2: The figure is too simplified without giving any detail. You should include more details about projection, geometric long convolution and gating and make them consistent with the context.
6. Line 144 -- 145: Add equation number. Also I think $\hat{x}_i$ is not defined here.
7. Line 169 -- 170: typo: "To this end,"
8. Line 191 -- 192: I don't think $F^H$ is defined.
9. Line 237 -- 239: Spherical harmonics with type-0 and type-1 vectors are the same as the representation in Cartesian space.
10. Section 3 should be reflected in figures.

---

> Experiments

1. Figure 5: Why is SE(3)-Hyena better than SE(3)-Transformer? I think some explanations would be great.
2. Line 463: Why using just 2 layers? In such small models, global context probably does not exist, resulting a potentially unfair comparison.
3. Line 473 -- 474: You can train a smaller SE(3)-Transformer to fit on GPU?
4. Please compare the proposed network with other previous works on QM9 and MD17 datasets.

---

> Reproducibility

Please submit the code when the work is under review instead of releasing upon acceptance, especially when parts of the paper are unclear and experiments are not well verified.

--- 

> Question

1. Is the memory complexity still $O(N^2)$? If not, please give a very short introduction to this in the paper (either in the main text or appendix).
2. Line 294 -- Line 302: The weighting factor depends on the index of input tokens. How is the permutation of tokens handled?
3. Line 269: You still use typical message passing (EGNN layer) to encode the context. Would this be the memory bottleneck instead of SE(3)-Hyena operator? If yes, I think other equivariant networks share the same memory and compute bottleneck as this work.
soundness: 1
strengths: 1. Using a sub-quadratic operator for global context in 3D atomistic modeling can be potentially a good idea.
summary: The authors proposed SE(3)-Hyena operator for modeling global interaction of atomistic systems. Unlike naive SE(3)-equivariant attention, the SE(3)-equivariant operator does not require quadratic computational complexity due to the usage of FFT. The reduced complexity is well-benchmarked on toy examples.
weaknesses: 1. The writing should be greatly improved. See Questions below for more details.
2. Experiments in the paper are very limited. The first two experiments basically tell little about how effective a network architecture can be. The third one does not show the benefit of SE(3)-Hyena except that the proposed method takes less memory.
3. Lack of comparisons to previous works on other better-benchmarked datasets such as MD17, QM9 and so on. Overall, how effective modeling global context is remains unclear.
4. The exact proposed architecture is missing or very hard to understand. I think a better visualization covering all the details in a high-level manner can be helpful.
5. The proposed architecture only uses type-0 and type-1 vectors and can only use 1 channel for type-1 vectors. These significantly limit applying the network to slightly larger datasets.