SOAP-BPNN

This is a Behler-Parrinello type neural network [1], which, instead of their original atom-centered symmetry functions, we use the Smooth overlap of atomic positions (SOAP) [2] as the atomic descriptors, computed with torch-spex.

Installation

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

pip install metatrain[soap_bpnn]

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

Default Hyperparameters

The description of all the hyperparameters used in soap_bpnn 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: soap_bpnn
  model:
    soap:
      max_angular: 6
      max_radial: 7
      cutoff:
        radius: 5.0
        width: 0.5
    bpnn:
      num_hidden_layers: 2
      num_neurons_per_layer: 32
      layernorm: true
    add_lambda_basis: true
    heads: {}
    zbl: false
    long_range:
      enable: false
      use_ewald: false
      smearing: 1.4
      kspace_resolution: 1.33
      interpolation_nodes: 5
  training:
    distributed: false
    distributed_port: 39591
    batch_size: 8
    num_epochs: 100
    warmup_fraction: 0.01
    learning_rate: 0.001
    log_interval: 5
    checkpoint_interval: 25
    scale_targets: true
    fixed_composition_weights: {}
    remove_composition_contribution: true
    fixed_scaling_weights: {}
    per_structure_targets: []
    num_workers: null
    log_mae: false
    log_separate_blocks: false
    best_model_metric: rmse_prod
    loss: mse

Model hyperparameters

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

ModelHypers.soap: SOAPConfig = {'cutoff': {'radius': 5.0, 'width': 0.5}, 'max_angular': 6, 'max_radial': 7}

Configuration of the SOAP descriptors.

ModelHypers.bpnn: BPNNConfig = {'layernorm': True, 'num_hidden_layers': 2, 'num_neurons_per_layer': 32}

Configuration of the neural network architecture.

ModelHypers.add_lambda_basis: bool = True

This boolean parameter controls whether or not to add a spherical expansion term of the same angular order as the targets, when they are tensorial.

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

The type of head (“linear” or “mlp”) to use for each target (e.g. heads: {“energy”: “linear”, “mtt::dipole”: “mlp”}). All omitted targets will use a MLP (multi-layer perceptron) head. MLP heads consists of one hidden layer with as many neurons as the SOAP-BPNN (see BPNNConfig.num_neurons_per_layer).

ModelHypers.zbl: bool = False

Whether to use the ZBL short-range repulsion as the baseline for the model. May be needed to achieve better description at the close-contact, repulsive regime.

ModelHypers.long_range: LongRangeHypers = {'enable': False, 'interpolation_nodes': 5, 'kspace_resolution': 1.33, 'smearing': 1.4, 'use_ewald': False}

Parameters related to long-range interactions.

May be needed to describe important long-range effects not captured by the short-range SOAP-BPNN model

with the following definitions needed to fully understand some of the parameters:

class metatrain.soap_bpnn.documentation.SOAPConfig

Configuration for the SOAP descriptors.

max_angular: int = 6

Maximum angular channels of the spherical harmonics when computing the SOAP descriptors.

max_radial: int = 7

Maximum radial channels of the spherical harmonics when computing the SOAP descriptors.

cutoff: SOAPCutoffConfig = {'radius': 5.0, 'width': 0.5}

Determines the cutoff routine of the atomic environment.

class metatrain.soap_bpnn.documentation.SOAPCutoffConfig

Cutoff configuration for the SOAP descriptor.

radius: float = 5.0

Should be set to a value after which most of interatomic is expected to be negligible. Note that the values should be defined in the position units of your dataset.

width: float = 0.5

The radial cutoff of atomic environments is performed smoothly, over another distance defined by this parameter.

class metatrain.soap_bpnn.documentation.BPNNConfig

Configuration for the BPNN architecture.

num_hidden_layers: int = 2

Controls the depth of the neural network. Increasing this generally leads to better accuracy from the increased descriptivity, but comes at the cost of increased training and evaluation time.

num_neurons_per_layer: int = 32

Controls the width of the neural network. Increasing this generally leads to better accuracy from the increased descriptivity, but comes at the cost of increased training and evaluation time.

layernorm: bool = True

Whether to use layer normalization before the neural network. Setting this hyperparameter to false will lead to slower convergence of training, but might lead to better generalization outside of the training set distribution.

Trainer hyperparameters

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

TrainerHypers.distributed: bool = False

Whether to use distributed training

TrainerHypers.distributed_port: int = 39591

Port for distributed communication among processes

TrainerHypers.batch_size: int = 8

The number of samples to use in each batch of training. This hyperparameter controls the tradeoff between training speed and memory usage. In general, larger batch sizes will lead to faster training, but might require more memory.

TrainerHypers.num_epochs: int = 100

Number of epochs.

TrainerHypers.warmup_fraction: float = 0.01

Fraction of training steps used for learning rate warmup.

TrainerHypers.learning_rate: float = 0.001

Learning rate.

TrainerHypers.log_interval: int = 5

Interval to log metrics.

TrainerHypers.checkpoint_interval: int = 25

Interval to save checkpoints.

TrainerHypers.scale_targets: bool = True

Normalize targets to unit std during training.

TrainerHypers.fixed_composition_weights: dict[str, dict[int, float]] = {}

Weights for atomic contributions.

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

TrainerHypers.remove_composition_contribution: bool = True

Whether to remove the atomic composition contribution from the targets by fitting a linear model to the training data before training the neural network.

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

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.

TrainerHypers.per_structure_targets: list[str] = []

Targets to calculate per-structure losses.

TrainerHypers.num_workers: int | None = None

Number of workers for data loading. If not provided, it is set automatically.

TrainerHypers.log_mae: bool = False

Log MAE alongside RMSE

TrainerHypers.log_separate_blocks: bool = False

Log per-block error.

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

Metric used to select best checkpoint (e.g., rmse_prod)

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

This section describes the loss function to be used. See the Loss functions for more details.

References

[1]

Jörg Behler and Michele Parrinello. Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces. Phys. Rev. Lett., 98(14):146401, April 2007. doi:10.1103/PhysRevLett.98.146401.

[2]

Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Phys. Rev. B, 87(18):184115, May 2013. doi:10.1103/PhysRevB.87.184115.