Shift and scale not used from arguments when fine-tuning foundation model

[pimdh]

Describe the bug
The code of configure_model suggests that when fine-tuning a foundation model, the atomic_inter_shift should be taken from args.mean. However, this is not what currently happens, due this call having default kwarg use_shift=True, use_scale=True, which copies the values from the loaded foundation model.

Perhaps it is the intended behaviour to copy these values from the foundation model, but then the code in configure_model is confusing.

To Reproduce

import numpy as np
import torch
from mace.tools.model_script_utils import configure_model
from mace.tools.multihead_tools import dict_to_namespace
from mace.tools.utils import AtomicNumberTable

heads = ["default"]
mean = [0.1]
std = [1.2]
z_table = AtomicNumberTable([1, 2])
atomic_energies = np.random.rand(len(z_table))
model_foundation = torch.load("mace_agnesi_small.model")
model_config = {
    "model": "ScaleShiftMACE",
    "loss": "weighted",
    "compute_energy": True,
    "compute_forces": True,
    "compute_stress": False,
    "compute_dipole": False,
    "mean": mean,
    "std": std,
    "scaling": "rms_forces_scaling",
    "foundation_filter_elements": True,
}
model, output_args = configure_model(
    dict_to_namespace(model_config), None, atomic_energies, model_foundation, heads, z_table
)
print(model.scale_shift)

Expected behavior
When running the above script (an excerpt from run_train), I expect the ScaleShift block to have shift 0.1 (and perhaps also scale=1.2). However, it shows that it uses the scale and shift from the foundation model ScaleShiftBlock(scale=0.8042, shift=0.0000).

[ilyes319]

Hey @pimdh, indeed that is the expected behavior, though we should make it more explicit in the doc. I want to add, that if you do change functional/code, the E0s will change in the model, meaning that there is a flexiblity to change the shifts in the model.

[wcwitt]

I think this is less a matter for documentation and more that run_train should error when conflicting input is given. I've seen a version of this as well, and it's pretty confusing, e.g. one can specify what feels like a full model architecture which might then be completely overwritten if fine-tuning is used.

Hard to imagine a solution that isn't spaghetti though. Are there any good tools for automatically determining which inputs are compatible/incompatible?