Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes an E(n) equivariant network acting on pointclouds with edges. The method is simple to compute and shows competitive performance on various tasks.

The authors claim that their attention method is universal, and that it is more expressive than "dot product scalarization" attention used in prior wok.
strength_and_weaknesses: ## Strength
- The method is evaluated on a wide variety of tasks.
- The method performs competitively.

## Weakness
- The core spatial attention method is a small variation of / combination of pre-existing methods.

### Representation of theorem 1
I think that theorem 1 is misrepresented in the paper. The authors write
> spatial attention is capable of universally approximating any functions defined on local node environment while preserving E(n)-invaraince/equivariance in arbitrary n-dimensional space.

I would interpret this as saying that any invariant/equivariant function from the neighbourhood features and coordinates can be represented by spatial attention. I would thus expect a theorem like:

> Given $N$ neighbours, $C$-dimensional features and $n$ spatial dimensions, let $V \subset \mathbb R^{C^N}$ be the subset of the edge features where all edges have distinct values. Then for any function $g : V \times \mathbb R^{n^N} \times \mathbb R^n \to \mathbb R$, that is invariant under permutations and $E(n)$, there is a spatial attention layer that matches $g$ arbitrarily well on all inputs.

However, this is not what is shown in the proof. The proof states that for **one given value of edge features** $\mathbb R^{C^N}$, and one invariant function of the coordinates $g: \mathbb R^{n^N} \times \mathbb R^n \to \mathbb R$, there is a corresponding spatial attention layer. The spatial attention is thus not shown to be universally express any function from the node neighbourhood with node features. Other than the paper suggests, the theorem thus also has no implication about whether the SAKE network can represent any equivariant function of pointclouds with node features. I think the claim that the layer is a universal approximator is thus incorrect.

Like theorem 1, I think that remark 3 wrongly claims universality.

### Remark 2
The DPS method is not so clearly defined, but for equation $\phi^{DPS}(v; \lambda_k, \lambda_q)=(\phi^{SA}(v; \lambda_+)^2 - \phi^{SA}(v; \lambda_-)^2)/2$ (Sec 8.2) to make sense, it seems like:

$\phi^{DPS}(v; \lambda_k \lambda_q)_i = \sum\_{uu'}\lambda\_{q,i,u}\lambda\_{k,i,u'} \langle f\_u , f\_{u'} \rangle$

where $f_u=f(e_{uv})\in \mathbb R^n$ and $\lambda\_{q,i,u}=\lambda\_{q,i}(h\_{e\_{uv}}) \in \mathbb R$ is the intended defintion.

If so, the claim in Sec 8.2 that with orthogonal edges $\phi^{DPS}(v)=0$ is clearly false. For example, if there is one neighbour with edge $e$, $N_\lambda=1, \lambda=1, f(e)=e, \mu(c)=c$, then it looks like $\phi^{DPS}(v)=||e||^2$. Thus the claim that the spatial attention is more expressive than DPS, doesn't appear correct.

A forteriori, it appears to me that a construction similar to the proof in Sec 8.1 can be used for $\phi^{DPS}$ to show that it is just as universal as the spatial attention method as defined in Theorem 1. Appropriately choosing $\lambda$, one can construct a list of all inner products $\langle x_v - x_u , x_v - x_{u'} \rangle$, from which all pairwise distances between the neighbours can be constructed via $||x_u - x_{u'}||^2=||x_v - x_{u}||^2+||x_v - x_{u'}||^2 - 2 \langle x_v- x_u, x_v - x_{u'} \rangle$. Subsequently, the same argument  as made in sec 8.1 can be made to show that this construction is universal as defined in theorem 1.

Also with a definition of $\phi^{DSA}$ that's more similar to conventional attention, like $\phi^{DPS}(v; \lambda_k \lambda_q)_{uu'} = \sum\_{i}\lambda\_{q,i,u}\lambda\_{k,i,u'} \langle f\_u , f\_{u'} \rangle$, choosing $\lambda=1,f\_u=e\_{vu}$, it is universal as defined in theorem 1.

In conclusion, regardless of how exactly $\phi^{DSA}$ is defined, it doesn't appear any less expressive than the the proposed spatial attention method. I thus don't think Remark 2 is correct.

## Typo
- p2, first paragraph "invaraince"

clarity,_quality,_novelty_and_reproducibility: Clarity: There appear to be many inaccuracies in the formulae. This makes it very hard to understand the details of the proposed method and follow the proofs.
- Above equation (6), the type $\phi: \mathcal X \times \mathcal H \to \mathcal H$ does not seem consistent with definition (6), which is more like $\phi: (\mathcal X \times \mathcal H)^{\mathcal N(v)} \to \mathcal H$.
- Similarly, $\mu: N_\lambda \to \mathcal H$ should probably be $\mu: \mathbb R^{N_\lambda} \to \mathcal H$.
- In theorem 1, the function $g$ should take a set of neighbours as input, not a single neighbour.
- As stated above, the definition of $\phi^{DPS}$ is not clear.
- In section 8.2, $\lambda_{u, i}$ in the definition of $\lambda_+$ and $\lambda_-$ should be $\lambda_{q, i}$

Quality: I have serious concerns about the correctness of both theoretical claims.

Novelty: The method is a small variation of pre-existing work.

Reproducibility: The source code is provided.
summary_of_the_review: I have serious concerns about both theoretical claims in this paper: that the spatial attention mechanism is universal, and that it is more expressive than different attentional methods. I thus recommend rejection of this paper.

---
After their revision, the theoretical issues are fixed. I now think the paper is a valuable contribution to the literature on equivariant pointcloud networks and should be accepted. I weakly recommend acceptance, as the novelty is still limited.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper exploits several kinds of attention mechanisms (Euclidean and semantic attention, spatial attention) upon EGNN [6]. Necessary theoretical claims have been made to justify the expressivity of the proposed spatial attention. Experiments are conducted on invariant tasks and equivariant tasks.
strength_and_weaknesses: Strengths:

1. The design of spatial attention in Eq. (5) is well motivated and novel. As illustrated in the proof of Theorem 1, if certain condition meets, the spatial attention is able to recover the local environment of each node. The proof of Theorem 1 is interesting by recovering the distance matrix.

2. The part of Related Work is well organized, providing a good guidance of understanding of how recent progress is made.

3. The experimental comparison with various methods have been conducted. 

Weaknesses:

1. The condition $h_{e_{vu_i}}\neq h_{e_{vu_j}}$ in Theorem 1 is over-restricted and easily broken when two neighbor nodes share the same embedding (the same node features and local context). The authors claimed that the inequality holds when the system is not strictly symmetrical. Unfortunately, many practical molecular systems are symmetric such as benzene and many other materials such as crystals. 

2. How to build a permutation-invariant and E(n)-invariant (or equivariant) function has been discussed in [A], which is not cited, unfortunately. In [A], Propositions 10-11 have presented the universal form of message passing from local neighbors up to the O(n) equivariant. More importantly, the results by [A] does not require the condition $h_{e_{vu_i}}\neq h_{e_{vu_j}}$ that is indispensable in this paper. The authors are suggested to compare the difference with [A] which seems more elegant. 

3. Moreover, the proof of Theorem 1 only shows the results on invariant function other than equivariant function. In [B] and its citated references, the universal form of equivariant message passing can be derived from TFN and GemNet, both of which again do not require any restriction of node embedding. 

4. It is unclear why "The distinctness condition corresponds to the full-rank requirement in GMN".

5. The authors have shown experimental comparisons on MD17 and QM9, which are already well explored, and the results of the proposed method are not clearly superior to other methods. The authors are strongly suggested to conduct performance on other recent and more challenging tasks, such as the trajectory prediction setting of MD17 and the benchmark OC20-22. 

6. The ablation studies are insufficient and unclear to justify the contribution of this paper. It seems from Algorithm 1 that the main novelty lies in the proposal of spatial attention Eq. 6 along with other minor techniques. The authors should detail how each proposed component affects the performance and provide necessary analyses.




[A] Scalars are universal: Equivariant machine learning, structured like classical physics, NIPS 2021.

[B] GemNet: Universal Directional Graph Neural Networks for Molecules, NIPS 2021.
clarity,_quality,_novelty_and_reproducibility: + This paper is overall well written.

- The novelty is marginal and mainly lies in the proposal of spatial attention mechanism (Eq.6), which as mentioned above is over-restricted compared to recent methods. 

- The results of EGNN in the table in Section 11 are much worse than other methods, which seems quite weird and unconvincing.

summary_of_the_review: The authors are tackling a variable problem on enhancing the expressivity of equivariant GNNs on physical modeling. However, given the limited novelty, technicality, and insufficient ablations, I suggest weak reject for its current version.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: Some of the earlier successful methods for physical modelling (like Schutt et al., 2017 and Bartok et. al, 2013) are unable to fully account for the geometry of the system. To model particle arrangements, they rely on using radial information (which captures inter-particle distances). It should be straightforward to see that such methods can't fully capture the geometry of node environments. For instance, one might have particular arrangements which are quite different relative to each other, even when the distance from a node remains unaltered.  More recent methods use a truncated series (resembling a multipole expansion) that uses spherical harmonics to encode higher-order interactions between particles -- and then incorporate it in a message-passing formulation. Such methods have proven to be quite powerful and have been developing quite rapidly over the past few years. However, they do tend to be somewhat expensive computationally. This paper attempts to address this by proposing a method (SAKE) that relies on a simple attention-based formulation that uses linear combinations of edge vectors while being equivariant. More specifically, SAKE uses the _norm_ of linear combinations of edge features to specify the node environment. The method and some of the contributions are summarized below: 

The background on equivariance and a high-level description of the message-passing formalism is provided in sections 2.1 and 2.2. The problem is formulated as implementing a function f_\theta: \mathcal{X} x \mathcal{H} \mapsto \mathcal{X} x \mathcal{H}. The space \mathcal{X} x \mathcal{H} forms a joint space of n-dimensional coordinates, and D-dimensional semantic embeddings. The approach that is most similar in terms of applicability and import is that of Satorras et al. The framework for spatial attention is defined in section 4. It uses an equivariant function on the edge embeddings. These embeddings are used to generate attention weights \lambda. The authors suggest taking N_{\lambda} such linear combinations, computing their norms and then further concatenating them (multiple heads), and mapping it to the \mathcal{H} space (by function \mu). Both the \lambda and \mu functions are modelled via MLPs. This simple spatial attention module is E(n) invariant. It is also universal for E(n) invariant functions, as encapsulated in theorem 1. The edge embedding is computed by borrowing ideas from Schutt et. al. and Satorras et al. (equation 7). In addition, there are two more attention modules to promote anisotropy (described on page 5). Finally, there is a similar term as in Satorras et al. that models updating a fictitious velocity (eq 10 and 11). 

SAKE is equivariant on the geometric space and invariant on the embedding space. The experiments are divided into three parts. For invariant modelling, the authors consider standard tasks on MD17, ISO17, and QM9.  On MD17, the approach is competitive, and beats all but one of the baselines, but with a significant time gain. While only three baselines are considered for ISO17, SAKE is reported to be better than the competition, while again reportedly faster. Then there are two equivariant tasks (charged N-body and walking motion) where a similar trend is reported. Finally, the method is used to describe a normalizing flow method and is shown to be competitive with the approaches of Kohler et al. and Satorras et al. 2022. 

strength_and_weaknesses: Strengths: 
- The paper addresses an important problem. The method is well-motivated and quite simple and elegant.
- Experimental results are quite encouraging. The presented approach is competitive with various methods that were proposed recently -- although it is not the best across the board -- but has a significant advantage in terms of run-time. 
- I think the formulation also suggests approaches to optimize for the architecture, which can improve the experimental results further. 

Weaknesses:
- I think the writing of the paper is a little bit clunky. 
- The novelty of the method is somewhat limited. But I think the experimental results make the case for it as far as I am concerned.
- The code is not properly anonymized. 

Minor comments:
- At the beginning of page 3, \mathcal{H} is D dimensional, but at the beginning of section 4, it is written to be C dimensional. Might want to remove this inconsistency. 
- Please re-write the abstract. For instance, the first line is too long and seems to have some punctuation missing. 

Some typos
- Page 2, line 2: invaraince --> invariance.
- Page 2, line 2 in paragraph 2: "keep track" --> keeps track. 
- Page 3, first paragraph of section 3: "Compared to SAKE, these models are not capable in equivariant modeling" --> "Compared to SAKE, these models are not capable of equivariant modeling".
- Page 4, below equation 6: "function that operate on the edge vector" --> "function that operates on the edge vector"
- Page 4, below remark 1: "With all local degrees of freedom are incorporated in spatial attention" --> "With all local degrees of freedom incorporated in spatial attention,"
- Page 5, above section 6: "In relation to spherical harmonics-based models." --> "Relation to spherical harmonics-based models."
- Page 7, section 6.2: "model can predict the evolving of a physical system sufficiently long after initial conditions" --> "model can predict the evolution of a physical system sufficiently long after initial conditions"
- Page 8, next to table 5: "invaraince" --> invariance. 
- Page 8, next to table 5: "as the composing" --> "as the composition"
- Page 8: "Note that this reduce" --> "Note that this reduces"
clarity,_quality,_novelty_and_reproducibility: The writing of the paper could be better but is generally quite clear. To the best of my knowledge, the proposed method is novel (although it relies very heavily on two papers). It seems to have major advantages in terms of scalability and generally returns competitive performance across a range of tasks. 
summary_of_the_review: See the above summary.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper introduces spatial attention into E(n)-equivariant neural networks to describe the node environments for the modeling of many-body systems. The architecture builts on the EGNN introduced by Satorras et al. and introduces spatial attention to improve the performance in several  invariant and equivariant prediction tasks.
strength_and_weaknesses: Strength:
- The introduced spatial attention improves the expressiveness of the EGNN and preserves the E(n)-Equivariance of the original model.
- The paper conducts extensive experiments to demonstrate the improved performance in invariant property prediction, equivariant property prediction, and density estimation in equivariant normalizing flows.
- Model inference is significantly faster than architectures based on spherical harmonics.
- Detailed ablation study is performed to compare SAKE and EGNN. 

Weakness:

- The ablation results are confusing. EGNN is significantly worse than SAKE (for Aspirin, SAKE has MAE of 9.9 and EGNN has an MAE of 298.05). It looks like the EGNN is not properly trained. Further, the authors need to provide more detailed descriptions for the ablation study. What does “No update” mean? 
- What is the equivariant function f in Eq. 6? I cannot find a detailed explanation of the function f in the paper. Is it an identity function or other equivariant function?
- The Proof 8.1 only considers the special case of f = I. Will Theorem 1 still hold if f is not the identity function?
- Some inference time comparisons are incomplete. In Table 1, the inference time of SchNet, sGDML, PaiNN is missing. In Table 5, the inference time of GNN, EGNN is also missing. How much overhead does the attention mechanism introduce compared with GNN and EGNN? 
- In QM9, the model performs less competitively in some intensive property prediction tasks. The authors hypothesize that it can be mitigated by size-invariant pooling functions. This should be a straightforward improvement by using a mean pooling rather than sum pooling. I suggest providing additional results using the mean pooling.

clarity,_quality,_novelty_and_reproducibility: - Overall, the paper is well-written and easy to follow. There are some technical details missing which makes it hard to read occasionally.
- The novelty of the spatial attention architecture seems marginal. It extends the equivariant update in Eq. 4 of the EGNN paper from node coordinates to node features. 
- The authors provided the code to reproduce the results in anonymous repo.

summary_of_the_review: Overall, the paper introduces a novel equivariant spatial attention layer, extending the EGNN architecture. The authors demonstrated improved performance in a broad range of tasks. However, the novelty of the spatial attention layer seems marginal. More detailed ablation study is needed to clarify the usefulness of the introduced spatial attention layer.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper introduces an efficient equivariant network design based on spatial attention. The method performs competitively relative to various recent methods on a number of quantum chemistry benchmarks while being much faster. Since speed is one of the main selling points of ML for chemistry, this is an important result. Reviewers appreciate the simplicity and effectiveness of the approach, as well as the extensive empirical validation. Theorem 1 gives a (limited) universality result for spatial attention layers.

Several reviewers noted that the spatial attention method is a small variation of existing methods, and thus not very novel. However, given that the exact method was not known, and performs well in terms of speed and accuracy, I consider this to be a valuable contribution regardless. Another minor concern is writing quality, and I encourage the authors to continue to polish the paper.
justification_for_why_not_higher_score: Limited novelty, writing quality
justification_for_why_not_lower_score: The method is fast and accurate, and will be of high interest to the ML for quantum chemistry community
Decision: Accept: poster