[BUG] Polarizability cannot be correctly trained with `tf` backend

[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

[Yi-FanLi]

I want to mention that the shift in the prediction of polarizability appears only when shift_diag=True. Setting shift_diag=False can make accurate fitting of the training data.

Also, the shift only appears when doing inference with a compressed model. An uncompressed model does not cause this issue.

[njzjz]

Reason: constant_matrix is not recovered when the model is recovered. It should be handled in the same way as bias_atom_e in the energy model.

[Yi-FanLi]

By using #4581 , the trained model now gives the following output:

This test result shows that #4581 fixes this issue.