MACE

[abhshkdz]

Hey all, here's an initial attempt at integrating the MACE model into the OCP codebase.

Checklist of some of the code changes and details we discussed:

• Layer-wise weight decay with default set to 5.e-7 (following this)
• EMA with default decay set to 0.99 (following this)
• AdamW with amsgrad = True, sweeping over initial_lr = [1.e-2, 5.e-3, 1.e-3]
• Gradients clipped to max norm = 10 (following this)
• Switched off ocp normalize_labels since the ScaleShiftBlock in MACE model does the same thing.
• Reduce weight decay to 1.e-10 or lower.
• Correct scale and shift statistics. Scale should be RMS forces, and shift should be mean energy per atom, both computed on the training set.
• SWA scheduler with energy / force loss coefficient switch
• E0 initialized to per-atom fitted energies (currently initialized to 0s)
• Differences b/w the ocp force loss (mean L2) vs. mace force loss (MSE).

On code organization, I've tried to use as much of ACEsuit/mace-core as possible, but there's still quite a bit of duplication between this and ACEsuit/mace, which is non-ideal. Let me know if you guys have thoughts on how to better organize this and/or feel free to directly push to this branch.

I've kicked off a few training runs on the OC20 S2EF-2M set with the above changes. Would be great to also establish parity on the smaller dataset.