Finetuning on PET model and rotational symmetrization

[dominicvarghese]

Hi everyone,

I am applying a foundational-model + fine-tuning workflow to KTaO₃ and I am seeing an unexpected structural distortion after fine-tuning. I would appreciate any insight into whether this behavior is expected or if I may be making a mistake in my setup.

Using VASP with the PBEsol functional, I first relaxed cubic KTaO₃ and obtained a lattice parameter of 3.989 Å.
When I relaxed the structure using the pretrained OMAD-L model, the lattice parameter was 3.985 Å, which is in good agreement with DFT.

I then fine-tuned the MAD-S model using the following YAML configuration:

architecture:
  name: pet
  training:
    batch_size: 4
    num_epochs: 30
    learning_rate: 1e-3
    warmup_fraction: 0.01
    finetune:
      method: full
      read_from: pet-mad-s-v1.1.0.ckpt

training_set:
  systems:
    read_from: ../KTO/PBEsol_train_PBE/KTO_training_data_complete.xyz
    reader: ase
    length_unit: angstrom
  targets:
    energy/finetune:
      quantity: energy
      key: energy
      unit: eV
      per_atom: false

      forces:
        key: forces
       
      virial:
        key: stress

validation_set: 0.1
test_set: 0.1

The training metrics were:

Train;
Energy RMSE: 20.4 meV/atom
Forces RMSE: 26.1 meV/Å
Virial RMSE: 26.9 meV/atom

Validation
Energy RMSE: 19.4 meV/atom
Forces RMSE: 29.5 meV/Å
Virial RMSE: 31.3 meV/atom

Test
Energy RMSE: 17.6 meV/atom
Forces RMSE: 30.4 meV/Å
Virial RMSE: 26.5 meV/atom

However, when I use this fine-tuned model to relax KTaO₃, the structure no longer remains cubic. Instead, it relaxes to a rhombohedral structure with a lattice parameter of approximately 5.25 Å and bond angles changing from 90° to ~62.5°.

Is there a mistake I am making in this fine-tuning process?

Thanks
Dominic

[abmazitov]

Thank you for sharing your issue, @dominicvarghese! Before going deep into details, could you tell what script do you use to run the relaxation after fine-tuning the model?

[ceriottm]

Hm. the large-ish errors, that are the same for train/val/test and for small and large models stink of some inconsistency in the data / unit problems. Try to look at the parity plots and identify if there are patterns that can help identify the issue.

[dominicvarghese]

Thanks for the reply. There seems to be some inconsistency with the data. This is the parity plot I got:

The different lines in the energy plot says there is some error in the dataset. For creating the fine-tuning data, I did DFT calculations with VASP PBEsol potentials and did NPT/NVT runs for different isotropic strains at few different temperatures.
This is the Energy vs Volume plot showing the different configurations covered in the training set.

The data seems consistent as the VASP tags involved in DFT and AIMD were same. Also, I did use the same dataset to fine-tune on MACE MATPES-R2SCAN model and that gave me small error with training for energy(RMSE: 2 meV/atom) and force (RMSE: 16 meV/A). I am also attaching the parity plots, I got from that training as well here:

Could it be I am missing some tags in the fine-tuning part?

[ceriottm]

I have a suspicion. Try to use stress: stress rather than "virial" in the input file. stress and virial are two different things, and there's a lot of confusion about the definition (basically one is scaled by volume and has the units of a pressure, the other by number of atoms and has the units of an energy) and depending on what you have computed you might have chosen the wrong convention. since the key is called "stress" I'd assume you've the stress there and not the virial. This would explain the shift in energy values for different cell volumes.

[dominicvarghese]

Thanks for the suggestion. Changing the virial to stress indeed fixed the energy errors.

Now the errors are,

Training:  energy/finetune RMSE (per atom): 0.87173 meV 
                forces[energy/finetune] RMSE: 15.195 meV/A
                virial[energy/finetune] RMSE (per atom): 2.5424 meV

Validation:  energy/finetune RMSE (per atom): 0.97406 meV 
                  forces[energy/finetune] RMSE: 15.479 meV/A 
                  virial[energy/finetune] RMSE (per atom): 2.2099 meV

Testing: energy/finetune RMSE (per atom): 0.84885 meV 
             forces[energy/finetune] RMSE: 15.674 meV/A 
             virial[energy/finetune] RMSE (per atom): 2.5176 meV 

The parity is as follows:

Thanks for the suggestions @ceriottm, @frostedoyster, @abmazitov !

[ceriottm]

@lab-cosmo/cosmonauts we need to document better the difference between stress and virial and explain when to use either.

[dominicvarghese]

Hi,

Since I have the fine-tuned model for the KTaO3 dataset I am using, I was doing the full geometry relaxation of the Pm-3m structure to get the relaxed lattice parameters using the potential.

I used the following code in ASE to do the relaxation:

import numpy as np
import torch
from ase.io import read, write
from ase.filters import FrechetCellFilter, ExpCellFilter
import ase.optimize as opt
from pymatgen.core import Structure
from ase import units
from ase.md.langevin import Langevin
import numpy as np
import time

from upet.calculator import UPETCalculator
from metatomic.torch.ase_calculator import MetatomicCalculator

# --- Parameters ---
ASE_OPTIMIZER = "BFGS"         # Faster than FIRE
ASE_FILTER = "frechet"         # Better convergence for cell relaxations
FORCE_MAX = 0.00001              # eV/Å
MAX_STEPS = 500
INPUT_FILE = "POSCAR"
OUTPUT_FILE = "fully_relaxed_kto.cif"

atoms = read(INPUT_FILE)
print(f"Read {len(atoms)} atoms from {INPUT_FILE}")

calculator = MetatomicCalculator(
    model="model-ft.pt",             # local fine-tuned file
    device="cuda",                   # Use GPU
    variants={"energy": "finetune"},
    extensions_directory=None      
)

atoms.calc = calculator

# --- Choose filter and optimizer ---
filter_cls = {"frechet": FrechetCellFilter, "exp": ExpCellFilter}[ASE_FILTER]
optimizer_cls = {
    "GOQN": opt.GoodOldQuasiNewton,
    "FIRE": opt.fire.FIRE,
    "BFGS": opt.BFGS,
    "LBFGS": opt.LBFGS,
}[ASE_OPTIMIZER]

# --- Run relaxation ---
ucf = filter_cls(atoms)
print("\nStarting full structural relaxation...")
with optimizer_cls(ucf, logfile="relax.log") as optimizer:
    optimizer.run(fmax=FORCE_MAX, steps=MAX_STEPS)

# --- Save results ---
relaxed_atoms = getattr(ucf, "atoms", atoms)
write("fully_relaxed_kto.cif", relaxed_atoms, format="cif")
write("CONTCAR-finetune", relaxed_atoms, format="vasp", vasp5=True, direct = True)
print("\nRelaxation complete. Structures saved to CONTCAR")

When I use the standard structure from Materials Project for relaxation;

Input:

Values of a,b,c,alpha,beta,gamma:
3.99488 3.99488 3.99488 90.00000 90.00000 90.00000

The relaxed output is as follows:

Values of a,b,c,alpha,beta,gamma:
3.98954 3.98962 3.98971 89.99949 90.00039 89.99819

Using the default tolerance in FINDSYM to check the symmetry gives a spacegroup different from the Pm-3m.

Is there anything that needs to be added in the relaxation script so that I can get the perfect cubic structure with a = b = c, alpha = beta = gamma = 90?

[ppegolo]

Hi Dominic,

This behavior is quite expected because the architecture of PET is unconstrained and you're relaxing the cell without any symmetry constraint (have a look here for further details and some comments). If you want to keep a symmetric structure, you should constrain the relaxation with something like:

from ase.filters import FixSymmetry

atoms.set_constraint(FixSymmetry(atoms))
ucf = FrechetCellFilter(atoms, mask=[True]*3 + [False]*3)

# optimize as usual
...

atoms.set_constraint()  # remove symmetry constraints

This way, the space group will not change during relaxation.
Otherwise, if you want to let it fully relax and find the final space group, you usually need to increase a bit the tolerance of the symmetry searcher, because PET will always break exact symmetry (this has really little effect on actual properties, but I understand that having a symmetric structure might be useful in some cases).

[dominicvarghese]

Hi everyone,

I am currently fine-tuning a foundation model (e.g., PET-OAM, trained on PBE) using a smaller dataset calculated with a different functional (e.g., PBEsol).
In other architectures like MACE, this domain shift is typically handled via refitting $E_0$ - subtracting the specific atomic self-energies for the new functional. Could you confirm if the PET architecture (specifically within the metatomic implementation) follows a similar approach? In the yaml file used in fine-tuning, we dont specify any atomisation energy anywhere.
I want to ensure I am correctly accounting for the energy baseline shift between the foundation model's training data and my PBEsol dataset.

[dominicvarghese]

Thanks for the reply. To clarify, in my calculations, I am inheriting the head weights of the foundational model.

I also have a question regarding an observation during relaxations. I apply a periodic displacement to the atoms, fix one set of atoms, and relax the remaining ones to induce a strain gradient. Ideally, applying equal but opposite displacements should result in the exact same dipole moment (in magnitude).

However, when I do this with the PET model, these dipole values differ slightly. When I run the exact same setup with a strictly equivariant model (like MACE), the values match perfectly.

Is this slight difference because the force is fitted as a separate entity rather than strictly as an energy derivative in this model? Or is this just the expected numerical noise from the unconstrained architecture?

[frostedoyster]

I think this is expected due to the unconstrained nature of the model (@ppegolo can you confirm?).

In any case, you seem to care about these symmetry-breaking effects... I would then highly recommend you to try our symmetrization options: try to feed rotational_average_order=3 to the UPETCalculator and see if things improve (you might have to upgrade the upet package to the latest version). If you care about speed, you can try to lower it to 2 or 1 and/or play around with the batch size for the symmetrization (rotational_average_batch_size parameter).

Is this slight difference because the force is fitted as a separate entity rather than strictly as an energy derivative in this model?

Unrelated, but the forces are the derivatives of the energies by default. You can change that by setting non_conservative=True.

Let us know how it goes

[dominicvarghese]

Thanks for the suggestions @frostedoyster . I updated to the latest upet version(0.1.2).

What is the new command to call the fine-tuned model with the rotational_average_order key?

Earlier I was using the metatomic calculator (MetatomicCalculator(model="../../../model-ft.pt",device="cuda",variants={"energy": "finetune"},extensions)

I tried using the UPET calculator UPETCalculator(model="../../../model-ft.pt", device="cuda", variants={"energy": "finetune"}, rotational_average_order=3, rotational_average_batch_size=4), but this is not the correct way to call it(error loading models).

[frostedoyster]

I think UPETCalculator needs the .ckpt file (and not .pt).

Otherwise, if you want to stick to the lower-level MetatomicCalculator, you can just wrap your calculator inside SymmetrizedCalculator (documentation here)

[dominicvarghese]

UPETCalculator( checkpoint_name="../../../model-ft.ckpt",...) gave

Error loading model: UPETCalculator.init() got an unexpected keyword argument 'checkpoint_name'

and

UPETCalculator( model="../../../model-ft.ckpt",...) gave

Error loading model: Model ../../../model-ft.ckpt is not available. Please select one of the following: ['pet-mad-s', 'pet-omat-xs', 'pet-omat-s', 'pet-omat-m', 'pet-omat-l', 'pet-omat-xl', 'pet-oam-l', 'pet-oam-xl', 'pet-omad-xs', 'pet-omad-s', 'pet-omad-l', 'pet-omatpes-l', 'pet-spice-s', 'pet-spice-l']

[frostedoyster]

The argument's name is checkpoint_path. Have a look at the definition of the class
https://github.com/lab-cosmo/upet/blob/main/src/upet/calculator.py

EDIT: thanks for letting us know, the docs were wrong and we've fixed them. You can also use the calculators from metatomic, including the symmetrized one: they're more API-stable in general

[frostedoyster]

@dominicvarghese Does everything work for you now? Can we mark the discussion as resolved? You can always open a new discussion if you want to discuss things further

[dominicvarghese]

@frostedoyster Sorry for this late reply. Yes, the command was working and I am getting the same value for the symmetric structures (more time-consuming with the rotational averaging).

predictions will be averaged over a quadrature of the O(3) group based on a Lebedev grid of the specified order (here 3). Higher orders lead to more accurate equivariance, but also increase the computational cost.

Is there any resource so that I can understand more on what this means? (Is this similar to the value of l in equivariant graph models)

[frostedoyster]

Yes, they're the same notation and they represent the same thing (although of course what we're doing here is different). We essentially average over a "smart" grid of different orientations of the structure. An LLM can probably get you started on good resources to study Lebedev grids, if you're interested in the specifics of the rotational grid

[frostedoyster]

Issues solved (both virial vs stress and rotational symmetrization)
I'm marking as solved just to keep discussions as specific and easy to read as possible, please open a new one if you encounter any other issues or if you have further questions!