Force prediction issues near defects

[BorisWasBezet]

Dear,

I’m using NequIP to train a model specifically for predicting forces in defected structures. While the model performs well on non-defected atoms, the accuracy of force predictions on atoms near the defect is significantly lower.

I’m already using a higher weight for the force RMSE in the loss function:

loss:
target: nequip.train.EnergyForceMetrics
coeffs:
per_atom_energy_rmse: 1.0
forces_rmse: 5.0

However, this doesn’t seem to sufficiently improve predictions for the defect region.
Is there a recommended way to emphasize or reweight the contribution of atoms near a defect during training?

Alternatively, is there a loss function or training strategy better suited for systems with localized anomalies like defects?

Any advice or guidance would be greatly appreciated.

Kind regards,
Boris