FlashMD

FlashMD is a method for the direct prediction of positions and momenta in a molecular dynamics simulation, presented in [1]. When compared to traditional molecular dynamics methods, it predicts the positions and momenta of atoms after a long time interval, allowing the use of much larger time steps. Therefore, it achieves a significant speedup (10-30x) compared to molecular dynamics using MLIPs. The FlashMD architecture implemented in metatrain is based on The PET architecture.

Installation

To install FlashMD and its dependencies, you can run the following from the root of the repository:

pip install metatrain[flashmd]

Default Hyperparameters

The default hyperparameters for the FlashMD architecture are:

architecture:
  name: experimental.flashmd
  model:
    predict_momenta_as_difference: false
    cutoff: 4.5
    cutoff_width: 0.2
    d_pet: 128
    d_head: 128
    d_node: 512
    d_feedforward: 256
    num_heads: 8
    num_attention_layers: 2
    num_gnn_layers: 3
    normalization: RMSNorm
    activation: SwiGLU
    transformer_type: PreLN
    featurizer_type: feedforward
    long_range:
      enable: false
      use_ewald: false
      smearing: 1.4
      kspace_resolution: 1.33
      interpolation_nodes: 5
  training:
    timestep: null
    masses: {}
    distributed: false
    distributed_port: 39591
    batch_size: 16
    num_epochs: 1000
    warmup_fraction: 0.01
    learning_rate: 3e-4
    weight_decay: null
    log_interval: 1
    checkpoint_interval: 100
    scale_targets: true
    fixed_composition_weights: {}
    fixed_scaling_weights: {}
    per_structure_targets: []
    num_workers: null
    log_mae: false
    log_separate_blocks: false
    best_model_metric: rmse_prod
    grad_clip_norm: 1.0
    loss: mse

Hyperparameters

In order to choose and tune hyperparameters for FlashMD, you can refer to the PET architecture documentation, as most of the hyperparameters are shared between the two architectures. Here, we will only discuss the hyperparameters that are specific to FlashMD.

model

  • predict_momenta_as_difference: This parameter controls whether the model will predict future momenta directly or as a difference between the future and the current momenta. Setting it to true will help when predicting relatively small timesteps (when compared to the momentum autocorrelation time), while setting it to false is beneficial when predicting large timesteps. Defaults to false.

training

  • timestep: The time interval (in fs) between the current and the future positions and momenta that the model must predict. This option is not used in the training, but it is registered in the model and it will be used to validate that the timestep used during inference in MD engines is the same as the one used during training. This hyperparameter must be provided by the user.

  • masses: A dictionary mapping atomic species to their masses (in atomic mass units). Indeed, it should be noted that FlashMD models, as implemented in metatrain, are not transferable across different isotopes. The masses are not used during training, but they are registered in the model and they will be used during inference to validate that the masses used in MD engines are the same as the ones used during training. If not provided, masses from the ase.data module will be used. These correspond to masses averaged over the natural isotopic abundance of each element.

References