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

[kavanase]

I would say this is an open question in the field, of how best to train MLIPs on defect structures. The main issue is that atoms near the defect usually make up a small fraction of your dataset (and thus you are likely in the 'low-data' regime for these configurations).
You could try some form of loss re-weighting to try and mitigate this, but I haven't seen any established successful approaches for this.