Low RMSE Error for Stress, but failed Parity Plot

[mstapelberg]

Hi,

I am finding difficulties training a MACE potential with Energies, Forces, and Stresses from scratch with the following command:

#!/bin/bash
#
#SBATCH --job-name=mace_vcrtiwzr_gnome_15k_real_stress
#SBATCH --output=dist-%j.out
#SBATCH --error=dist-%j.err
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=10
#SBATCH --nodes=1
#SBATCH --gres=gpu:4
#SBATCH --constraint=rtx6000 # or v100s
#SBATCH --time=5-06:00:00
#SBATCH -p regular

# Load Conda environment
source /home/myless/.mambaforge/etc/profile.d/conda.sh

# Activate the 'allegro' environment
conda activate mace-11.7

cd "$SLURM_SUBMIT_DIR"

export OMP_NUM_THREADS=10
#--wandb_log_hypers='name model num_interactions num_channels max_L correlation E0s r_max lr train_file valid_file test_dir statistics_file num_workers batch_size valid_batch_size max_num_epochs loss error_table default_dtype device seed' \
srun python /home/myless/Packages/mace_all/mace/scripts/run_train.py \
    --name='vcrtiwzr_vac_real_stress_e1_f10_s100' \
    --model='MACE' \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=2 \
    --correlation=3 \
    --E0s='average' \
    --loss='weighted' \
    --forces_weight=10.0 \
    --energy_weight=1.0 \
    --stress_weight=100.0 \
    --compute_stress=True \
    --compute_forces=True \
    --r_max=5.0 \
    --lr=0.005 \
    --train_file='./processed_data_stress/train' \
    --valid_file='./processed_data_stress/val' \
    --test_dir='./processed_data_stress/test' \
    --statistics_file='./processed_data_stress/statistics.json' \
    --swa \
    --start_swa=400 \
    --swa_energy_weight=1000 \
    --num_workers=8 \
    --batch_size=8 \
    --valid_batch_size=8 \
    --max_num_epochs=600 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --scheduler='ReduceLROnPlateau' \
    --scheduler_patience=50 \
    --lr_scheduler_gamma=0.99 \
    --error_table='PerAtomRMSEstressvirials' \
    --default_dtype='float32' \
    --device='cuda' \
    --distributed \
    --restart_latest \
    --seed=2222 \
    --wandb \ # removed subsequent wandb data for anonymity

My training/validation results are the following:

255 2024-08-05 08:26:30.947 INFO: Epoch 498: loss=0.0027, RMSE_E_per_atom=0.5 meV, RMSE_F=14.4 meV / A, RMSE_stress_per_atom=0.2 meV / A^3
256 2024-08-05 08:30:46.659 INFO: Training complete
257 2024-08-05 08:30:46.661 INFO: Computing metrics for training, validation, and test sets
258 2024-08-05 08:30:46.666 INFO: Loading checkpoint: checkpoints/vcrtiwzr_vac_real_stress_e1_f10_s100_run-2222_epoch-396.pt
259 2024-08-05 08:30:47.671 INFO: Loaded model from epoch 396
260 2024-08-05 08:30:47.672 INFO: Evaluating train ...
261 2024-08-05 08:32:15.036 INFO: Evaluating valid ...
262 2024-08-05 08:32:25.561 INFO:
263 +-------------+---------------------+------------------+-------------------+---------------------------------------+
264 | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
265 +-------------+---------------------+------------------+-------------------+---------------------------------------+
266 |    train    |         3.4         |       15.8       |       28.74       |                  28.9                 |
267 |    valid    |         3.4         |       17.0       |       30.42       |                  29.4                 |
268 +-------------+---------------------+------------------+-------------------+---------------------------------------+
269 2024-08-05 08:32:25.562 INFO: Saving model to checkpoints/vcrtiwzr_vac_real_stress_e1_f10_s100_run-2222.model
270 2024-08-05 08:32:25.916 INFO: Compiling model, saving metadata to vcrtiwzr_vac_real_stress_e1_f10_s100_compiled.model
271 2024-08-05 08:32:27.233 INFO: Loading checkpoint: checkpoints/vcrtiwzr_vac_real_stress_e1_f10_s100_run-2222_epoch-498_swa.pt
272 /home/myless/.mambaforge/envs/mace-11.7/lib/python3.11/site-packages/torch/jit/_check.py:172: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
273   warnings.warn("The TorchScript type system doesn't support "
274 2024-08-05 08:32:28.329 INFO: Loaded model from epoch 498
275 2024-08-05 08:32:28.329 INFO: Evaluating train ...
276 2024-08-05 08:33:55.577 INFO: Evaluating valid ...
277 2024-08-05 08:34:06.096 INFO:
278 +-------------+---------------------+------------------+-------------------+---------------------------------------+
279 | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
280 +-------------+---------------------+------------------+-------------------+---------------------------------------+
281 |    train    |         0.4         |       12.8       |       23.23       |                  12.4                 |
282 |    valid    |         0.5         |       14.4       |       25.84       |                  13.1                 |
283 +-------------+---------------------+------------------+-------------------+---------------------------------------+
284 2024-08-05 08:34:06.096 INFO: Saving model to checkpoints/vcrtiwzr_vac_real_stress_e1_f10_s100_run-2222_stagetwo.model
285 2024-08-05 08:34:06.450 INFO: Compiling model, saving metadata vcrtiwzr_vac_real_stress_e1_f10_s100_stagetwo_compiled.model
286 2024-08-05 08:34:07.632 INFO: Done

However when making a parity plot, I get the following results (very impressive for Energy and Force, but less so for Stress):

I made sure to convert my vasp stresses from kbar to eV/Ang^3 by dividing each stress tensor component by 160.21766208.

Here's an example header for my xyz frames:

63
config_type=slice_sample REF_energy=-565.87409334 REF_stress="0.000547630010704997 9.056429741656608e-05 -8.401071283438865e-05 9.056429741656608e-05 0.0006958034373534656 -8.92535805001306e-05 -8.401071283438865e-05 -8.92535805001306e-05 0.0008651355799386785" Lattice="9.937269 0.018474 -3.50533 -4.955916 8.534349 -3.514767 -0.002201 -0.015225 10.509815" Properties=species:S:1:pos:R:3:REF_force:R:3:Z:I:1

I'm not entirely sure what is causing this issue, but I've seen a similar issue while using Nequip and Allegro. I checked my energy cutoff, and got the following results:

I've not had issues with this using CHGNet (however it was only with V-Cr-Ti, however I don't anticipate the addition of W and Zr to affect the ability to train stresses). So it's probably something I'm doing wrong with setting up these models!

Thanks again!
Myles

[pobo95]

Maybe it's because you used E0s as average. My discussion #317 can be the answer.

[gabor1]

you are trining with float32. not recommended!

[gabor1]

your relative force errors are not so great either. could be due to the float32, also could be due to your very high stress weights. if you are using the latest version of the code, a stress weight bug was fixed, so you no longer need large stress weights.

[mstapelberg]

I am indeed using the latest version of the code; I will adjust my stress weights to be more in line with my energy and force weights.
The 1000 : 10 : 100 weights for E : F : S seemed to work better anyways. Would it be wise toI perhaps try a 1:10:1 first to lower my relative force errors?
Thank you,
Myles

[gabor1]

yes I think so. we use two-phase training, so first you have higher force weights (like 1:10:1) and then for the last third of your epochs, you can increase the energy weights a lot, this gives very good energy errors without worsening the force errors too much

[mstapelberg]

Understood! Will give this a try; thank you again for your suggestions.

[gabor1]

The average E0s are unlikely to cause the problem with the stress. Although it is high recommended to use isolated atom energies as E0s, because then you are training for the binding energy, rather than the deviation from the average energy. This is especially important when you have more varied stoichiometry test data, because the average energy is then extremely dependent on the stoichiometry. Using isolated atom E0s will lead to much more stable potentials.

[mstapelberg]

Thank you very much for your reply and explanation Prof. Csanyi, I will generate isolated atom results and switch to float64 training. I'll reply here with updates in the next couple of days.

My dataset covers the V-Cr-Ti-W-Zr space (my datapoints for reference)

For additional context, I am training a model on structures, energies, forces, and stresses with a monovacancy (63 atoms/unit cell) in them for calculating the vacancy formation energy of materials in my composition space.

[ilyes319]

@mstapelberg Can you please post your full log file?

[mstapelberg]

Hi @ilyes319 , Thank you very much for your reply! I have attached both logs (my 0-500 and 500-600 epochs):
train_log_000-500.txt
train_log_500-600.txt

Thank you,
Myles

[gabor1]

also make sure that your data is k-point converged. what energy cutoff are you using?

[mstapelberg]

I selected an energy cutoff of 360 eV, using a 6 x 6 x 6 kpoint mesh. I was worried about my energy cutoff being a bit low for the stress convergence, but when I varied ecut for the stress tensors (see plot in initial post) the stresses went a bit whacky as I increased the energy cut past 450 eV.

I did a K-Point convergence in pure vanadium here:

Here is my energy cutoff convergence test (also in pure vanadium)

[mstapelberg]

One potential issue I see is perhaps my addition of elements like tungsten necessitate a larger energy cutoff (typically I use 540 eV for tungsten) than with just vanadium alone? I'll do an energy cut convergence test on a supercell with high amounts of tungsten and see if my selection of 360 eV still works.

[gabor1]

your stress curves are very disturbing, convergence should only improve with high cutoff, not become worse.

for k-point convergence again look at stress, not energy. and consider using a fixed k-point density (or spacing), rather than a fixed k-grid. typically sufficient spacings are 0.2 or lower (in VASP units)

[mstapelberg]

I see. Yes I was quite confused by this result. Thank you very much for your suggestion. I’ll take a look!

[gabor1]

one difference between mace and other, more nonlinear architectures is that if you have errors in your data (unconverted dft for example), mace will refuse to fit it well.

[mstapelberg]

I see thank you for this explanation! By unconverted do you mean for the units of stress?

I divided each stress tensor element by 160.21766208 to go from kbar to eV/angstrom^3 ; is that incorrect?

[gabor1]

sorry I meant "unconverged" but auto-correct changed it.

[mstapelberg]

Sorry do you mean unconverged from a k point, ecut, ismear, etc situation?

Or from an ionic step situation? As I was using intermediate ionic steps to train my model on out of equilibrium structures

[gabor1]

I meant k-point and ecut

[gabor1]

you should always train on out-of-equilibrium structures

[Mathieu-Istas]

I made sure to convert my vasp stresses from kbar to eV/Ang^3 by dividing each stress tensor component by 160.21766208.

Isn't the factor 1602? I get:
1 kbar = 10^8 Pa = 10^8 J / m^3 = 1 / 1.602 * 10^8 * 10^19 eV / m^3 = 1 / 1.602 * 10^8 * 10^19 * 10^-30 eV / A^3 = 1 / 1.602 * 10^-3 eV/A^3

[gabor1]

No it just changes it when I type space. But I'll explore..

[mstapelberg]

@bernstei Thank you for your suggestion! I had started this journey roughly around last October when the read function was not working with OUTCARs and made my DIY version and just stuck with it. I now will no longer do so. Thanks!

[bernstei]

I'm curious - the read function has been working for my OUTCARs (+/- a few specific bugs) for years. What was the recent failure?

[mstapelberg]

I think it was a temporary bug that took a week or so to fix. There was some parse error if I vaguely remember? Can try to find the issue itself, but by that time they had fixed it I had something that was reasonably working (hadn't needed to convert stresses at that point though!) so I had done a work around, but your suggestion works now! Thank you.

[bernstei]

Eh, no need. Just was surprised, since I thought that particular capability has been fairly stable

[mstapelberg]

Hi folks, update here on some convergence testing once again, but instead of in pure Vanadium I did it in 24V-10Cr-10Ti-10W-10Zr / V-15at%Cr-15at%Ti-15at%W-15at%Zr (the most alloyed material I would look at in my analysis)

Varying ECut:

Varying KSpacing:

I originally was using an Encut of 360 eV and a 6 x 6 x 6 kpoint mesh as that seemed sufficient for pure Vanadium. It now looks like a kspacing of 0.15 for my unit cell created a 5 x 5 x 5 grid (with 64 atoms). It seems that my energy cut was indeed too low, but my kpoints mesh was sufficient.

Sounds like I should re-generate more training data in the next couple of days and re-train a few models, including isolated atoms? If that seems to fix my issue I'll close. Thanks again for everyone's assistance!

[mstapelberg]

Hi @gabor1 @bernstei , I hope you are doing well! I have just gotten some updated results.

I have retrained a small model on ~1000 structures with higher energy cutoffs and a converged k-point mesh. I used 550 eV energy cutoff and a KSPACING = 0.15.
Energy-Cutoff:

K-Points:

Here are my logs and parity plots:

261 2024-08-15 02:58:12.617 INFO: Evaluating valid ...
262 2024-08-15 02:58:14.673 INFO:
263 +-------------+---------------------+------------------+-------------------+---------------------------------------+
264 | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
265 +-------------+---------------------+------------------+-------------------+---------------------------------------+
266 |    train    |         7.8         |       6.0        |        4.43       |                  16.2                 |
267 |    valid    |         8.8         |       7.4        |        5.72       |                  14.3                 |
268 +-------------+---------------------+------------------+-------------------+---------------------------------------+
269 2024-08-15 02:58:14.673 INFO: Saving model to checkpoints/vcrtiwzr_perf_fep_e1_f50_s10_run-2222.model
270 2024-08-15 02:58:15.286 INFO: Compiling model, saving metadata to vcrtiwzr_perf_fep_e1_f50_s10_compiled.model
271 2024-08-15 02:58:16.524 INFO: Loading checkpoint: checkpoints/vcrtiwzr_perf_fep_e1_f50_s10_run-2222_epoch-490_swa.pt
272 /home/myless/.mambaforge/envs/mace-11.7/lib/python3.11/site-packages/torch/jit/_check.py:172: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
273   warnings.warn("The TorchScript type system doesn't support "
274 2024-08-15 02:58:17.648 INFO: Loaded model from epoch 490
275 2024-08-15 02:58:17.649 INFO: Evaluating train ...
276 2024-08-15 02:58:30.078 INFO: Evaluating valid ...
277 2024-08-15 02:58:32.137 INFO:
278 +-------------+---------------------+------------------+-------------------+---------------------------------------+
279 | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
280 +-------------+---------------------+------------------+-------------------+---------------------------------------+
281 |    train    |         3.8         |       5.9        |        4.35       |                  8.2                  |
282 |    valid    |         2.8         |       7.4        |        5.78       |                  5.3                  |
283 +-------------+---------------------+------------------+-------------------+---------------------------------------+
284 2024-08-15 02:58:32.137 INFO: Saving model to checkpoints/vcrtiwzr_perf_fep_e1_f50_s10_run-2222_stagetwo.model
285 2024-08-15 02:58:32.749 INFO: Compiling model, saving metadata vcrtiwzr_perf_fep_e1_f50_s10_stagetwo_compiled.model
286 2024-08-15 02:58:33.954 INFO: Done

Full log here:
full_log_vcrtiwzr_fep.log

I haven't included isolated atoms yet, would that be the next step?

[bernstei]

Are the values in your virial/stress parity plot consistent with the errors that the MACE fit returns in the tables?

[mstapelberg]

Hello! Thank you very much for your help with this. I have updated the parity plot:

Seems like the rmse is in line with what I saw during training. For further context I had split my dataset into 80-10-10 split and this plot contains the 10% reserved for testing. It seems in line with the training + validation results. But the stresses still seem quite off... However unlike forces and energies, I don't have good intuition as for what a "good enough" stress error is. Is there perhaps a reference you'd recommend I read about to learn more and build that intuition? Thank you again for all your help Dr. Bernstein!

[bernstei]

Those virial errors are very small by my standards. Although the lack of correlation isn't encouraging, if you actually don't have any appreciable stresses, maybe the parity plot is just indicating that both calculated and references stresses are very small. You could add more large stress data to your fitting database (assuming you care about it). Take a small subset of your cells, scale them by maybe 1-3 %, check the (reference) stresses and see if it spans a much wider range. If so, you could try to fit that and hopefully you'll get a more sensible looking parity plot.

[mstapelberg]

I see! This dataset did just include ISIF=3 relaxations on perfect supercells (125 atoms) that had their atomic positions slightly distorted by up to 0.1 Angstrom. My overall end goal is to create vacancies in each of these structures and then calculate the vacancy migration energy using NEB. I then want to eventually train a model to calculate vacancy migration energy barriers.

In addition to increasing the size of my supercells, perhaps I could try creating a few vacancies and adding them to this dataset? Thanks again!
Myles

[mstapelberg]

Hello, @bernstei @gabor1 @ilyes319

I hope you both are doing well!

I have a bit of a conundrum here. I have fit MACE to the exact same data set (used same train, val, and test) that I used to fit a CHGNet model. Here is the CHGNet model's (using their standard settings, batch_size = 8, lr=0.005) Parity Plot:

Here is the MACE:

The CHGNet model is somehow able to find the stress trend? Whereas Mace gets very impressive energies and forces, but misses quite a large amount of the stresses. Is there any idea why this is happening? I imagine this is simply user error, but even after altering weights, rmax, precision, number of channels, etc I am still getting similar results.

Thank you again for all your help!
Best,
Myles

[gabor1]

yes, this is clearly not a hyper parameter issue, something is going wrong here.

[gabor1]

are your force RMSE showing the wrong units? surely not 13 eV/A

[mstapelberg]

Ah yes. Those are indeed incorrect units; they should be in meV/Angstrom. Thank you for catching that, I'll update my script accordingly!

[gabor1]

you say you have tried different weights. have you tried to increase the stress weights? how much?

Also, your stress plot shows two different clusters. can you inspect which structures are in which cluster, does that reveal any pattern?

[mstapelberg]

Also to note, the varying stress weights were done on my new dataset (let's call it VCrTiWZr Converged) that I am putting together, which is different from the CHGNet vs Mace comparison (this was an older dataset).

[gabor1]

if there was no improvement in the stresses as you increased the weights, then there is a data problem. do try fitting without stresses entirely, and please post the parity plot of stresses of that model

[mstapelberg]

Understood! To confirm, by fitting do you mean re-training without compute-stress, changing the error_table output to just look at energies and forces, and using ASE's numeric_stress feature to calculate the stress tensor with finite difference?
Then re-post a parity plot with those results?

I'll also try flipping the sign on the stress tensors and training with stresses in parallel as @bernstei suggested!

Thank you,
Myles

[gabor1]

it's best to create an input file without any stress data (or rename the stress key to gibberish in all frames, or give a gibberish key to MACE so that it doesn't find any stresses). after you train like that, you can just get the stresses from ASE as normal, you don't need finite difference. a MACE model will always give you analytical stresses even if you trained without stresses.

[mstapelberg]

Understood. I'll rename the stress_key when I pre-process my data but leave the compute_stress=true. Very handy to know that MACE can analytically calculate the stresses! Thank you again!

[bernstei]

There was a discussion above about the stress conversion factor. Are you now using ASE's built-in parsing of OUTCAR and then writing with ase.io.write to the extxyz format, or are you still doing it yourself (albeit with the correct magnitude factor)? If the latter, are you confident about the sign? Might be interesting to run, empirically, with a flipped stress sign, as that is a common error.

[gabor1]

one way we test this is that we train without stress, and then plot the predicted stress against the reference. if you see a negative correlation or an order of magnitude difference, that points to unit conversion errors

[mstapelberg]

I verified that my parser was providing the same stress out. I read in the stress tensor in kB, then divide by 1602.1766208. This gives me the same stress tensor as if I were to extract from ase (albeit the sign was flipped). I can try multiplying the stress tensors by -1 so that the compression/tension conventions are identical.

The ASE built-in Outcar parsing was about 1-2 OOM slower than my custom parser which was noticeable after about 6 hours spent reading my outcars using ase. I will finish running the ase parsing --> train on those results --> and compare.

thank you both again for your help!