Questions w.r.t tensor product in Allegro

[thu-wangz17]

Hi, this is an amazing work! I have two questions w.r.t the tensor product in Allegro.

1. In the paper, tensor product is conducted to obtain $V_{n, (l_1, p_1, l_2, p_2)\rightarrow(l_{out}, p_{out})}^{ij, L}$. In my understanding, the neighborhoods are firstly aggregated and then the tensor product will only conducted once. Since in the original formula (i.e. Equation 16), the number of tensor product operation is the number of the neighborhood, which is larger than one. Thus the method could reduce the computational cost greatly, Is my understanding right?
2. To obtain type-$l_{out}$ outputs, type-$l_1$ and type-$l_2$ features should be coupled with CG-coefficient. Do these CG-coefficient be included in $w_{n, l_2, p_2}^{ik, L}$? In the paper, $V$ has omited the $m$ index, so do the Equation 16-18 include the summation of all the $m$ index? If so, the weight $w_{n, l_2, p_2}^{ik, L}$ is actually $\sum_{m_1, m_2} w_{n, m_1, m_2, p_2}^{ik, L}$ for a fixed $l_1$ and $l_2$?

[simonbatzner]

Hi @sakuraiiiii, thanks for reaching out, answers below:

1. Your understanding is correct. In a naive way - eqn 16 - one would have to compute a TP between every feature $V_{ij}$ and every spherical harmonics edge embedding $Y_{ik}$. We use the fact that the TP is bilinear to pull the TP out of the sum, which reduces it to a single TP. Before it was a sum over TPs between edge and edge. Now it's a single TP of edge vs sum over edges (i.e. environment).

2. That's correct, we use the CG coefficients to decompose the results of the tensor product into a set of irreps. But no, these CG coefficient are not part of $w^{ik, L}_{n, l_2, p_2}$, that is only a weight for each edge in the environment. After this sum, we compute the TP on a per-feature-basis, so if you have n features and as a simple example imagine they were all one l=1, p=odd and your output was l=0, p=even (so a $1o \otimes 1o \rightarrow 0e$ TP corresponding to a dot product), then you'd compute n of these tensor products, all on a per-feature basis, so feature n=1 of the edge feature is talking to feature n=1 of the environment embedding, n=2 talking to n=2, etc. There is no separate weight along the m-dimension. The CG coefficients are implicit in the $ \otimes$ as we define it in eqn 4.

I hope this makes sense + answers your question.

[thu-wangz17]

Very clear. Thank you very much.