Merge pull request #4 from anyangml/copilot/add-switch-rmse-normalization-fix

feat(loss): add intensive parameter for backward-compatible RMSE normalization

Author: anyangml

deepmd/dpmodel/loss/ener.py

@@ -90,6 +90,13 @@ class EnergyLoss(Loss):
         If true, use L2 norm of force vectors for loss calculation when loss_func='mae' or use_huber is True.
         Instead of computing loss on force components, computes loss on ||F_pred - F_label||_2.
         This treats the force vector as a whole rather than three independent components.
+    intensive : bool
+        If true, energy and virial losses are computed as intensive quantities,
+        normalized by the square of the number of atoms (1/N^2). This ensures the loss
+        value is independent of system size and consistent with per-atom RMSE reporting.
+        If false (default), uses the legacy normalization (1/N), which may cause the loss to scale
+        with system size. The default is false for backward compatibility with models trained
+        using deepmd-kit <= 3.0.1.
     **kwargs
         Other keyword arguments.
     """
@@ -116,6 +123,7 @@ def __init__(
         huber_delta: float | list[float] = 0.01,
         loss_func: str = "mse",
         f_use_norm: bool = False,
+        intensive: bool = False,
         **kwargs: Any,
     ) -> None:
         # Validate loss_func
@@ -155,6 +163,7 @@ def __init__(
         self.use_huber = use_huber
         self.huber_delta = huber_delta
         self.f_use_norm = f_use_norm
+        self.intensive = intensive
         if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"):
             raise RuntimeError(
                 "f_use_norm can only be True when use_huber or loss_func='mae'."
@@ -256,11 +265,15 @@ def call(
 
         loss = 0
         more_loss = {}
+        # Normalization exponent controls loss scaling with system size:
+        # - norm_exp=2 (intensive=True): loss uses 1/N² scaling, making it independent of system size
+        # - norm_exp=1 (intensive=False, legacy): loss uses 1/N scaling, which varies with system size
+        norm_exp = 2 if self.intensive else 1
         if self.has_e:
             if self.loss_func == "mse":
                 l2_ener_loss = xp.mean(xp.square(energy - energy_hat))
                 if not self.use_huber:
-                    loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
+                    loss += atom_norm_ener**norm_exp * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm_ener * energy,
@@ -335,7 +348,7 @@ def call(
                     xp.square(virial_hat_reshape - virial_reshape),
                 )
                 if not self.use_huber:
-                    loss += atom_norm**2 * (pref_v * l2_virial_loss)
+                    loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * virial_reshape,
@@ -525,7 +538,7 @@ def serialize(self) -> dict:
         """
         return {
             "@class": "EnergyLoss",
-            "@version": 2,
+            "@version": 3,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -546,6 +559,7 @@ def serialize(self) -> dict:
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
             "f_use_norm": self.f_use_norm,
+            "intensive": self.intensive,
         }
 
     @classmethod
@@ -563,6 +577,10 @@ def deserialize(cls, data: dict) -> "Loss":
             The deserialized loss module
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        version = data.pop("@version")
+        check_version_compatibility(version, 3, 1)
         data.pop("@class")
+        # Backward compatibility: version 1-2 used legacy normalization
+        if version < 3:
+            data.setdefault("intensive", False)
         return cls(**data)

deepmd/dpmodel/loss/ener_spin.py

@@ -50,6 +50,13 @@ class EnergySpinLoss(Loss):
         if true, the energy will be computed as \sum_i c_i E_i
     loss_func : str
         Loss function type: 'mse' or 'mae'.
+    intensive : bool
+        If true, energy and virial losses are computed as intensive quantities,
+        normalized by the square of the number of atoms (1/N^2). This ensures the loss
+        value is independent of system size and consistent with per-atom RMSE reporting.
+        If false (default), uses the legacy normalization (1/N), which may cause the loss to scale
+        with system size. The default is false for backward compatibility with models trained
+        using deepmd-kit <= 3.0.1.
     **kwargs
         Other keyword arguments.
     """
@@ -69,6 +76,7 @@ def __init__(
         limit_pref_ae: float = 0.0,
         enable_atom_ener_coeff: bool = False,
         loss_func: str = "mse",
+        intensive: bool = False,
         **kwargs: Any,
     ) -> None:
         valid_loss_funcs = ["mse", "mae"]
@@ -89,6 +97,7 @@ def __init__(
         self.start_pref_ae = start_pref_ae
         self.limit_pref_ae = limit_pref_ae
         self.enable_atom_ener_coeff = enable_atom_ener_coeff
+        self.intensive = intensive
         self.has_e = self.start_pref_e != 0.0 or self.limit_pref_e != 0.0
         self.has_fr = self.start_pref_fr != 0.0 or self.limit_pref_fr != 0.0
         self.has_fm = self.start_pref_fm != 0.0 or self.limit_pref_fm != 0.0
@@ -117,6 +126,10 @@ def call(
         loss = 0
         more_loss = {}
         atom_norm = 1.0 / natoms
+        # Normalization exponent controls loss scaling with system size:
+        # - norm_exp=2 (intensive=True): loss uses 1/N² scaling, making it independent of system size
+        # - norm_exp=1 (intensive=False, legacy): loss uses 1/N scaling, which varies with system size
+        norm_exp = 2 if self.intensive else 1
 
         if self.has_e:
             energy_pred = model_dict["energy"]
@@ -130,7 +143,7 @@ def call(
                 energy_pred = xp.sum(atom_ener_coeff * atom_ener_pred, axis=1)
             if self.loss_func == "mse":
                 l2_ener_loss = xp.mean(xp.square(energy_pred - energy_label))
-                loss += atom_norm * (pref_e * l2_ener_loss)
+                loss += atom_norm**norm_exp * (pref_e * l2_ener_loss)
                 more_loss["rmse_e"] = self.display_if_exist(
                     xp.sqrt(l2_ener_loss) * atom_norm, find_energy
                 )
@@ -238,7 +251,7 @@ def call(
             diff_v = virial_label - virial_pred
             if self.loss_func == "mse":
                 l2_virial_loss = xp.mean(xp.square(diff_v))
-                loss += atom_norm * (pref_v * l2_virial_loss)
+                loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
                 more_loss["rmse_v"] = self.display_if_exist(
                     xp.sqrt(l2_virial_loss) * atom_norm, find_virial
                 )
@@ -326,7 +339,7 @@ def serialize(self) -> dict:
         """Serialize the loss module."""
         return {
             "@class": "EnergySpinLoss",
-            "@version": 1,
+            "@version": 2,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -340,12 +353,17 @@ def serialize(self) -> dict:
             "limit_pref_ae": self.limit_pref_ae,
             "enable_atom_ener_coeff": self.enable_atom_ener_coeff,
             "loss_func": self.loss_func,
+            "intensive": self.intensive,
         }
 
     @classmethod
     def deserialize(cls, data: dict) -> "EnergySpinLoss":
         """Deserialize the loss module."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 1, 1)
+        version = data.pop("@version")
+        check_version_compatibility(version, 2, 1)
         data.pop("@class")
+        # Backward compatibility: version 1 used legacy normalization
+        if version < 2:
+            data.setdefault("intensive", False)
         return cls(**data)

deepmd/pd/loss/ener.py

@@ -61,6 +61,7 @@ def __init__(
         use_huber: bool = False,
         huber_delta: float | list[float] = 0.01,
         f_use_norm: bool = False,
+        intensive: bool = False,
         **kwargs: Any,
     ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
@@ -119,6 +120,13 @@ def __init__(
         f_use_norm : bool
             If True, use L2 norm of force vectors for loss calculation.
             Not implemented in PD backend, only for serialization compatibility.
+        intensive : bool
+            If true, energy and virial losses are computed as intensive quantities,
+            normalized by the square of the number of atoms (1/N^2). This ensures the loss
+            value is independent of system size and consistent with per-atom RMSE reporting.
+            If false (default), uses the legacy normalization (1/N), which may cause the loss to scale
+            with system size. The default is false for backward compatibility with models trained
+            using deepmd-kit <= 3.0.1.
         **kwargs
             Other keyword arguments.
         """
@@ -161,6 +169,7 @@ def __init__(
         self.inference = inference
         self.use_huber = use_huber
         self.huber_delta = huber_delta
+        self.intensive = intensive
         (
             self._huber_delta_energy,
             self._huber_delta_force,
@@ -218,6 +227,10 @@ def forward(
         # more_loss['log_keys'] = []  # showed when validation on the fly
         # more_loss['test_keys'] = []  # showed when doing dp test
         atom_norm = 1.0 / natoms
+        # Normalization exponent controls loss scaling with system size:
+        # - norm_exp=2 (intensive=True): loss uses 1/N² scaling, making it independent of system size
+        # - norm_exp=1 (intensive=False, legacy): loss uses 1/N scaling, which varies with system size
+        norm_exp = 2 if self.intensive else 1
         if self.has_e and "energy" in model_pred and "energy" in label:
             energy_pred = model_pred["energy"]
             energy_label = label["energy"]
@@ -243,7 +256,7 @@ def forward(
                         l2_ener_loss.detach(), find_energy
                     )
                 if not self.use_huber:
-                    loss += atom_norm**2 * (pref_e * l2_ener_loss)
+                    loss += atom_norm**norm_exp * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * energy_pred,
@@ -414,7 +427,7 @@ def forward(
                     l2_virial_loss.detach(), find_virial
                 )
             if not self.use_huber:
-                loss += atom_norm**2 * (pref_v * l2_virial_loss)
+                loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
             else:
                 l_huber_loss = custom_huber_loss(
                     atom_norm * model_pred["virial"].reshape([-1]),
@@ -564,7 +577,7 @@ def serialize(self) -> dict:
         """
         return {
             "@class": "EnergyLoss",
-            "@version": 2,
+            "@version": 3,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -585,6 +598,7 @@ def serialize(self) -> dict:
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
             "f_use_norm": self.f_use_norm,
+            "intensive": self.intensive,
         }
 
     @classmethod
@@ -602,8 +616,12 @@ def deserialize(cls, data: dict) -> "TaskLoss":
             The deserialized loss module
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        version = data.pop("@version")
+        check_version_compatibility(version, 3, 1)
         data.pop("@class")
+        # Handle backward compatibility for older versions without intensive
+        if version < 3:
+            data.setdefault("intensive", False)
         return cls(**data)
 
 

deepmd/pt/loss/ener.py

@@ -61,6 +61,7 @@ def __init__(
         use_huber: bool = False,
         f_use_norm: bool = False,
         huber_delta: float | list[float] = 0.01,
+        intensive: bool = False,
         **kwargs: Any,
     ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
@@ -120,6 +121,13 @@ def __init__(
             The threshold delta (D) used for Huber loss, controlling transition between
             L2 and L1 loss. It can be either one float shared by all terms or a list of
             three values ordered as [energy, force, virial].
+        intensive : bool
+            If true, energy and virial losses are computed as intensive quantities,
+            normalized by the square of the number of atoms (1/N^2). This ensures the loss
+            value is independent of system size and consistent with per-atom RMSE reporting.
+            If false (default), uses the legacy normalization (1/N), which may cause the loss to scale
+            with system size. The default is false for backward compatibility with models trained
+            using deepmd-kit <= 3.0.1.
         **kwargs
             Other keyword arguments.
         """
@@ -163,6 +171,7 @@ def __init__(
         self.inference = inference
         self.use_huber = use_huber
         self.f_use_norm = f_use_norm
+        self.intensive = intensive
         if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"):
             raise RuntimeError(
                 "f_use_norm can only be True when use_huber or loss_func='mae'."
@@ -225,6 +234,10 @@ def forward(
         # more_loss['log_keys'] = []  # showed when validation on the fly
         # more_loss['test_keys'] = []  # showed when doing dp test
         atom_norm = 1.0 / natoms
+        # Normalization exponent controls loss scaling with system size:
+        # - norm_exp=2 (intensive=True): loss uses 1/N² scaling, making it independent of system size
+        # - norm_exp=1 (intensive=False, legacy): loss uses 1/N scaling, which varies with system size
+        norm_exp = 2 if self.intensive else 1
         if self.has_e and "energy" in model_pred and "energy" in label:
             energy_pred = model_pred["energy"]
             energy_label = label["energy"]
@@ -250,7 +263,7 @@ def forward(
                         l2_ener_loss.detach(), find_energy
                     )
                 if not self.use_huber:
-                    loss += atom_norm**2 * (pref_e * l2_ener_loss)
+                    loss += atom_norm**norm_exp * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * energy_pred,
@@ -432,7 +445,7 @@ def forward(
                         l2_virial_loss.detach(), find_virial
                     )
                 if not self.use_huber:
-                    loss += atom_norm**2 * (pref_v * l2_virial_loss)
+                    loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * model_pred["virial"].reshape(-1),
@@ -599,7 +612,7 @@ def serialize(self) -> dict:
         """
         return {
             "@class": "EnergyLoss",
-            "@version": 2,
+            "@version": 3,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -620,6 +633,7 @@ def serialize(self) -> dict:
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
             "f_use_norm": self.f_use_norm,
+            "intensive": self.intensive,
         }
 
     @classmethod
@@ -637,8 +651,12 @@ def deserialize(cls, data: dict) -> "TaskLoss":
             The deserialized loss module
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        version = data.pop("@version")
+        check_version_compatibility(version, 3, 1)
         data.pop("@class")
+        # Handle backward compatibility for older versions without intensive
+        if version < 3:
+            data.setdefault("intensive", False)
         return cls(**data)