SPACE (Experimental)¶

https://codecov.io/gh/metatensor/metatrain/branch/main/graph/badge.svg?component=space

SPACE (Smooth Physical Architecture with Compact Equivariants) is a physics-inspired equivariant neural network architecture. Compared to, for example, MACE and GRACE, it uses a geometrically motivated basis and a fast and elegant tensor product implementation. The tensor product used in SPACE leverages a equivariant representation that differs from the typical spherical one. You can read more about it here: https://pubs.acs.org/doi/10.1021/acs.jpclett.4c02376.

Installation¶

To install this architecture along with the metatrain package, run:

pip install metatrain[space]

where the square brackets indicate that you want to install the optional dependencies required for space.

Additional outputs¶

feature: the internal SPACE features, before the different heads for each target.
mtt::aux::{target}_last_layer_features: The features for a given target, taken before the last linear layer of the corresponding head.

Default Hyperparameters¶

The description of all the hyperparameters used in space is provided further down this page. However, here we provide you with a yaml file containing all the default hyperparameters, which might be convenient as a starting point to create your own hyperparameter files:

architecture:
  name: experimental.space
  model:
    num_tensor_products: 6
    num_gnn_layers: 3
    cutoff: 8.0
    num_neighbors_adaptive: 16
    cutoff_width: 1.0
    num_element_channels: 128
    force_rectangular: false
    radial_basis:
      max_eigenvalue: 25.0
      element_scale: 0.7
      mlp_depth: 3
      mlp_expansion_ratio: 4
    initial_scaling: 1.0
    message_scaling: 0.1
    final_scaling: 1.0
    use_sphericart: false
    mlp_head_num_layers: 1
    mlp_head_expansion_ratio: 4
    tensor_product_expansion_ratio: 2
    heads: {}
    zbl: false
  training:
    compile: false
    distributed: null
    distributed_port: 39591
    batch_size: 8
    max_atoms_per_batch: null
    min_atoms_per_batch: 0
    num_epochs: 1000
    learning_rate: 0.003
    warmup_fraction: 0.01
    gradient_clipping: 1.0
    ema_decay: 0.999
    log_interval: 1
    checkpoint_interval: 25
    scale_targets: true
    atomic_baseline: {}
    fixed_scaling_weights: {}
    num_workers: null
    per_structure_targets: []
    log_separate_blocks: false
    log_mae: false
    best_model_metric: rmse_prod
    loss: mse
    finetune:
      read_from: null
      method: full
      config: {}
      inherit_heads: {}

Tuning hyperparameters¶

The default hyperparameters above will work well in most cases, but they may not be optimal for your specific use case. There is good number of parameters to tune, both for the model and the trainer. Here, we provide a list of the parameters that are in general the most important (in decreasing order of importance) for the SPACE architecture:

ModelHypers.radial_basis: RadialBasisHypers = {'element_scale': 0.7, 'max_eigenvalue': 25.0, 'mlp_depth': 3, 'mlp_expansion_ratio': 4}

Hyperparameters for the radial basis functions.

Raising``max_eigenvalue`` from its default will increase the number of spherical irreducible representations (irreps) used in the model, which can improve accuracy at the cost of computational efficiency. Increasing this value will also increase the number of radial basis functions (and therefore internal features) used for each irrep.

ModelHypers.num_element_channels: int = 128

Number of channels per element.

This determines the size of the embedding used to encode the atomic species, and it increases or decreases the size of the internal features used in the model.

TrainerHypers.num_epochs: int = 1000

Number of epochs to train the model.

A larger number of epochs might lead to better accuracy. In general, if you see that the validation metrics are not much worse than the training ones at the end of training, it might be a good idea to increase this value.

TrainerHypers.batch_size: int = 8

Batch size for training.

Decrease this value if you run into out-of-memory errors during training. You can try to increase it if your structures are very small (less than 20 atoms) and you have a good GPU.

ModelHypers.num_gnn_layers: int = 3

Number of GNN layers.

Increasing this value might increase the accuracy of the model (especially on larger datasets), at the expense of computational efficiency.

TrainerHypers.learning_rate: float = 0.003

Learning rate for the optimizer.

You can try to increase this value (e.g., to 0.01) if training is stable and slow or decrease it (e.g., to 0.001 or less) if you see divergence in the first few epochs and/or instabilities.

ModelHypers.cutoff: float = 8.0

Cutoff radius for neighbor search.

This should be set to a value after which most of the interactions between atoms is expected to be negligible. A lower cutoff will lead to faster models.

ModelHypers.force_rectangular: bool = False

Makes the number of channels per irrep the same.

This might improve accuracy with a limited increase in computational cost.

Model hyperparameters¶

The parameters that go under the architecture.model section of the config file are the following:

ModelHypers.num_tensor_products: int = 6¶

Number of tensor products per GNN layer.

ModelHypers.num_gnn_layers: int = 3¶

Number of GNN layers.

Increasing this value might increase the accuracy of the model (especially on larger datasets), at the expense of computational efficiency.

ModelHypers.cutoff: float = 8.0¶

Cutoff radius for neighbor search.

This should be set to a value after which most of the interactions between atoms is expected to be negligible. A lower cutoff will lead to faster models.

ModelHypers.num_neighbors_adaptive: int | None = 16¶

Target number of neighbors for the adaptive cutoff scheme.

This parameter activates the adaptive cutoff functionality. Each atomic environment has a different cutoff, that is chosen such that the number of neighbors is approximately equal to this value. This can be useful to have a more uniform number of neighbors per atom, especially in sparse systems. Setting it to None disables this feature and uses all neighbors within the fixed cutoff radius.

ModelHypers.cutoff_width: float = 1.0¶

Width of the cutoff smoothing function.

ModelHypers.num_element_channels: int = 128¶

Number of channels per element.

This determines the size of the embedding used to encode the atomic species, and it increases or decreases the size of the internal features used in the model.

ModelHypers.force_rectangular: bool = False¶

Makes the number of channels per irrep the same.

This might improve accuracy with a limited increase in computational cost.

ModelHypers.radial_basis: RadialBasisHypers = {'element_scale': 0.7, 'max_eigenvalue': 25.0, 'mlp_depth': 3, 'mlp_expansion_ratio': 4}¶

Hyperparameters for the radial basis functions.

Raising``max_eigenvalue`` from its default will increase the number of spherical irreducible representations (irreps) used in the model, which can improve accuracy at the cost of computational efficiency. Increasing this value will also increase the number of radial basis functions (and therefore internal features) used for each irrep.

ModelHypers.initial_scaling: float = 1.0¶

Scaling for the initial features.

ModelHypers.message_scaling: float = 0.1¶

Scaling for message passing.

ModelHypers.final_scaling: float = 1.0¶

Final scaling factor applied to the model outputs.

ModelHypers.use_sphericart: bool = False¶

Whether to use spherical Cartesian coordinates.

ModelHypers.mlp_head_num_layers: int = 1¶

Number of layers in the heads for MLP heads.

ModelHypers.mlp_head_expansion_ratio: int = 4¶

Expansion ratio for the hidden layers of the MLP head.

ModelHypers.tensor_product_expansion_ratio: int = 2¶

Expansion ratio for the tensor product iterations.

ModelHypers.heads: dict[str, Literal['linear', 'mlp']] = {}¶

Heads to use in the model, with options being “linear” or “mlp”.

ModelHypers.zbl: bool = False¶

Whether to use the ZBL potential in the model.

with the radial basis hyperparameters being:

class metatrain.experimental.space.documentation.RadialBasisHypers[source]¶

Hyperparameter concerning the radial basis functions used in the model.

max_eigenvalue: float = 25.0¶

Maximum eigenvalue for the radial basis.

This single parameter controls the maximum angular momentum l_max and the number of radial basis functions per angular momentum n_max_l.

A sufficiently high value of max_eigenvalue must be chosen to cover the maximum rank of the target property. The following guide (only up to L = 10) can be used to choose a suitable value for max_eigenvalue:

`l_max`	`max_eigenvalue` range, `[low, high)`
0	`[8.00, 12.97)`
1	`[12.97, 22.36)`
2	`[22.36, 35.32)`
3	`[35.32, 51.93)`
4	`[51.93, 72.21)`
5	`[72.21, 96.20)`
6	`[96.20, 123.87)`
7	`[123.87, 155.24)`
8	`[155.24, 190.32)`
9	`[190.32, 229.09)`
10	`[229.09, 271.55)`

element_scale: float = 0.7¶: Scaling factor for the element-dependent radial lengthscales.

mlp_depth: int = 3¶: Depth of the radial MLP. Must be at least 2.

mlp_expansion_ratio: int = 4¶: Expansion ratio for the radial MLP hidden layers.

Trainer hyperparameters¶

The parameters that go under the architecture.trainer section of the config file are the following:

TrainerHypers.compile: bool = False¶

Whether to use torch.compile during training.

This can lead to significant speedups, but it will cause a compilation step at the beginning of training which might take up to 5-10 minutes, mainly depending on max_eigenvalue. Note that this option does not work at the moment with adaptive cutoffs.

TrainerHypers.distributed: NotRequired[bool | None] = None¶

Whether to use distributed training. When not set, distributed training is enabled automatically when running under more than one SLURM task. Setting this option explicitly is deprecated.

TrainerHypers.distributed_port: int = 39591¶

Port for DDP communication.

TrainerHypers.batch_size: int = 8¶

Batch size for training.

Decrease this value if you run into out-of-memory errors during training. You can try to increase it if your structures are very small (less than 20 atoms) and you have a good GPU.

TrainerHypers.max_atoms_per_batch: int | None = None¶

If set, use greedy atom-count packing instead of fixed batch_size. Structures are accumulated into each batch until adding another would exceed this limit, producing variable numbers of structures per batch. Supported with any dataset type. When set, batch_size is ignored for constructing training and validation batches (it is still used internally for composition model and scaler fitting).

TrainerHypers.min_atoms_per_batch: int = 0¶

Minimum total number of atoms required to keep a batch when max_atoms_per_batch is set. Batches whose total atom count falls below this threshold are discarded during packing. Defaults to 0 (no minimum).

TrainerHypers.num_epochs: int = 1000¶

Number of epochs to train the model.

A larger number of epochs might lead to better accuracy. In general, if you see that the validation metrics are not much worse than the training ones at the end of training, it might be a good idea to increase this value.

TrainerHypers.learning_rate: float = 0.003¶

Learning rate for the optimizer.

You can try to increase this value (e.g., to 0.01) if training is stable and slow or decrease it (e.g., to 0.001 or less) if you see divergence in the first few epochs and/or instabilities.

TrainerHypers.warmup_fraction: float = 0.01¶

Fraction of training steps for learning rate warmup.

TrainerHypers.gradient_clipping: float | None = 1.0¶

Gradient clipping value. If None, no clipping is applied.

TrainerHypers.ema_decay: float | None = 0.999¶

Decay factor for exponential moving average of model parameters. If None, EMA is not used.

TrainerHypers.log_interval: int = 1¶

Interval to log metrics during training.

TrainerHypers.checkpoint_interval: int = 25¶

Interval to save model checkpoints.

TrainerHypers.scale_targets: bool = True¶

Normalize targets to unit std during training.

If true, a single scale is computed for each target, given by the uncentered standard deviation across all values in the dataset for that target.

For targets with more than one property (i.e. > 1 block or >= 1 block with > 1 property), per-property scales are also computed, and used to re-scale model predictions.

See also Target scaling.

TrainerHypers.atomic_baseline: Dict[str, float | Dict[int, float]] | str = {}¶

The baselines for each target.

By default, metatrain will fit a linear model (CompositionModel) to compute the least squares baseline for each atomic species for each target.

However, this hyperparameter allows you to provide your own baselines, either as a dictionary or as a path to a pre-trained composition model checkpoint. The value of the hyperparameter should either be:

a dictionary where the keys are the target names, and the values are either (1) a single baseline to be used for all atomic types, or (2) a dictionary mapping atomic types to their baselines.

a string path to a .ckpt file from a pre-trained composition model.

For example:

atomic_baseline: {"energy": {1: -0.5, 6: -10.0}} will fix the energy baseline for hydrogen (Z=1) to -0.5 and for carbon (Z=6) to -10.0, while fitting the baselines for the energy of all other atomic types, as well as fitting the baselines for all other targets.

atomic_baseline: {"energy": -5.0} will fix the energy baseline for all atomic types to -5.0.

atomic_baseline: {"mtt:dos": 0.0} sets the baseline for the “mtt:dos” target to 0.0, effectively disabling the atomic baseline for that target.

atomic_baseline: "/path/to/model.ckpt" loads a pre-trained composition model checkpoint, overriding the default least-squares fit.

This atomic baseline is substracted from the targets during training, which avoids the main model needing to learn atomic contributions, and likely makes training easier. When the model is used in evaluation mode, the atomic baseline is added on top of the model predictions automatically.

Note

This atomic baseline is a per-atom contribution. Therefore, if the property you are predicting is a sum over all atoms (e.g., total energy), the contribution of the atomic baseline to the total property will be the atomic baseline multiplied by the number of atoms of that type in the structure.

Note

If a MACE model is loaded through the mace_model hyperparameter, the atomic baselines in the MACE model are used by default for the target indicated in mace_head_target. If you want to override them, you need to set explicitly the baselines for that target in this hyperparameter.

TrainerHypers.fixed_scaling_weights: Dict[str, float | Dict[int, float]] | str = {}¶

Weights for target scaling.

This is passed to the fixed_weights argument of Scaler.train_model, see its documentation to understand exactly what to pass here.

Apart from those options, one can pass a path to a model checkpoint. If that is the checkpoint of a Scaler model, the pre-trained scaler will be loaded. When passing a checkpoint for the scaler, atomic_baseline must also be a checkpoint for a composition model.

TrainerHypers.num_workers: int | None = None¶

Number of workers for data loading.

TrainerHypers.per_structure_targets: list[str] = []¶

List of targets to calculate per-structure losses.

TrainerHypers.log_separate_blocks: bool = False¶

Whether to log per-block error during training.

TrainerHypers.log_mae: bool = False¶

Whether to log MAE alongside RMSE during training.

TrainerHypers.best_model_metric: Literal['rmse_prod', 'mae_prod', 'loss'] = 'rmse_prod'¶

Metric used to select the best model checkpoint.

TrainerHypers.loss: str | dict[str, LossSpecification] = 'mse'¶

Loss function used for training.
TrainerHypers.finetune: NoFinetuneHypers | FullFinetuneHypers | LoRaFinetuneHypers | HeadsFinetuneHypers = {'config': {}, 'inherit_heads': {}, 'method': 'full', 'read_from': None}¶
Finetuning hyperparameters.

Setting read_from to a path enables finetuning from a pretrained SPACE checkpoint. The method can be one of:

"full": all parameters are trainable.

"heads": only the heads and last layers are trained; the rest of the model is frozen.

"lora": LoRA adapters are injected into the linear layers and only those are trained.

config may be left empty, in which case the SPACE module names are used (linear_layer for LoRA; the heads and last layers of the base model for "heads"). Example for heads-only finetuning:
finetune:
  read_from: /path/to/pretrained.ckpt
  method: heads
Example for LoRA finetuning:
finetune:
  read_from: /path/to/pretrained.ckpt
  method: lora
  config:
    rank: 8
    alpha: 16