How to fine-tune for accurate NPT liquid water density?

[sunnywoo257]

Hi everyone,

I'm trying to fine-tune a MACE foundation model to reproduce accurate liquid water density in NPT MD (~300 K, 1 atm), but I'm running into difficulties and would really appreciate some advice.

Setup

• Foundation model: MACE-MP-0 (r_max = 6 Å, 2 interaction layers)
• Fine-tuning data: 50 liquid-water configurations (7 Å clusters), computed with ORCA
• MD: LAMMPS NPT, 600 H2O, checking the equilibrium density

What I've observed

Method	Water density
Raw foundation model (no fine-tuning)	~0.92 g/cm3
Naive fine-tuning, lr = 1.5e-4	0.976 g/cm3
Multihead replay, lr = 1.5e-4	0.93 g/cm3
Multihead replay, lr = 1.5e-4 (pure-water frames removed from the replay set)	0.88 g/cm3

With naive fine-tuning, I was able to raise the NPT water density by increasing the learning rate. However, once the model was in that state, the MD of the system I actually want to study — which used to run stably with the original foundation model — no longer ran stably. I believe this is catastrophic forgetting.

Multihead replay, on the other hand, keeps the MD stable, but it consistently lowers the water density to ~0.88–0.93, regardless of learning rate, replay ratio (real_pt_data_ratio_threshold 0.1–0.5), or number of epochs (30/50). I also noticed that the replay set (subsampled from the pretraining corpus) is mostly small gas-phase water clusters/dimers, so I tried removing all pure-water frames from it, but the density actually got worse (0.88).

Questions

1. Starting from the foundation model and training on my ORCA data, I want the NPT liquid water density to come close to the experimental value. What is the best way to achieve this?

2. When I continue training a model that reproduces the experimental water density well, I don't want to lose the foundation model's predictive performance on elements other than water. Is there a way to do this?

3. To achieve the above, I used a multihead approach (replay, etc.) — is this a reasonable/justified choice?

4. Ultimately, I would like a model that both predicts the experimental water density well in a water-only NPT run, AND runs NPT/NVT MD as stably as the foundation model does when my system of interest is placed in water. Could you recommend a good fine-tuning strategy for this?

Thank you in advance for any advice, and thank you for developing and maintaining MACE!

[OmriShf]

Hi! I'm not a developer or an expert in any way (I'm just stuck dealing with a very similar problem), so take my comment with a grain of salt, but I think I have a partial explanation. I assume the problems arise mostly from energy discontinuity between the water and everything else, due to the multy-fidelity of the training data. the density of water mace-mp-0 produces is the density of pbe water, to get an accurate density for water you would have to use a different functional or maybe just compress the water vary hard with the barostat? anyway, It is not actually part of the same dataset as the mp, and I dont think the water samples give enough information for the model to extrapolate the energy rescale of the head to other elements. Maybe you could consider using MACE-MH-1 (r2SCAN) instead, I think r2scan water are supposed to be better described (but i'm not sure, i haven't done it myself yet). If that doesn't work the solution might be to recalculate some additional systems from the mp, espacially ones with elements you want to use, using the same functional and basis, and training a head on the water + that.

btw, could I ask how you calculate bulk water using MACE? do you use multiple gpus through lammps? I have been trying that for a while and it doesn't seem to be possible. if you are calculating droplets in vacuum note that the density should also effected by laplace pressure etc.

[sunnywoo257]

Hi!

Thank you for taking the time to reply — I really appreciate it.

Having looked into this myself, I had also suspected, as you note, that the issue stems from mixing multiple levels of functional fidelity within the training data. That said, I'm not yet convinced that recomputing my system with a matched functional/basis set and training a dedicated head under the multi-head fine-tuning scheme would fully resolve the density discrepancy — that's something I'd still want to verify empirically before drawing a conclusion.

As for the implementation side: I'm running MACE within LAMMPS via the ML-IAP interface. My usage has been limited to a single GPU, so I can't speak to the multi-GPU case, but the single-GPU workflow has been working reliably for me.

Thanks again for the thoughtful response — I'll keep this in mind going forward.