Skip to content

feat(loss): support three-value huber delta #5329

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

Open
OutisLi wants to merge 2 commits into
base: master
Choose a base branch
from
Open

feat(loss): support three-value huber delta #5329

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
dbc652e
feat(loss): support three-value huber delta
OutisLi Mar 19, 2026
a7a1edd
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 20, 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
33 changes: 23 additions & 10 deletions deepmd/dpmodel/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@
from deepmd.utils.data import (
DataRequirementItem,
)
from deepmd.utils.loss import (
resolve_huber_deltas,
)
from deepmd.utils.version import (
check_version_compatibility,
)
Expand Down Expand Up @@ -74,8 +77,10 @@ class EnergyLoss(Loss):
- For absolute prediction errors within D: quadratic loss (0.5 * (error**2))
- For absolute errors exceeding D: linear loss (D * |error| - 0.5 * D)
Formula: loss = 0.5 * (error**2) if |error| <= D else D * (|error| - 0.5 * D).
huber_delta : float
The threshold delta (D) used for Huber loss, controlling transition between L2 and L1 loss.
huber_delta : float | list[float]
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].
loss_func : str
Loss function type for energy, force, and virial terms.
Options: 'mse' (Mean Squared Error, L2 loss, default) or 'mae' (Mean Absolute Error, L1 loss).
Expand All @@ -84,6 +89,7 @@ class EnergyLoss(Loss):
f_use_norm : bool
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.
**kwargs
Other keyword arguments.
"""
Expand All @@ -107,7 +113,7 @@ def __init__(
limit_pref_gf: float = 0.0,
numb_generalized_coord: int = 0,
use_huber: bool = False,
huber_delta: float = 0.01,
huber_delta: float | list[float] = 0.01,
loss_func: str = "mse",
f_use_norm: bool = False,
**kwargs: Any,
Expand Down Expand Up @@ -153,6 +159,11 @@ def __init__(
raise RuntimeError(
"f_use_norm can only be True when use_huber or loss_func='mae'."
)
(
self._huber_delta_energy,
self._huber_delta_force,
self._huber_delta_virial,
) = resolve_huber_deltas(huber_delta)
if self.use_huber and (
self.has_pf or self.has_gf or self.relative_f is not None
):
Expand Down Expand Up @@ -215,7 +226,10 @@ def call(

if self.relative_f is not None:
force_hat_3 = xp.reshape(force_hat, (-1, 3))
norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), (-1, 1)) + self.relative_f
norm_f = (
xp.reshape(xp.linalg.vector_norm(force_hat_3, axis=1), (-1, 1))
+ self.relative_f
)
diff_f_3 = xp.reshape(diff_f, (-1, 3))
diff_f_3 = diff_f_3 / norm_f
diff_f = xp.reshape(diff_f_3, (-1,))
Expand Down Expand Up @@ -250,7 +264,7 @@ def call(
l_huber_loss = custom_huber_loss(
atom_norm_ener * energy,
atom_norm_ener * energy_hat,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_e * l_huber_loss
more_loss["rmse_e"] = self.display_if_exist(
Expand All @@ -276,7 +290,7 @@ def call(
l_huber_loss = custom_huber_loss(
xp.reshape(force, (-1,)),
xp.reshape(force_hat, (-1,)),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
else:
force_diff_3 = xp.reshape(force_hat - force, (-1, 3))
Expand All @@ -286,7 +300,7 @@ def call(
l_huber_loss = custom_huber_loss(
force_diff_norm,
xp.zeros_like(force_diff_norm),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
loss += pref_f * l_huber_loss
more_loss["rmse_f"] = self.display_if_exist(
Expand Down Expand Up @@ -317,7 +331,7 @@ def call(
l_huber_loss = custom_huber_loss(
atom_norm * virial_reshape,
atom_norm * virial_hat_reshape,
delta=self.huber_delta,
delta=self._huber_delta_virial,
)
loss += pref_v * l_huber_loss
more_loss["rmse_v"] = self.display_if_exist(
Expand All @@ -336,7 +350,6 @@ def call(
if self.has_ae:
atom_ener_reshape = xp.reshape(atom_ener, (-1,))
atom_ener_hat_reshape = xp.reshape(atom_ener_hat, (-1,))

if self.loss_func == "mse":
l2_atom_ener_loss = xp.mean(
xp.square(atom_ener_hat_reshape - atom_ener_reshape),
Expand All @@ -347,7 +360,7 @@ def call(
l_huber_loss = custom_huber_loss(
atom_ener_reshape,
atom_ener_hat_reshape,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_ae * l_huber_loss
more_loss["rmse_ae"] = self.display_if_exist(
Expand Down
24 changes: 17 additions & 7 deletions deepmd/pd/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@
from deepmd.utils.data import (
DataRequirementItem,
)
from deepmd.utils.loss import (
resolve_huber_deltas,
)
from deepmd.utils.version import (
check_version_compatibility,
)
Expand Down Expand Up @@ -56,7 +59,7 @@ def __init__(
loss_func: str = "mse",
inference: bool = False,
use_huber: bool = False,
huber_delta: float = 0.01,
huber_delta: float | list[float] = 0.01,
f_use_norm: bool = False,
**kwargs: Any,
) -> None:
Expand Down Expand Up @@ -109,8 +112,10 @@ def __init__(
- For absolute prediction errors within D: quadratic loss (0.5 * (error**2))
- For absolute errors exceeding D: linear loss (D * |error| - 0.5 * D)
Formula: loss = 0.5 * (error**2) if |error| <= D else D * (|error| - 0.5 * D).
huber_delta : float
The threshold delta (D) used for Huber loss, controlling transition between L2 and L1 loss.
huber_delta : float | list[float]
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].
f_use_norm : bool
If True, use L2 norm of force vectors for loss calculation.
Not implemented in PD backend, only for serialization compatibility.
Expand Down Expand Up @@ -156,6 +161,11 @@ def __init__(
self.inference = inference
self.use_huber = use_huber
self.huber_delta = huber_delta
(
self._huber_delta_energy,
self._huber_delta_force,
self._huber_delta_virial,
) = resolve_huber_deltas(huber_delta)
if self.use_huber and (
self.has_pf or self.has_gf or self.relative_f is not None
):
Expand Down Expand Up @@ -238,7 +248,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_norm * energy_pred,
atom_norm * energy_label,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_e * l_huber_loss
rmse_e = l2_ener_loss.sqrt() * atom_norm
Expand Down Expand Up @@ -305,7 +315,7 @@ def forward(
l_huber_loss = custom_huber_loss(
force_pred.reshape([-1]),
force_label.reshape([-1]),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
loss += pref_f * l_huber_loss
rmse_f = l2_force_loss.sqrt()
Expand Down Expand Up @@ -409,7 +419,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_norm * model_pred["virial"].reshape([-1]),
atom_norm * label["virial"].reshape([-1]),
delta=self.huber_delta,
delta=self._huber_delta_virial,
)
loss += pref_v * l_huber_loss
rmse_v = l2_virial_loss.sqrt() * atom_norm
Expand Down Expand Up @@ -440,7 +450,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_ener_reshape,
atom_ener_label_reshape,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_ae * l_huber_loss
rmse_ae = l2_atom_ener_loss.sqrt()
Expand Down
27 changes: 19 additions & 8 deletions deepmd/pt/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@
from deepmd.utils.data import (
DataRequirementItem,
)
from deepmd.utils.loss import (
resolve_huber_deltas,
)
from deepmd.utils.version import (
check_version_compatibility,
)
Expand Down Expand Up @@ -57,7 +60,7 @@ def __init__(
inference: bool = False,
use_huber: bool = False,
f_use_norm: bool = False,
huber_delta: float = 0.01,
huber_delta: float | list[float] = 0.01,
**kwargs: Any,
) -> None:
r"""Construct a layer to compute loss on energy, force and virial.
Expand Down Expand Up @@ -112,8 +115,11 @@ def __init__(
f_use_norm : bool
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.
huber_delta : float
The threshold delta (D) used for Huber loss, controlling transition between L2 and L1 loss.
This treats the force vector as a whole rather than three independent components.
huber_delta : float | list[float]
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].
**kwargs
Other keyword arguments.
"""
Expand Down Expand Up @@ -162,6 +168,11 @@ def __init__(
"f_use_norm can only be True when use_huber or loss_func='mae'."
)
self.huber_delta = huber_delta
(
self._huber_delta_energy,
self._huber_delta_force,
self._huber_delta_virial,
) = resolve_huber_deltas(huber_delta)
if self.use_huber and (
self.has_pf or self.has_gf or self.relative_f is not None
):
Expand Down Expand Up @@ -244,7 +255,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_norm * energy_pred,
atom_norm * energy_label,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_e * l_huber_loss
rmse_e = l2_ener_loss.sqrt() * atom_norm
Expand Down Expand Up @@ -308,7 +319,7 @@ def forward(
l_huber_loss = custom_huber_loss(
force_pred.reshape(-1),
force_label.reshape(-1),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
else:
force_diff_norm = torch.linalg.vector_norm(
Expand All @@ -320,7 +331,7 @@ def forward(
l_huber_loss = custom_huber_loss(
force_diff_norm,
torch.zeros_like(force_diff_norm),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
loss += pref_f * l_huber_loss
rmse_f = l2_force_loss.sqrt()
Expand Down Expand Up @@ -426,7 +437,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_norm * model_pred["virial"].reshape(-1),
atom_norm * label["virial"].reshape(-1),
delta=self.huber_delta,
delta=self._huber_delta_virial,
)
loss += pref_v * l_huber_loss
rmse_v = l2_virial_loss.sqrt() * atom_norm
Expand Down Expand Up @@ -474,7 +485,7 @@ def forward(
l_huber_loss = custom_huber_loss(
atom_ener_reshape,
atom_ener_label_reshape,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_ae * l_huber_loss
rmse_ae = l2_atom_ener_loss.sqrt()
Expand Down
24 changes: 17 additions & 7 deletions deepmd/tf/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
from deepmd.utils.data import (
DataRequirementItem,
)
from deepmd.utils.loss import (
resolve_huber_deltas,
)
from deepmd.utils.version import (
check_version_compatibility,
)
Expand Down Expand Up @@ -87,8 +90,10 @@ class EnerStdLoss(Loss):
- For absolute prediction errors within D: quadratic loss (0.5 * (error**2))
- For absolute errors exceeding D: linear loss (D * |error| - 0.5 * D)
Formula: loss = 0.5 * (error**2) if |error| <= D else D * (|error| - 0.5 * D).
huber_delta : float
The threshold delta (D) used for Huber loss, controlling transition between L2 and L1 loss.
huber_delta : float | list[float]
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].
loss_func : str
Loss function type. Options: 'mse' or 'mae'.
Not implemented in TF backend, only for serialization compatibility.
Expand Down Expand Up @@ -118,7 +123,7 @@ def __init__(
limit_pref_gf: float = 0.0,
numb_generalized_coord: int = 0,
use_huber: bool = False,
huber_delta: float = 0.01,
huber_delta: float | list[float] = 0.01,
loss_func: str = "mse",
f_use_norm: bool = False,
**kwargs: Any,
Expand Down Expand Up @@ -162,6 +167,11 @@ def __init__(
)
self.use_huber = use_huber
self.huber_delta = huber_delta
(
self._huber_delta_energy,
self._huber_delta_force,
self._huber_delta_virial,
) = resolve_huber_deltas(huber_delta)
if self.use_huber and (
self.has_pf or self.has_gf or self.relative_f is not None
):
Expand Down Expand Up @@ -351,7 +361,7 @@ def build(
l_huber_loss = custom_huber_loss(
atom_norm_ener * energy,
atom_norm_ener * energy_hat,
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_e * l_huber_loss
more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy)
Expand All @@ -362,7 +372,7 @@ def build(
l_huber_loss = custom_huber_loss(
tf.reshape(force, [-1]),
tf.reshape(force_hat, [-1]),
delta=self.huber_delta,
delta=self._huber_delta_force,
)
loss += pref_f * l_huber_loss
more_loss["l2_force_loss"] = self.display_if_exist(
Expand All @@ -375,7 +385,7 @@ def build(
l_huber_loss = custom_huber_loss(
atom_norm * tf.reshape(virial, [-1]),
atom_norm * tf.reshape(virial_hat, [-1]),
delta=self.huber_delta,
delta=self._huber_delta_virial,
)
loss += pref_v * l_huber_loss
more_loss["l2_virial_loss"] = self.display_if_exist(
Expand All @@ -388,7 +398,7 @@ def build(
l_huber_loss = custom_huber_loss(
tf.reshape(atom_ener, [-1]),
tf.reshape(atom_ener_hat, [-1]),
delta=self.huber_delta,
delta=self._huber_delta_energy,
)
loss += pref_ae * l_huber_loss
more_loss["l2_atom_ener_loss"] = self.display_if_exist(
Expand Down
8 changes: 6 additions & 2 deletions deepmd/utils/argcheck.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3108,7 +3108,11 @@ def loss_ener() -> list[Argument]:
"- For absolute errors exceeding D: linear loss D * (\\|error\\| - 0.5 * D) \n\n"
"Formula: loss = 0.5 * (error**2) if \\|error\\| <= D else D * (\\|error\\| - 0.5 * D). "
)
doc_huber_delta = "The threshold delta (D) used for Huber loss, controlling transition between L2 and L1 loss. "
doc_huber_delta = (
"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]. "
)
doc_loss_func = (
"Loss function type for energy, force, and virial terms. "
"Options: 'mse' (Mean Squared Error, L2 loss, default) or 'mae' (Mean Absolute Error, L1 loss). "
Expand Down Expand Up @@ -3258,7 +3262,7 @@ def loss_ener() -> list[Argument]:
),
Argument(
"huber_delta",
float,
[float, list[float]],
optional=True,
default=0.01,
doc=doc_huber_delta,
Expand Down
Loading
Loading