Model training for StressForceOutput

[wkylee14]

Hi,

Thanks for making the Allegro repo. be public. I was just wondering if you have any guidance or though on preparing a config file when training the Allegro that predicts stress tensor outputs as well in addition to forces and total potential energies (StressForceOutput).

I've tried several attempts for my dataset with the following configurations, but losses for stress tensor decreases very slowly and marginally while losses for forces or energies keep decreasing after a certain number of training epochs.

Attempt 1: Applying PerAtomMSELoss
loss_coeffs:
forces: 1.
total_energy:
- 1.
- PerAtomMSELoss
stress:
- 1.
- PerAtomMSELoss

Attempt 2: Assigning more weights to loss for stress tensor predictions
loss_coeffs:
forces: 1.
total_energy:
- 1.
- PerAtomMSELoss
stress: 100.

Attempt 3: Assigning simple MSE loss function for stress tensor
loss_coeffs:
forces: 1.
total_energy:
- 1.
- PerAtomMSELoss
stress: 1.

Otherwise, do you recommend not to add loss for stress tensor?
Any recommendation or guidance when I use Allegro to predict stress tensors, forces, and potential energies would be welcome!

Kind regards,

[Linux-cpp-lisp]

Hi @wkylee14 ,

Thanks for your interest in our code!

stress should use a normal, and not PerAtom loss. Stress training issues are often linked to incorrect labels, either due to DFT issues, unit conversion issues, or an incorrect sign convention. (We follow the convention stress = (-1 / volume) * virial as discussed in various other threads on the nequip repo: https://github.com/mir-group/nequip/blob/main/nequip/nn/_grad_output.py#L346-L349).