(v0.3.5). Implementing model="MACE", but .out file still indicates "ScaleShiftMACE"

[ankitmathanker]

I am trying to train a model using --model="MACE", but when I check the .out file while is appear that the initially it reads --model="MACE", but before run starts it is inidicating ScaleShiftMACE. I have attached folder which contains job.out file and submission script.
model_00.zip

<mace.tools.slurm_distributed.DistributedEnvironment object at 0x200197cd2b80>
2025-01-06 07:26:11.592 INFO: Process group initialized: True
2025-01-06 07:26:11.592 INFO: Processes: 1
2025-01-06 07:26:11.592 INFO: MACE version: 0.3.5
2025-01-06 07:26:11.593 INFO: Configuration: Namespace(config=None, name='MACE', seed=0, log_dir='logs', model_dir='.', checkpoints_dir='checkpoints', results_dir='results', downloads_dir='downloads', device='cuda', default_dtype='float32', distributed=True, log_level='INFO', error_table='PerAtomRMSE', model='MACE', r_max=5.0, radial_type='bessel', num_radial_basis=8, num_cutoff_basis=5, pair_repulsion=False, distance_transform='None', interaction='RealAgnosticResidualInteractionBlock', interaction_first='RealAgnosticResidualInteractionBlock', max_ell=3, correlation=3, num_interactions=2, MLP_irreps='16x0e', radial_MLP='[64, 64, 64]', hidden_irreps='64x0e + 64x1o + 64x2e', num_channels=None, max_L=None, gate='silu', scaling='rms_forces_scaling', avg_num_neighbors=1, compute_avg_num_neighbors=True, compute_stress=False, compute_forces=True, train_file='/p/gpfs1/mathanker1/model_benchmarking/avg_M_64/00_iter/00_dft/train.xyz', valid_file=None, valid_fraction=0.1, test_file='/p/gpfs1/mathanker1/model_benchmarking/avg_M_64/00_iter/00_dft/test_last.xyz', test_dir=None, multi_processed_test=False, num_workers=0, pin_memory=True, atomic_numbers=None, mean=None, std=None, statistics_file=None, E0s='average', keep_isolated_atoms=False, energy_key='energy', forces_key='forces', virials_key='virials', stress_key='stress', dipole_key='dipole', charges_key='charges', loss='weighted', forces_weight=100.0, swa_forces_weight=100.0, energy_weight=1.0, swa_energy_weight=1000.0, virials_weight=1.0, swa_virials_weight=10.0, stress_weight=1.0, swa_stress_weight=10.0, dipole_weight=1.0, swa_dipole_weight=1.0, config_type_weights='{"Default":1.0}', huber_delta=0.01, optimizer='adam', batch_size=1, valid_batch_size=10, lr=0.01, swa_lr=0.001, weight_decay=5e-07, amsgrad=True, scheduler='ReduceLROnPlateau', lr_factor=0.8, scheduler_patience=50, lr_scheduler_gamma=0.9993, swa=True, start_swa=800, ema=True, ema_decay=0.99, max_num_epochs=1000, patience=2048, foundation_model=None, foundation_model_readout=True, eval_interval=2, keep_checkpoints=False, save_all_checkpoints=True, restart_latest=True, save_cpu=False, clip_grad=10.0, wandb=False, wandb_project='', wandb_entity='', wandb_name='', wandb_log_hypers=['num_channels', 'max_L', 'correlation', 'lr', 'swa_lr', 'weight_decay', 'batch_size', 'max_num_epochs', 'start_swa', 'energy_weight', 'forces_weight'])
2025-01-06 07:26:11.593 INFO: CUDA version: 11.8, CUDA device: 0
2025-01-06 07:26:11.741 INFO: Error accessing Git repository: /p/gpfs1/mathanker1/model_benchmarking/avg_M_64/00_iter/01_train/00
2025-01-06 07:26:11.855 INFO: Since ASE version 3.23.0b1, using energy_key 'energy' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_energy'. You need to use --energy_key='REF_energy', to tell the key name chosen.
2025-01-06 07:26:11.861 INFO: Since ASE version 3.23.0b1, using forces_key 'forces' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_forces'. You need to use --forces_key='REF_forces', to tell the key name chosen.
2025-01-06 07:26:11.868 INFO: Since ASE version 3.23.0b1, using stress_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen.
2025-01-06 07:26:11.878 INFO: Loaded 21 training configurations from '/p/gpfs1/mathanker1/model_benchmarking/avg_M_64/00_iter/00_dft/train.xyz'
2025-01-06 07:26:11.878 INFO: Using random 10.0% of training set for validation
2025-01-06 07:26:11.918 INFO: Since ASE version 3.23.0b1, using energy_key 'energy' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_energy'. You need to use --energy_key='REF_energy', to tell the key name chosen.
2025-01-06 07:26:11.920 INFO: Since ASE version 3.23.0b1, using forces_key 'forces' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_forces'. You need to use --forces_key='REF_forces', to tell the key name chosen.
2025-01-06 07:26:11.922 INFO: Since ASE version 3.23.0b1, using stress_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen.
2025-01-06 07:26:11.925 INFO: Loaded 5 test configurations from '/p/gpfs1/mathanker1/model_benchmarking/avg_M_64/00_iter/00_dft/test_last.xyz'
2025-01-06 07:26:11.925 INFO: Total number of configurations: train=19, valid=2, tests=[Default: 5]
2025-01-06 07:26:11.927 INFO: AtomicNumberTable: (1, 8, 11, 17, 29, 45, 79)
2025-01-06 07:26:11.927 INFO: Atomic Energies not in training file, using command line argument E0s
2025-01-06 07:26:11.927 INFO: Computing average Atomic Energies using least squares regression
2025-01-06 07:26:11.933 INFO: Atomic energies: [-6.242603671703923, -3.1213018358519413, -4.075917353519044, -3.3008814990745616, -2.4822627893532454, -5.7818687197892, -2.233969179303094]
2025-01-06 07:26:12.124 INFO: WeightedEnergyForcesLoss(energy_weight=1.000, forces_weight=100.000)
2025-01-06 07:26:14.059 INFO: Average number of neighbors: 40.78620910644531
2025-01-06 07:26:14.059 INFO: Selected the following outputs: {'energy': True, 'forces': True, 'virials': False, 'stress': False, 'dipoles': False}
2025-01-06 07:26:14.105 INFO: Building model
2025-01-06 07:26:14.106 INFO: Hidden irreps: 64x0e + 64x1o + 64x2e
2025-01-06 07:26:16.688 INFO: Using stochastic weight averaging (after 800 epochs) with energy weight : 1000.0, forces weight : 100.0 and learning rate : 0.001
2025-01-06 07:26:16.911 WARNING: Cannot find checkpoint with tag 'MACE_run-0' in 'checkpoints'
2025-01-06 07:26:16.913 INFO: ScaleShiftMACE(
  (node_embedding): LinearNodeEmbeddingBlock(
    (linear): Linear(7x0e -> 64x0e | 448 weights)
  )
  (radial_embedding): RadialEmbeddingBlock(
    (bessel_fn): BesselBasis(r_max=5.0, num_basis=8, trainable=False)
    (cutoff_fn): PolynomialCutoff(p=5.0, r_max=5.0)
  )
  (spherical_harmonics): SphericalHarmonics()
  (atomic_energies_fn): AtomicEnergiesBlock(energies=[-6.2426, -3.1213, -4.0759, -3.3009, -2.4823, -5.7819, -2.2340])
  (interactions): ModuleList(
    (0): RealAgnosticInteractionBlock(
      (linear_up): Linear(64x0e -> 64x0e | 4096 weights)
      (conv_tp): TensorProduct(64x0e x 1x0e+1x1o+1x2e+1x3o -> 64x0e+64x1o+64x2e+64x3o | 256 paths | 256 weights)
      (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 256]
      (linear): Linear(64x0e+64x1o+64x2e+64x3o -> 64x0e+64x1o+64x2e+64x3o | 16384 weights)
      (skip_tp): FullyConnectedTensorProduct(64x0e+64x1o+64x2e+64x3o x 7x0e -> 64x0e+64x1o+64x2e+64x3o | 114688 paths | 114688 weights)
      (reshape): reshape_irreps()
    )
    (1): RealAgnosticResidualInteractionBlock(
      (linear_up): Linear(64x0e+64x1o+64x2e -> 64x0e+64x1o+64x2e | 12288 weights)
      (conv_tp): TensorProduct(64x0e+64x1o+64x2e x 1x0e+1x1o+1x2e+1x3o -> 192x0e+320x1o+320x2e+256x3o | 1088 paths | 1088 weights)
      (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 1088]
      (linear): Linear(192x0e+320x1o+320x2e+256x3o -> 64x0e+64x1o+64x2e+64x3o | 69632 weights)
      (skip_tp): FullyConnectedTensorProduct(64x0e+64x1o+64x2e x 7x0e -> 64x0e | 28672 paths | 28672 weights)
      (reshape): reshape_irreps()
    )
  )

[ilyes319]

Hey @ankitmathanker, see that issue: #633, the naming might be confusing.

[wcwitt]

To add to this briefly, the specific explanation for @ankitmathanker's observation is that --model="MACE" creates a ScaleShiftMACE with zero shift and a fixed scale (if I remember correctly), which is by design and not a problem aside from the potential confusion.