Skip to content

Fix: Square atom_norm in non-Huber energy and virial loss calculations. #5332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
anyangml wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Fix: Square atom_norm in non-Huber energy and virial loss calculations. #5332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f261151
fix: Square atom_norm in non-Huber energy and virial loss calculations.
anyangml Mar 23, 2026
247f053
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions deepmd/dpmodel/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ def call(
if self.loss_func == "mse":
l2_ener_loss = xp.mean(xp.square(energy - energy_hat))
if not self.use_huber:
loss += atom_norm_ener * (pref_e * l2_ener_loss)
loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
else:
l_huber_loss = custom_huber_loss(
atom_norm_ener * energy,
Expand Down Expand Up @@ -312,7 +312,7 @@ def call(
xp.square(virial_hat_reshape - virial_reshape),
)
if not self.use_huber:
loss += atom_norm * (pref_v * l2_virial_loss)
loss += atom_norm**2 * (pref_v * l2_virial_loss)
else:
l_huber_loss = custom_huber_loss(
atom_norm * virial_reshape,
Expand Down
4 changes: 2 additions & 2 deletions deepmd/pd/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ def forward(
l2_ener_loss.detach(), find_energy
)
if not self.use_huber:
loss += atom_norm * (pref_e * l2_ener_loss)
loss += atom_norm**2 * (pref_e * l2_ener_loss)
else:
l_huber_loss = custom_huber_loss(
atom_norm * energy_pred,
Expand Down Expand Up @@ -404,7 +404,7 @@ def forward(
l2_virial_loss.detach(), find_virial
)
if not self.use_huber:
loss += atom_norm * (pref_v * l2_virial_loss)
loss += atom_norm**2 * (pref_v * l2_virial_loss)
else:
l_huber_loss = custom_huber_loss(
atom_norm * model_pred["virial"].reshape([-1]),
Expand Down
4 changes: 2 additions & 2 deletions deepmd/pt/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,7 @@ def forward(
l2_ener_loss.detach(), find_energy
)
if not self.use_huber:
loss += atom_norm * (pref_e * l2_ener_loss)
loss += atom_norm**2 * (pref_e * l2_ener_loss)
else:
Comment on lines 241 to 243
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are tests comparing losses across backends, but they don’t assert the intended natoms-scaling explicitly. Since this change fixes a subtle normalization bug (MSE energy/virial terms should scale with 1/natoms^2), consider adding a small regression test that checks the loss contribution scales correctly when natoms changes, so future refactors can’t accidentally revert to 1/natoms.

Copilot uses AI. Check for mistakes.
l_huber_loss = custom_huber_loss(
atom_norm * energy_pred,
Expand Down Expand Up @@ -421,7 +421,7 @@ def forward(
l2_virial_loss.detach(), find_virial
)
if not self.use_huber:
loss += atom_norm * (pref_v * l2_virial_loss)
loss += atom_norm**2 * (pref_v * l2_virial_loss)
else:
l_huber_loss = custom_huber_loss(
atom_norm * model_pred["virial"].reshape(-1),
Expand Down
6 changes: 4 additions & 2 deletions deepmd/tf/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,7 @@ def build(
more_loss = {}
if self.has_e:
if not self.use_huber:
loss += atom_norm_ener * (pref_e * l2_ener_loss)
loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Finish the TensorFlow normalization fix in the remaining loss variants.

EnerStdLoss is corrected here, but EnerSpinLoss.build() still uses the old single-factor normalization at Line 757 and Line 770, and EnerDipoleLoss.build() still does the same at Line 1039. Those paths will keep the old 1/N scaling after this PR.

Minimal follow-up diff 
-        if self.has_e:
-            l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
+        if self.has_e:
+            l2_loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
@@
-        if self.has_v:
-            l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
+        if self.has_v:
+            l2_loss += global_cvt_2_ener_float(atom_norm**2 * (pref_v * l2_virial_loss))
@@
-        l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
+        l2_loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)

Also applies to: 373-375

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@deepmd/tf/loss/ener.py` at line 349, EnerSpinLoss.build and
EnerDipoleLoss.build still use the old single-factor 1/N normalization; update
both to use the same squared per-atom normalization pattern used in EnerStdLoss
(i.e., multiply the loss term by atom_norm_ener**2 * (pref_e * l2_ener_loss)
instead of a single atom_norm_ener factor) so all energy loss variants apply the
corrected normalization consistently; locate the energy-loss accumulation in
EnerSpinLoss.build and EnerDipoleLoss.build and replace the old single-factor
scaling with the squared atom_norm_ener**2 * (pref_e * l2_ener_loss) pattern
used in EnerStdLoss.

else:
l_huber_loss = custom_huber_loss(
atom_norm_ener * energy,
Expand All @@ -370,7 +370,9 @@ def build(
)
if self.has_v:
if not self.use_huber:
loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
loss += global_cvt_2_ener_float(
atom_norm**2 * (pref_v * l2_virial_loss)
)
Comment on lines 348 to +375
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This updates EnerStdLoss’ non-Huber MSE scaling to use atom_norm2 / atom_norm_ener2, but the same file still has other non-Huber MSE loss accumulations using a linear atom_norm factor (e.g., EnerSpinLoss build around l2_loss += atom_norm_ener * ... and virial term with atom_norm * ..., and EnerDipoleLoss build around l2_loss += atom_norm_ener * ...). That leaves inconsistent normalization depending on which loss class is used, and also conflicts with the per-atom RMSE reporting (rmse_e computed as sqrt(l2_ener_loss)/natoms). Please update those remaining energy/virial MSE paths to square the atom normalization as well, or clarify why those classes should behave differently.

Copilot uses AI. Check for mistakes.
else:
l_huber_loss = custom_huber_loss(
atom_norm * tf.reshape(virial, [-1]),
Expand Down
Loading