Geometric Hyena Networks for Large-scale Equivariant Learning

Computational and performance benefits:

We appreciate the reviewer acknowledging the runtime and memory benefits of our method, as computational efficiency is the main purpose of our work.

Second, we believe there may be a misunderstanding regarding performance improvements. G-Hyena significantly outperforms the G-Transformer model—with non-overlapping error bars on Tc-Ribo in both the backbone and all-atom settings (-7% error reduction in both cases). It also shows substantial gains on the Protein-MD task in both backbone and all-atom scenarios (-26% and -36% error reduction, respectively). Since the primary architectural difference between G-Hyena and G-Transformer is the geometric long convolution, we attribute these performance improvements to that component. In addition, the associative recall experiment (Figure 3, bottom left) shows that G-Hyena scales more favorably with respect to the hidden dimension compared to G-Transformer. This improved scaling behavior is practically relevant, as it potentially enables more efficient model size reduction and further lowers memory usage. Taken together, we consider both the computational and performance advantages of our model to be significant in the context of the conducted experiments.

Non-equivariant state space model:

We agree that the extended comparison with the non-equivariant Hyena model will further strengthen the experiments. In addition to Hyena in the associative recall experiments, we test a standard Hyena model trained with data augmentation on Open Vaccine, Ribonanza-2k, Tc-ribo, and Protein-MD datasets.

Results suggest that standard non-equivariant Hyena trained with extensive data augmentation still significantly lags behind properly equivariant Geometric Hyena (with G-Hyena performance gain ranging from -3% error reduction on Tc-Ribo all-atom to -53% error reduction Ribonanza-2k backbone). Furthermore, the non-equivariant Hyena model struggles to generalize on the Protein-MD dataset where it underperforms significantly compared to equivariant methods. The results will be included in the revised version of the paper.

Presentation of the paper and clarifications:

• We will refine the text of the paper to correct grammatical errors and typos as soon as openreview allows revision. We will further streamline the writing in "invariant-equivariant subspace interaction" section and remove redundancies in the beginnings of sections 4.1 and 4.2.
• In Eq. (1), $\hat{\mathbf{x}}$ and $\hat{\mathbf{f}}$ refer to the output of Geometric Hyena.
• Indeed, the more precise domain and co-domain formulation would be $\Psi: (\mathbb{R}^{3} \times \mathbb{R}^{d})^{N} \rightarrow (\mathbb{R}^{3} \times \mathbb{R}^{d})^{N}$ as our model operates on a sequence of $N$ scalar vector feature tuples $(x_1, f_1), (x_2, f_2), \dots, (x_N, f_N)$ with each feature tuple being $\mathbb{R}^{3} \times \mathbb{R}^{d}$.
• The notation $F^H$ in Eq. (2) refers to the Hermitian transpose of the FFT matrix.
• The gating function indeed acts per token.
• We appreciate the reviewer's focus on clarity and precision of the presentation of the paper! We will fix typos, add clarification and adjust notation accordingly in the revised version.

Related work section:

Due to space constraints, a comprehensive review of related work could not be included in the main text. Instead, we focused on discussing the most essential prior work in the introduction, while providing a more detailed review in Supplementary Section A. This structure is not uncommon with multiple top-tier conference ML papers (e.g., [1,2,3]) adopting a similar format. If the reviewer prefers an explicit related work section in the main text, we are happy to include a shortened version there while retaining the extended related work in the supplementary material.

We thank the reviewer for their constructive feedback! We believe we have addressed all questions and comments. In light of this, we kindly ask the reviewer to consider raising the score.

[1] Gu et al. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. COLM 2024.
[2] Yi et al. Bridge the Modality and Capability Gaps in Vision-Language Model Selection. NeurIPS 2024.
[3] Razin et a. Implicit Regularization in Deep Learning May Not Be Explainable by Norms. NeurIPS 2020.

Application focus of Geometric Hyena:

Regarding the generalizability of our method to arbitrary geometric graphs, we refer the reviewer to the Limitations section (lines 424–439R), where we explicitly discuss the limitations of G-Hyena for point clouds and emphasize that our method is best suited for problems where a canonical ordering can be established. Also, note that Reviewer YW9h acknowledges in their review that we have noted this limitation explicitly. With this, we disagree that we overclaimed the contribution of our method to arbitrary geometric graphs. That said, we will revise the text to make our target application domain more explicit. Specifically, we will update the phrasing to: "The Geometric Hyena is designed for tasks that require modeling invariant and equivariant features in geometric graphs where canonical order can be established". This will be clarified in the revised version of the paper.

Second, we would like to point out that the class of geometric graphs with canonical ordering is quite broad and includes large bio-molecules such as proteins, antibodies, peptides, enzymes, and RNA - systems that are highly relevant to biosciences and drug discovery. In this context, we also remind the reviewer that our submission is under the Application-Driven Machine Learning track, which explicitly includes biosciences as a focus area in ICML call for papers (icml.cc/Conferences/2025/CallForPapers). In addition, large bio-molecules represent a well-established and rapidly growing application area by its own in the machine learning community, with numerous top-tier papers focused on this domain [1,2,3,4,5]. Viewed in this light, the reviewer’s remark that our work may “have a significant contribution to RNA-Seq analysis” also reflects its broader relevance to machine learning for biosciences. Finally, we note that our method applies beyond RNA, as demonstrated in the protein molecular dynamics task in Section 4.5.

Equiformer and VNT:

Our method is designed for large bio-molecules, such as RNA or proteins, which can consist of hundreds or thousands of atoms. In contrast, Equiformer and VNT are developed for small geometric systems. For example, the Ribonanza-2k dataset we use has bio-molecules with up to 11300 atoms, whereas the MD17 dataset used to evaluate Equiformer has molecules with only up to 25 atoms. The high dimensionality of our data introduces unique computational challenges that our method is specifically designed to address that do not arise in low-dimensional settings and that existing equivariant models are not equipped to handle at scale. As a result, it is also infeasible to directly reuse the hyperparameters from Equiformer paper (while VNT paper does not report exact optimal hyperparameters). Still following reviewer's suggestion, we tried running Equiformer with the smallest default configuration from the original paper and it runs out of memory on our data even with a batch size of 1 highlighting further the importance of our model for large-scale bio-molecular graphs. Finally, our initial reason for running the models to match the parameter count of G-Hyena was to ensure that we could fairly compare the benefits of these models by isolating the impact of model scale on performance which is a standard practice [6,7].

Associative recall:

We appreciate the reviewer’s attention to the clarity of presentation of the geometric associative recall task. To aid understanding, we illustrated how “each bigram (key and value) is connected by a rotation matrix” in Figure 4 of the supplementary material. To further improve clarity, we propose updating the caption of Figure 4 to: “A geometric sequence consists of key and value vector tokens, where consecutive key-value pairs form bigrams. Geometric associative recall requires retrieving the value vector corresponding to a query, where the query matches one of the keys in the sequence.”

Vector long convolution:

Vector long convolution is a novel component introduced in our work that we demonstrate how to implement efficiently in sub-quadratic time. It is not a standard definition.

We thank the reviewer for their constructive feedback! We believe we have addressed all questions and comments. In light of this, we kindly ask the reviewer to consider raising the score.

[1] Groth et al. Kermut: Composite kernel regression for protein variant effects. ICLR 2024.
[2] Jing et al. AlphaFold Meets Flow Matching for Generating Protein Ensembles. ICML 2024.
[3] Nori et al. RNAFlow: RNA Structure & Sequence Design via Inverse Folding-Based Flow Matching. ICML 2024.
[4] Gong et al. Evolution-Inspired Loss Functions for Protein Representation Learning. ICML 2024.
[5] Tan et al., Deciphering RNA Secondary Structure Prediction: A Probabilistic K-Rook Matching Perspective. ICML 2024.
[6] He et al. Deep Residual Learning for Image Recognition. CVPR 2015.
[7] Tan et al. Rethinking Model Scaling for Convolutional Neural Networks. ICML 2019.

We appreciate the reviewer's positive feedback on contributions of our paper, highlighting the novelty ("the extension is non-trivial and the method is clearly novel") of our method and its experimental validation ("the authors consider a strong set of experiments to test their methods").

Non-equivariant baselines:

We agree that the extended comparison with non-equivariant data augmentation baselines will further strengthen the experiments. We train and test Hyena and Transformer models with data augmentation on Open Vaccine, Ribonanza-2k, Tc-ribo and Protein-MD datasets.

Results suggest that on Open Vaccine and Ribonanza-2k data, non-equivariant models lag significantly behind Geometric Hyena (with G-Hyena delivering up to 34% error reduction), and on Tc-Ribo all-atom Transformer performs on par with. Yet, non-equivariant models struggle to generalize on the Protein-MD task where non-equivariant Hyena and Transformer perform poorly compared to Geometric Hyena.

Local methods with sequence-structured priors:

Extending local methods with sequence-structured prior presents a non-trivial direction for future work that extends beyond the scope of our paper. However, to foster future research, we conducted an extra experiment where we employed a Hyena layer to extract sequence features and then ran EGNN on top (Seq-EGNN). Alternatively, we can first run EGNN to extract geometric features and then run a sequence model on top (EGNN-Seq).

Our initial results suggest that adding sequence prior has little effect when EGNN is applied on top of sequence features. However, when the sequence model is applied on top of EGNN features, we observe slight improvement, suggesting potential for further research in this direction.

Future benchmarking on materials:

In our work, we focus specifically on large bio-molecules. We agree that extending experimental comparison to a broader domain, including materials, presents an important direction for future work, and we plan to do so in the future. We would appreciate if the reviewer could point us to relevant material benchmarks that include large molecules for our future work.

Higher-order Geometric Hyena:

Recent works [1,2] observed that higher-order representations provide diminishing performance improvement while significantly increasing computational requirements and memory footprint. This is also supported by the excellent performance of scalarization-based equivariant GNNs [3]. Based on this evidence, we decided not to proceed with a higher-order version of Geometric Hyena since our focus is on computational and memory efficiency.

Other comments and questions:

• Indeed, our model operates on a sequence of $N$ scalar vector feature tuples $(x_1, f_1), (x_2, f_2), \dots, (x_N, f_N)$ rather than unordered set, we thank the reviewer for pointing this out. With this, a more precise formulation of the domain and co-domain would be $\Psi: (\mathbb{R}^{3} \times \mathbb{R}^{d})^{N} \rightarrow (\mathbb{R}^{3} \times \mathbb{R}^{d})^{N}$.
• Projection terminology is used for alignment with the terminology in the Transformer and Hyena papers.
• In our implementation, we use standard discrete FFT with circular convolution. There is no explicit frequency truncation and all available frequencies are used up to the Nyquist limit.
• 190-192L: Yes, similar to the scalarization trick [3], this results in a lossy compression of information.
• In Protein-MD each protein has 855 and 3341 atoms in backbone and all-atom version. Atom identities are used as feature for all methods.
• All clarifications, discussion, and new results will be added to the revised version of the paper.

[1] Brandstetter et al. Geometric and Physical Quantities Improve E(3) Equivariant Message Passing. ICLR 2022.
[2] Wand et al. Rethinking the Benefits of Steerable Features in 3D Equivariant Graph Neural Networks. ICLR 2024.
[3] Satorras et al. E(n) Equivariant Graph Neural Networks. ICML 2021.