Training Instability on Single-Element DFT Dataset with Variable Structure Sizes

[NoorAldinAlz]

Body:
Hi all,

I'm training a machine learning interatomic potential using a dataset, which contains 61 DFT-calculated structures of elemental (one element only). Each structure has a single frame, with total energy and atomic forces. The number of atoms varies significantly across the structures.

During training, I’m observing significant fluctuations in the energy and force loss — even though all data was generated consistently using the PBE functional.

I’ve already tried:

Lowering the learning rate

Gradually ramping the force and energy weights in the loss function

…but the loss still shows unstable behavior.

I'm looking for best practices or strategies to stabilize training when working with:

Variable-size structures
Total energy per frame (not normalized)
Single-element DFT data

Any insights into how to better handle this type of dataset ? would be greatly appreciated.

Thanks!