Relaxation discrepency with ASE when using small MACE model

[akwarii]

Description

We are currently implementing and benchmarking TorchSim workflows within atomate2 to enable batching for elastic and phonon properties. While production-grade models (like mace-medium) show excellent agreement between the ASE and TorchSim backends, we have encountered a significant divergence in relaxed structures and elastic constants when using a smaller toy/test MACE model (MACE.model from the atomate2 test suite).

Both backends use the Frechet cell filter and FixSymmetry constaint.

The whole discussion can be found here: materialsproject/atomate2#1505

Currently, I suspect this could be due to subtle differences in:

1. Internal line-search or backtracking implementations under noisy/spurious force gradients.
2. The exact coupling of strain and atomic positions during the optimization step within the cell filter.

Any help would be appreciated!

Reproduction Data

System: Silicon primitive cell (Si2)

• Initial lattice parameter $a$: 3.866975 Å

1. Optimized Lattice Parameters

• ASE (test model): $a =$ 3.801120 Å
• TorchSim (test model): $a =$ 3.851698 Å
• ASE (mace-medium): $a =$ 3.866058 Å
• TorchSim (mace-medium): $a =$ 3.865821 Å

2. Computed Elastic Constants (GPa)

Backend & Model	$C_{11}$	$C_{12}$	$C_{44}$
ASE (mace test)	9.703	9.699	0.002
TorchSim (mace test)	7.650	7.647	0.334
ASE (mace medium)	126.967	65.841	67.297
TorchSim (mace medium)	127.030	65.890	67.330

[CompRhys]

Thanks for the report, because the discrepancy for the medium model that has a much better PES is matching I am not sure whether this is a true issue. Because you mention LBFGS I would say are no doubt so numerical differences between out LBFGS here and ase, I think we made an effort to get FIRE to match exactly but on the LBFGS I think @abhijeetgangan attempted to stay as close as possible to the referenced literature implementation.

Can you retry using FIRE with the ase flavour as that would hopefully match much more closely.

[akwarii]

If I mentioned LBFGS it was an error, I'm actually using the FIRE optimizer with the default parameters (for both ASE and TorchSim) and a force tolerance of 1e-5.

[CompRhys]

FIRE doesn't have a line search fwiw. I am not sure then, @t-reents put in a large amount of work to ensure that ts fire and ase fire match as such this is probably a constraints issue. #552 and #438 are where fix sym was implemented and then apparent bugs were fixed.

[t-reents]

@akwarii Indeed, we did quite a bit of benchmarking back then. Could you maybe try to rerun your structures with the small model but without fixing the symmetry? It would of course be interesting and relevant to know if we potentially missed some special cases when testing the agreement if ASE and TS

[akwarii]

@t-reents I just tested disabling FixSymmetry in both ASE and TS and I'm getting the same results as before. At this point I'm starting to wonder if I overlooked something

[t-reents]

Hi @akwarii
I looked into the issue and I can reproduce it. I already had a suspicion which was now confirmed: the issue is not related to the symmetry fixing, which seem to work as expected, but rather to the convergence function that you are defining. By default, the force convergence function does not include the cell_forces:

torch-sim/torch_sim/runners.py

Lines 485 to 487 in 87dad4d

	def generate_force_convergence_fn[T: MDState | FireState](
	force_tol: float = 1e-1, *, include_cell_forces: bool = False
	) -> Callable:

This is confirmed by the following results:

model    backend / convergence         a (Ang)   steps
------------------------------------------------------
test     ASE                          3.801120      24
test     TorchSim (+cell fmax)        3.801120      24
test     TorchSim (atom-only)         3.851698       5
medium   ASE                          3.857618      38
medium   TorchSim (+cell fmax)        3.857618      38
medium   TorchSim (atom-only)         3.856473       5

As you can see, as soon as you include the cell_forces, ASE and torch-sim agree.

Script to reproduce

import numpy as np
import torch
from ase.build import bulk
from ase.filters import FrechetCellFilter
from ase.optimize import FIRE

import torch_sim as ts
from torch_sim.models.mace import MaceModel


TEST_MODEL = "atomate2/test_data/forcefields/mace/MACE.model"
DEVICE, DTYPE, FMAX = torch.device("cpu"), torch.float64, 1e-5


def si2():
    return bulk("Si", "diamond", a=3.866975 * np.sqrt(2))


def relax_ase(calc):
    atoms = si2()
    atoms.calc = calc
    dyn = FIRE(FrechetCellFilter(atoms), logfile=None)
    dyn.run(fmax=FMAX, steps=1000)
    return atoms.cell.cellpar()[0], dyn.nsteps


def relax_ts(model, *, include_cell_forces):
    state = ts.io.atoms_to_state([si2()], device=DEVICE, dtype=DTYPE)
    state = ts.fire_init(state=state, model=model, cell_filter=ts.CellFilter.frechet)
    converged = ts.generate_force_convergence_fn(
        force_tol=FMAX, include_cell_forces=include_cell_forces
    )
    for step in range(1, 1001):
        state = ts.fire_step(state=state, model=model)
        if bool(converged(state).all()) and step >= 5:
            break
    return state.to_atoms()[0].cell.cellpar()[0], step


def load(kind):
    from mace.calculators import MACECalculator, mace_mp
    from mace.calculators.foundations_models import mace_mp as mace_mp_raw

    if kind == "medium":
        calc = mace_mp(model="medium", device="cpu", default_dtype="float64")
        raw = mace_mp_raw(
            model="medium", return_raw_model=True, device="cpu", default_dtype="float64"
        )
    else:
        calc = MACECalculator(model_paths=TEST_MODEL, device="cpu", default_dtype="float64")
        raw = torch.load(TEST_MODEL, map_location=DEVICE)
    model = MaceModel(model=raw, device=DEVICE, dtype=DTYPE, compute_stress=True)
    return calc, model


rows = []
for kind in ("test", "medium"):
    calc, model = load(kind)
    rows.append((kind, "ASE", *relax_ase(calc)))
    rows.append((kind, "TorchSim (+cell fmax)", *relax_ts(model, include_cell_forces=True)))
    rows.append((kind, "TorchSim (atom-only)", *relax_ts(model, include_cell_forces=False)))

head = "{:<8} {:<24} {:>12} {:>7}".format("model", "backend / convergence", "a (Ang)", "steps")
print("\n" + head + "\n" + "-" * len(head))
for kind, label, a, steps in rows:
    print("{:<8} {:<24} {:>12.6f} {:>7d}".format(kind, label, a, steps))

[t-reents]

I just tested disabling FixSymmetry in both ASE and TS and I'm getting the same results as before. At this point I'm starting to wonder if I overlooked something

What a timing, I was just typing my previous comment 😄

[akwarii]

Oh I didn't even knew this argument existed thanks a lot for the help! Also, I would expect such a thing to be enabled by default if a cell filter is used but I guess it doesn't fit with your current design?