Implementing TensorNet in modelforge

[wiederm]

Description

TensorNet is a message-passing neural network with a unique way to satisfy equivariance under the orthogonal group. Previously, two approaches have dominated:

1. use of Cartesian vector features that are equivariante (internal vector features and outputs transform according to O(3) symmetry), examples are SAKE and PaiNN
2. use of irreducible representations of the O(3) group as internal features (MACE, Allegro)

TensorNet uses Cartesian tensor representations and embeddings which are O(3) equivariant.

To re-implement this model, we need to modify our current Embedding module (which is built for embedding of scalar values) and implement the Interaction module.

Embedding

Following the default modelforge input preparation, the model gets as input (for a given radius $r_c$ within a neighborhood $N(d_{ij} < r_c)$:

• $r_{ij} = r_j - r_i$ , between central atom $i$ and neighbors $j$
• $d_{ij} =||r_{ij}||$

Additional properties that need to be calculated for the Tensornet embedding:

• normalized distance vectors: $r_{ij}' = r_{ij} / || r_{ij} ||$
• $I^{ij}$ (scalar features), $A^{ij}$ (vector features) and $S^{ij}$ (tensor features)

Scalar features

Scalar feature will be initialized as $I_{0}^{ij} = \text{Id}$, with $Id$ the identity matrix.

Vector features

Tensor features

The tensor feature will be initialized as the symmetric traceless tensor of the outer product of $r_{ij}'$: $S_0^{ij} = r_{ij}'r_{ij}'^T - \frac{1}{3}Tr( r_{ij}'r_{ij}'^T)Id$, with $Tr$ is the trace operator and $Id$ the identity matrix.

Interaction

Todos

• Embedding module
• Interaction module

Status

• Ready to go