Paper: Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products (ICLR-2024)

Reviewer comment (Writing):
One point of critique is related to the treatment of Theorem 3.1 in the manuscript. The text uses a range of specialized terms—such as "spherical tensor operator," "reduced matrix element," and "total angular momentum"—and introduces various notations that may not be readily accessible to readers unfamiliar with quantum mechanics. To enhance clarity and comprehension, I recommend that the authors elaborate on these concepts, particularly in their exposition of the Wigner-Eckart theorem.

Author response:
Thank you for supporting our work. We reply to all of your questions below, and we have carefully revised and updated our paper to include these new changes:

> **Elaborate on concepts in Theorem 3.1**

Thank you for this suggestion. We have revised our paper to thoroughly elaborate on the mentioned notations and concepts in Section 3.1, for improved clarity and readability. In detail, we added a new subsection (Appendix A.7, Page 19-26), which is composed of the following contents:

- Appendix A.7 Angular Momentum, Tensor Operators and Wigner-Eckart theorem
  - A.7.1 Preliminary
    - Hilbert space and Dirac notation ("bra-ket" notation)
    - Operators
    - Eigenkets, basis kets and matrix elements of operators
  - A.7.2 A Brief Overview of Angular Momentum in Quantum Mechanics
    - Momentum Operator and Position Operator
    - Angular Momentum Operator
    - Eigenstates and eigenvalues of Angular Momentum Operators
    - Addition of Angular Momentum Operators
  - A.7.3 Tensor Operators and Wigner-Eckart theorem
    - Rotation Operators
    - Vector Operators and Tensor Operators
    - Wigner-Eckart theorem

For completeness, we comprehensively introduce necessary mathematical concepts in quantum mechanics, and elaborate step-by-step on their exposition in the context of Section 3.1 in an easy-to-follow and self-consistent manner, especially for general audiences who are unfamiliar with quantum mechanics. In Section 3.1, we add a reference to Appendix A.7 for interested readers, in case they are curious about the details of the notations and concepts. This should allow our readers to have more context about the concepts in Section 3.1, and also provide information on the relationship between Clebsch-Gordan coefficients and Gaunt coefficients.

Reviewer comment (Experiments):
Another inquiry concerns the necessity of using high-degree irreps. In most E3-NN architectures, irreps with degrees no greater than L=2 are commonly employed. Consequently, I question whether the computational overhead incurred through the FFT in implementing the proposed tensor product is actually more costly than directly using the CG product for low-degree irreps.

Author response:
> **Necessity of using higher-degree irreps**

We would like to clarify the following points and the motivation for using irreps of higher degrees:

1) The reason that most E3NN architectures commonly use irreps with degrees no greater than $L=2$ lies in the *significant computational cost*, rather than it being unnecessary to use. In the literature of E3NN architectures, the $\mathcal{O}(L^6)$ computational complexity of computing the Tensor Products of irreps has been widely acknowledged as a major obstacle. In Section 1 of our paper, we refer the readers to see comments in more than 10 related works over almost three years, discussing such limitations, which indeed reflects *the difficulty of tackling this problem*. Due to this, most existing E3NN architectures *are substantially limited* and cannot use irreps with higher degrees in practice, not because there is no need. 

2) Using irreps of higher degrees is indeed *necessary*. Even though the computational costs limit the usage of irreps of higher degrees, there are multiple examples showing that higher degree irreps can improve performance. Recent advancements in E(3)-equivariant networks suggest that the empirical performance of state-of-the-art models based on the Tensor Products of irreps can be consistently boosted by using higher degrees of irreps [1,2]. Additionally, specific tasks require the model to use higher degrees of irreps. For example, the DFT Hamiltonian prediction task [3,4] requires the model to predict a molecular system's DFT Hamiltonian matrix, whose shape is determined by the number of orbitals used to describe the system (proportional to the number of electrons). To construct the predicted matrix, E(3)-equivariant networks based on irreps are utilized, and the shape of Hamiltonian matrices exactly determines the largest degree of irreps that are necessary to use. Therefore, it is necessary, and has great significance, to tackle this bottleneck efficiency issue so that we can appropriately use irreps of higher degrees.

3) Our approach still remains very efficient when we use lower degree irreps. We provide detailed statistics of the three settings from Figure 1:

  - Equivariant Feature Interaction: 7.9x and 17.4x times acceleration for $L=2$ and $L=3$ respectively (compared to the e3nn implementation)
  - Equivariant Convolution: 4.1x and 7.0x times acceleration for $L=2$ and $L=3$ respectively (compared to the eSCN implementation)
  - Equivariant Many-body Interaction: 193.6x and 288.7x times acceleration for $L=2$ and $L=3$ respectively (compared to the e3nn implementation)
  

[1] Passaro, S., & Zitnick, C. L. (2023). Reducing SO (3) Convolutions to SO (2) for Efficient Equivariant GNNs. *ICML 2023*.

[2] Liao, Y. L., Wood, B., Das, A., & Smidt, T. (2023). EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. *arXiv preprint arXiv:2306.12059*.

[3] Unke, O., Bogojeski, M., Gastegger, M., Geiger, M., Smidt, T., & Müller, K. R. (2021). SE (3)-equivariant prediction of molecular wavefunctions and electronic densities. *Advances in Neural Information Processing Systems*, *34*, 14434-14447.

[4] Yu, H., Xu, Z., Qian, X., Qian, X., & Ji, S. (2023). Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian. *ICML 2023*.

We thank you again for your efforts in reviewing our paper, and we have carefully revised our paper and replied to each of your comments. We look forward to your re-evaluation of our submission based on our responses and updated results.

Reviewer comment (Writing):
Could you give me more explanation about the *braket* notations used in Theorem 3.1? It seems not clearly introduced in the paper.

Author response:
> **Provide more explanations on the bracket notations used in Theorem 3.1**

We have updated the paper to provide a comprehensive overview in Appendix A.7 on Dirac notation (bra-ket notation) and additional quantum mechanics concepts. We provide additional explanation here and summarize what we added: 

In physics, bra-ket notation is used to describe objects in a Hilbert space $\mathcal{H}$. It is a convenient way to write out inner product operations $(\cdot,\cdot):\mathcal{H}\times\mathcal{H}\to\mathbb{C}$. Vectors in the Hilbert space $\mathcal{H}$ are written as kets: $|\alpha\rangle\in\mathcal{H}$, and each ket has the one-to-one corresponding bra denoted as $\langle\alpha|$, with which the inner product of two kets can be written as $(|\alpha\rangle,|\beta\rangle)\equiv\langle\alpha|\beta\rangle$.

To describe the mapping between kets in a Hilbert space, we use the concept of operators, i.e., $\hat{O}:\mathcal{H}\to\mathcal{H}$. We define the eigenkets (also called eigenstates) $|\alpha\rangle$ of an operator as $\hat{A}$ if $\hat{A}|\alpha\rangle=a|\alpha$, where the scalar $a$ is the corresponding eigenvalue. Each ket in a Hilbert space can also be expressed as a linear combination of the basis kets. For any operator $\hat{A}$ and fixed basis kets $|\alpha_i\rangle$, we can calculate its matrix elements via $\langle\alpha_i|\hat{A}|\alpha_j\rangle$.

Building on the above notation, we provide a detailed explanation on the bra-ket notation used in Theorem 3.1:

  - $|l,m\rangle$ denote the eigenstates of angular momentum operators. 

    The angular momentum operators involve two operators: the momentum operator $\hat{\mathbf{p}}=-i\hbar\left(\frac{\partial}{\partial x},\frac{\partial}{\partial y},\frac{\partial}{\partial z}\right)$ and the position operator $\hat{\mathbf{x}}=\left(\hat{x},\hat{y},\hat{z}\right)$, where the position operator $\hat{x}$ (also $\hat{y},\hat{z}$) satisfies $\hat{x}f(x)\equiv xf(x)$. The angular momentum operator is defined as $\hat{\mathbf{L}}=(\hat{L}_x,\hat{L}_y,\hat{L}_z)$, where $\hat{L}_x=\hat{y} \hat{p}_z-\hat{z} \hat{p}_y,\hat{L}_y=\hat{z} \hat{p}_x-\hat{x} \hat{p}_z,\hat{L}_z=\hat{x} \hat{p}_y-\hat{y} \hat{p}_x$. The square of $\hat{\mathbf{L}}$ is $\hat{\mathbf{L}}^2=\hat{L}_x^2+\hat{L}_y^2+\hat{L}_z^2$. The kets $|l,m\rangle$ are simultaneously the solutions of the following equations (1) $\hat{L}_z|l,m\rangle=\hbar m|l,m\rangle$; (2) $\hat{\mathbf{L}}^2|l,m\rangle=\hbar l(l+1)|l,m\rangle$. Thus, $|l,m\rangle$ denote the eigenstates of the angular momentum operators. Given $l\in\mathbb{N}$, we have $-l\leq m\leq l$, and the set of $|l,m\rangle$ forms a basis for the $2l+1$-dimensional Hilbert space.

    The explicit form of kets $|l,m\rangle$ is exactly the spherical harmonics, $Y^{(l)}_m$, introduced in Section 2 in our paper. In quantum mechanics, it is conventional to use the ket notation $|l,m\rangle$, and we can obtain the function value of $Y^{(l)}_m$ with input $(\theta,\phi)$ via $\langle\theta,\phi|l,m\rangle\equiv Y^{(l)}_m(\theta,\phi)$.

  - $\langle l,m|$ denotes the bra corresponding to the eigenstates $|l,m\rangle$.

  - $\langle J,M|T^{(j)}_m|l,n\rangle$ denotes the matrix elements of spherical tensor operator $T^{(j)}_m$. This is when the basis corresponds to eigenkets of the angular momentum operators. As previously introduced, by fixing $J,l$, there are $2J+1$ and $2l+1$ basis kets (specified as $|J,M\rangle$ and $|l,n\rangle$) for the $2J+1$-dimensional and $2l+1$-dimensional Hilbert spaces respectively. By definition, spherical tensor operators of rank $j$ are families $\mathbf{T}^{(j)}=(T^{(j)}\_{-j},...,T^{(j)}_j)^\top$ of $2j+1$ operators $T^{j}_m$ that transform by rotations, like spherical harmonics. Therefore, there are $(2J+1)\times(2j+1)\times(2l+1)$ matrix elements of the spherical tensor operator $\mathbf{T}^{(j)}$ under the $|J,M\rangle$ and $|l,n\rangle$ basis.

  - $\langle J || \mathbf{T}^{(j)}||l\rangle$ denotes the reduced matrix elements, which are constants that only depend on the value of $J,j,l$.

Reviewer comment (Experiments):
How to determine in practice the truncation, ie. the degree L, in the 2D Fourier expansion (above Eq. (5))?

Author response:
> **How to determine, in practice, the truncation in the 2D Fourier expansion**

In practice, the truncation of the 2D Fourier expansion, i.e., the degree $L$, is set to the same value as the maximum degree of irreps.

Reviewer comment (Experiments):
Is it possible to choose other basis systems instead of Fourier? May I know how choosing different systems affects the result of this paper?

Author response:
> **Choosing other basis systems, and possible effects**

This is a good point. As introduced in Section 3, our efficient approach derives from the new perspective on the Tensor Products of irreps via Gaunt coefficients: the tensor product of irreps via the Gaunt coefficients is equivalent to multiplication between spherical functions represented by linear combinations of spherical harmonic bases. Given this, it is natural to use 2D Fourier basis to represent spherical functions via a change of basis, because we can leverage efficient computational methods like FFTs for acceleration. In fact, there also exist other 2D bases to represent spherical functions, e.g, 2D Haar Wavelets, which are possible to use instead of a Fourier basis. In our work, we provide the following criteria for choosing the different basis:
- *Approximation quality*: the multiplication results between spherical functions should not differ when using different bases, i.e., the approximation error of the change of basis should be small enough. In our work, we provided empirical evidence on the satisfactory approximation quality of the 2D Fourier basis in the sanity check experiment with SEGNN on the N-body simulation task. In Figure 1, we can see that our Gaunt Tensor Product achieves comparable, and even better performance compared to the standard implementation, which verifies the approximation quality of using a 2D Fourier basis.
- *Computational complexity*: another critical point to consider is the computational complexity of calculating the multiplication between spherical functions after the change of basis. In Section 3.2 of our paper, we provide a thorough analysis on the computational complexity by choosing the 2D Fourier basis, which requires the following steps to be efficient enough to achieve acceleration:
  - Converting coefficients from spherical harmonics to a new basis.
  - Perform multiplication between spherical functions in the new basis.
  - Converting coefficients from new basis back to spherical harmonics.

  Since the computational complexity of the standard implementation of CG tensor products is $\mathcal{O}(L^6)$, the above three steps should thus have lower complexity than this.

We thank you again for your efforts in reviewing our paper, and we have carefully revised our paper and replied to each of your comments. We look forward to your re-evaluation of our submission based on our responses and updated results.

Reviewer comment (Theory):
The current derivation and application are closely tied to $\text{O}(3)$. Is it feasible to extend this convolutional perspective beyond $\text{O}(3)$, e.g., to other Lie groups like $\text{SU}(2)$? It would be beneficial to see discussions on the aspect of its generalization.

Author response:
> **Extending our work to other groups like SU(2)**

This is a good point. In the literature, other Lie groups also describe important symmetries exhibited in real-world applications, e.g., $SU(2^N)$ in quantum mechanics, $SU(3)$ in quantum chromodynamics, and $SO^{+}(1,3)$. Equivariant models respecting these symmetries are demonstrated to be powerful in these tasks [1, 2]. It is thus interesting to extend our methodology to these scenarios. In practice, there are two points we need to check:
- The framework to respect symmetries: The equivariant networks in our paper respect symmetries based on the Tensor Products of irreps. To extend our approach to other Lie groups, we need to also generalize the Tensor Products of Irreps for these groups. This includes (1) constructing irreps of the target Lie group; (2) decomposing the Tensor Products of irreps by using generalized Clebsch-Gordan coefficients of the targeted group; (3) computing the tensor product results.
- Relationship between Clebsch-Gordan coefficients and Gaunt coefficients: For other Lie groups, it is crucial to check the relationship between the generalized version of Clebsch-Gordan coefficients to Gaunt coefficients, which determines whether we can express the computation of tensor products as multiplications between functions expanded by bases.

In our preliminary study, we have several interesting results that correspond to the above two points. First, we noticed that there exists a recent work [3] that proposed a general framework to respect symmetries to Reductive Lie Groups including ($SO^+(1,3)$,$SU(N)$ and etc.), which generalizes the Tensor Product of irreps and demonstrates the feasibility to other groups. The presented framework exactly matches our first point, thus meaning that it could be possible to apply our approach here. Moreover, to check the relationship between the Clebsch-Gordan coefficients and Gaunt coefficients, we use the Wigner-Eckart theorem. In fact, we can generalize the Wigner-Eckart theorem presented in Section 3.1 of our paper to other Lie groups, which is also supported in [4]. The provided generalized Wigner-Eckart theorem serves as a powerful tool to examine the relationships between different groups. From these preliminary results, we believe it is feasible to further explore the possibility of our methodology for other Lie groups, and leave this for future work.

Reviewer comment (Experiments):
I would like to see the inference time of EquiformerV2 model with and without the proposed Gaunt tensor products.

Author response:
> **Inference time comparison between EquiformerV2 with and without the proposed Gaunt Tensor Product operations**

We provide inference time comparisons here. In our experiments, we use the EquiformerV2 as the base architecture on the OC20 S2EF task to demonstrate the effectiveness of our approach. We apply our Gaunt Tensor Product to construct an efficient Equivariant Feature Interaction operation called Selfmix, and add it to each layer of the EquiformerV2 model. As we introduced in our paper, the equivariant operations in EquiformerV2 cannot be used to accelerate the computation of operations in the Equivariant Feature Interaction class. In contrast, our Gaunt Tensor Product can provide significant acceleration here. Empirically, this operation in EquiformerV2 with the Gaunt Tensor Product is 5.14x times faster than that without Gaunt Tensor Product on the OC20 S2EF task, demonstrating the efficiency of our approach.

Reviewer comment (Presentation):
Please adjust Equation (23)(26)(36) as they are written in the margins of the paper.

Author response:
> **Equations in margins of the paper**

Thank you for pointing this out. We have fixed this in our revised paper.