Training for vector and tensor variables

[chuyuanchen]

Dear developers, currently tensor information in Allegro is abandoned in the end and a scalar energy is produced. However, I think the equivariant latent space can be useful to train vector and tensor variables which have rotation symmetry? For example, the atom wise born effective charge (a rank 2 tensor), or the polarization vector, which describe the linear response of the systems (force and energy) to external electric field and would allow dynamics simulation under external field. Is there any plan to implement such feature to train for vectors and higher tensors? I think all is needed is to pick up the correct irreps in the final layer (e.g., instead of "0e", output "0e+1o" for vector) and modify the loss function, does it make sense to you? Thank you very much.

[cw-tan]

We do not have plans to implement higher-order tensor prediction, but would be happy to collaborate if anybody wants to implement it. The tricky part is that we currently use a DenseNet structure for the scalar track, so there is some uncertainty for how the tensor features should be propagated, or if one should just extract the last layer's tensor features. So it's not just "administrative" implementation to enable tensor predictions, but also some questions about architectural design that requires discussion, testing, etc.