Goto

Collaborating Authors

 irreducible representation




TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

Neural Information Processing Systems

The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.


Appendix

Neural Information Processing Systems

The introduction of convolution and attention to the space of rays in 3D required additional geometric representations for which there was no space in the main paper to elaborate. We will introduce here all the necessary notations and definitions. We have accompanied this presentation with examples of specific groups to elucidate the abstract concepts needed in the definitions. Figure 10: The visualization of Plรผcker coordinates: A ray xcan be denoted as (d,m)where x is any point on the ray x, and dis the direction of the ray x. mis defined as x d. Given the action of the group G on a homogeneous space X, and given x0 as the origin of X, the stabilizer group H of x0 in G is the group that leaves x0 intact, i.e., H = {h G|hx0 = x0}. The group, G, can be partitioned into the quotient space (the set of left cosets) G/H and X is isomorphic to G/H since all group elements in the same coset transform x0 to the same element in X, that is, for any element g gH we have g x0 = gx0. Example 1. SE(3) acting on the ray space R: Take SE(3) as the acting group and the ray space R as its homogeneous space. We use Plรผcker coordinates to parameterize the ray space R: any x R can be denoted as (d,m), where d S2 is the direction of the ray, and m = x d where x is any point on the ray, as shown in figure 10. R is the quotient space SE(3)/(SO(2) R)up to isomorphism. Example 2. SE(3) acting on the 3DEuclidean space R3: R3 is isomorphic to SE(3)/SO(3). Consider another case when SE(3) acts on the homogeneous space R3; for any g = (R,t) SE(3) and x R3, gx = Rx+t. If the fixed origin is [0,0,0]T, the stabilizer subgroup is H = SO(3) since any rotation g = (R,0)leaves [0,0,0]T unchanged. The last example is SO(3) acting on the homogeneous space sphere S2. Given the fixed origin point as [0,0,1]T, the stabilizer group is SO(2).





A General Framework for Equivariant Neural Networks on Reductive Lie Groups

Neural Information Processing Systems

Convolutional Neural Networks (CNNs) (LeCun et al., 1989) have become a widely used and powerful tool for computer vision tasks, in large part due to their ability to achieve translation