[BUG] Polarizability cannot be correctly trained

[Yi-FanLi]

Bug summary

I am trying to train a polarizability model with the example in the GitHub repo. When training the global-only polarizability model, I obtain the following learning curve:

Although the training data seem to be fitted, the trained model is predicting the diagnoal components of polarizability with an abnormal constant shift:

The non-diagonal components do not show the constant shift:

Moreover, the model cannot work on the validation dataset at all:

If I extract the diagonal terms of polarizability and fit them with a dipole model, the suspicious constant shift does not appear:

while the model still does not work at all for the validation data:

Moreover, the behavior of the atomic polarizability model is very similar to the global polarizability, i.e., showing the suspicious constant shift as well as failure to generalize on the validation set:

Therefore, I suspect that there should be a bug in the training of the polarization model. I am not sure whether the issue shown in the validation set stems from the dataset or the code, but there might also be a bug in the general tensor part.

DeePMD-kit Version

DeePMD-kit v3.0.1.dev55+g3e9cf881

Backend and its version

TensorFlow v2.18.0

How did you download the software?

pip

Input Files, Running Commands, Error Log, etc.

dp train polar_input.json

Steps to Reproduce

I used examples in the GitHub repo, but the tasks I used can also be found here:

polar_issue.tar.gz

Further Information, Files, and Links

No response