TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

Summary

The paper presents TensorNet, a novel O(3)-equivariant message-passing neural network for efficient representation of molecular systems in scientific research. This model utilizes Cartesian tensor atomic embeddings to simplify feature mixing via matrix product operations. By decomposing tensors into rotation group irreducible representations, it enables independent processing of scalars, vectors, and tensors when required. TensorNet outperforms higher-rank spherical tensor models in performance, utilizing fewer parameters, even with a single interaction layer for small molecule potential energies. Additionally, TensorNet can accurately predict vector and tensor molecular quantities on top of potential energies and forces, greatly reducing the model's computational cost. Therefore, TensorNet provides a promising framework for developing state-of-the-art equivariant models with enhanced efficiency and computational affordability.

Strengths

1. This paper adeptly presents TensorNet, a new learning model that not only establishes state-of-the-art performance but does so with a remarkable reduction in the number of parameters utilized. This marks a significant leap in model efficiency without sacrificing performance quality, setting a new benchmark in the field. 2. The majority of the empirical outcomes display a marked enhancement over existing methods, with the implementation of TensorNet consistently yielding superior results. The experimental evidence provided substantiates the model's efficacy, reinforcing the robustness and applicability of this innovative approach in real-world scenarios.

Weaknesses

The paper's exposition of the model architecture can be somewhat challenging to comprehend due to its complex nature. A potential improvement would be the inclusion of intuitive diagrams or visual aids within the main body of the text, not just in the Appendix. Simplified illustrations, possibly even a step-by-step visual guide, could greatly enhance the reader's understanding of the architecture and make the methodology more accessible to a broader audience.

Questions

1. The proposed architecture (illustrated in Fig. A1) appears quite intricate. Could you provide more insight or intuitive reasoning behind the specific design choices in the model? What fundamental principles or considerations have influenced the complexity and uniqueness of this architecture? 2. Could you delve deeper into the primary constraints or drawbacks of the TensorNet model compared to the existing methods in the field? What are the potential areas where other models might still hold an advantage, and how does TensorNet aim to address these challenges in future iterations or improvements?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This paper proposes TensorNet, an O(3)-equivariant neural network architecture for molecules. Using the decomposition of a 3x3 matrix into a scalar, vector, and matrix shown in Eq. (2), TensorNet efficiently computes the interaction of O(3)-equivariant features up to l=2, where l is the degree (frequency) of the O(3) representation. The performance is evaluated on several standard benchmarks, such as qm9, which show that TensorNet achieves equivalent or better performance than baselines.

Strengths

1. TensorNet is a novel method. I have never seen the construction of an equivariant net based on the decomposition (2). 2. The performance is evaluated with several different datasets and in terms of different metrics (e.g., prediction error, computation speed). 3. The performance is comparable to or better than existing approaches.

Weaknesses

1. The paper has room to improve in terms of presentation. I'm unfamiliar with the chemistry (molecule) domain, and some parts seem challenging to understand without expertise. For example, a vector r_ij is defined on line 175, but the mathematical definition is not described. Also, the meaning of the cutoff radius is unexplained. Another point is that there is no reference nor citation to Eq. (2), the core equation of this paper. These are not well known in the machine learning community, and it could be better to introduce them in plain words. 2. The limitations of the proposed method are not explicitly discussed. One limitation is that TensorNet cannot capture higher degree (l>2) information of O(3). 3. It is argued that some existing methods have a downside: "the computation of tensor products in most of these models containing higher-rank tensors and pseudotensors can be expensive" (lines 111-112). However, only one method (ET) is compared with the proposed method in terms of computational cost. Also, there is no theoretical evaluation of the complexity e.g., using big-O notation.

Questions

n/a

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Limitations are not explicitly addressed. No particular concern for potential negative societal impact.

Summary

This paper proposes a cartesian tensor representation for efficient learning of molecular potentials. It enables the feature mixing process to be a simple matrix product operation. In addition, the matrix product operation is simplified by cost-effective decomposition techniques. Experimental results demonstrate that the proposed method can effectively reach a comparable performance with a much smaller number of parameters.

Strengths

(1) The proposed method is very technically solid and this reviewer also thinks the efficient computation of equivariant architectures becomes increasingly important; (2) The extension of torchMD-net with efficient decomposition techniques is reasonably motivated, just like low-rank decomposition techniques for large language models; (3) Experimental settings are very solid, covering enough number of experimental settings, which fully support the effectiveness of the proposed method.

Weaknesses

(1) The presentation is not that easy to follow and the notation is a little bit complicated, which brings additional hardness for readers to understand the core idea. In addition, although the paper is more about mathematical techniques utilization, no figure illustration provided is still very tough for readers to quickly capture the core idea; (2) The proposed method is more like an application of efficient tensor decomposition techniques to the existing equivariant network architectures. Without molecular domain-specific insights somehow lowers the significance of the proposed method.

Questions

It seems the proposed architecture is an extension on TorchMD-net. Is it possible to apply this refinement to other equivariant models? If not, then this limitation will significantly lower its significance.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA.

Summary

The paper introduces Tensor3 Net, a message-passing neural network architecture designed for molecular systems representation. Tensor3 Net leverages rank-2 Cartesian tensor representations and O(3)-equivariance. The tensors are decomposed into rotation group irreducible representations, enabling separate processing of scalars, vectors, and tensors when necessary.

Strengths

The paper is well-written and the idea of building a O(3) equivariant model with Cartesian Tensor is valid, given the computation complexity of models such as e3nn, based on spherical harmonics.

Weaknesses

The one main weakness is that the author only model rank-2 tensors, which is relatively arbitrary. What it should be done is to introduce a network with general rank cartesian tensors (these will then be 3 x 3 x 3 x ... x 3 tensors). I see the value of adding the l=2 components in a cartesian way too marginal for publication at Neurips.

Questions

See weaknesses.

Rating

2: Strong Reject: For instance, a paper with major technical flaws, and/or poor evaluation, limited impact, poor reproducibility and mostly unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

1 poor

Limitations

The authors do not discuss the limitations of their approach.

Response

Thank you for the response. The rebuttal comments make sense to me (especially the reason of rank 2) and mostly resolve my concerns. I will raise my score.