Theoretical Background

OMatG implements the stochastic interpolants (SIs) framework for the modeling and generation of inorganic crystalline materials. SIs are a unifying framework for generative modeling that encompasses flow-matching and diffusion-based methods as specific instances, while offering a more general and flexible approach enabling the design of a broad class of novel generative models.

A stochastic interpolant $x_t = x(t, x_0, x_1, z)=\alpha(t),x_0 + \beta(t),x_1 + \gamma(t),z$ bridges samples $x_0$ from a (trivial) base distribution to samples $x_1$ from the target data distribution. Here, $t\in[0, 1]$ represents time and the random variable $z$ is drawn from a Gaussian distribution. The functional forms of $\alpha$, $\beta$, and $\gamma$ are flexible, only subject to a few constraints that, amongst other things, ensure that $x(t=0, x_0, x_1, z) = x_0$ and $x(t=1, x_0, x_1, z) = x_1$.

The time-dependent probability density of the stochastic process $x_t$ can be realized either via deterministic sampling through an ordinary differential equation (ODE) or stochastic sampling through a stochastic differential equation (SDE), only requiring a sample $x_0$ from the base distribution. This enables generative modeling by evolving samples from the base distribution to samples from the data distribution. Here, the required velocity term $b^\theta(t, x)$ for both ODE- and SDE-based sampling can be learned from data by "averaging over many paired samples ($x_0, x_1$)." For SDE-based sampling, an additional denoiser $z^\theta(t, x)$ can be learned likewise.

The flexibility of the SI framework stems from the ability to tailor the choice of interpolants and choosing between deterministic (ODE) and stochastic (SDE) sampling schemes (see Fig. 1 that visualizes the tunable components of the SI framework for bridging samples from a base distribution (gray particles) to samples from a target distribution (purple particles); figure taken from the OMatG paper.).

OMatG defines a crystalline material of $N$ atoms by its unit cell that is described by three lattice vectors $\mathbf{L} \in \mathbb{R}^{3\times3}$, $N$ fractional coordinates $\mathbf{X}\in[0,1)^{3\times N}$ with periodic boundary conditions, and $N$ discrete atomic species $\mathbf{A}\in\mathbb{Z}^N_{>0}$. During training an generation, all three components ${\mathbf{A}, \mathbf{X}, \mathbf{L}}$ are considered simultaneously. The SI framework is applied to the continuous structural representations ${\mathbf{X}, \mathbf{L}}$ while the discrete atomic species $\mathbf{A}$ are treated with discrete flow matching.

Configuration Files

Machine-learning models implemented with PyTorch Lightning rely on three essential parts:

1. Trainer: The training engine.
2. LightningDataModule: Handles data loading and preprocessing.
3. LightningModule: Defines the model and training logic.

Configuration files of OMatG thus generally contain specifications for these three parts.

Trainer

OMatG uses the standard PyTorch Lightning Trainer. Its parameters are specified in the trainer section of the configuration file, for example:

trainer:
  callbacks:  # List of callbacks to be used during training.
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        filename: "best_val_loss_total"
        save_top_k: 1
        monitor: "val_loss_total"
        save_weights_only: true
  accelerator: "gpu"
  gradient_clip_val: 0.5
  gradient_clip_algorithm: "value"
  num_sanity_val_steps: 0
  precision: "32-true"
  max_epochs: 2000
  enable_progress_bar: true

Note that it is possible to initialize specialized classes in the configuration file by specifying the class_path and init_args. The init_args dictionary contains the arguments that are passed to the constructor of the class.

In addition to the trainer, one should specify the optimizer and (optionally) the learning rate scheduler in their own sections:

optimizer:
  class_path: torch.optim.AdamW
  init_args:
    lr: 0.001
    weight_decay: 0.01
lr_scheduler:
  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
  init_args:
    T_max: 2000
    eta_min: 1e-07

LightingDataModule

The data section of the configuration constructs the OMGDataModule (see omg/datamodule/dataloader.py). It mainly expects the train_dataset, val_dataset, and predict_dataset sections. Each of these sections should construct an OMGTorchDataset (see omg/datamodule/dataloader.py again). This can be done based on lmdb files:

data:
  train_dataset:
    class_path: omg.datamodule.dataloader.OMGTorchDataset
    init_args:
      dataset:
        class_path: omg.datamodule.datamodule.DataModule
        init_args:
          lmdb_paths:
           - "data/mp_20/train.lmdb"
      niggli: False
  val_dataset:
    class_path: omg.datamodule.dataloader.OMGTorchDataset
    init_args:
      dataset:
        class_path: omg.datamodule.datamodule.DataModule
        init_args:
          lmdb_paths:
           - "data/mp_20/val.lmdb"
      niggli: False
  predict_dataset:
    class_path: omg.datamodule.dataloader.OMGTorchDataset
    init_args:
      dataset:
        class_path: omg.datamodule.datamodule.DataModule
        init_args:
          lmdb_paths:
           - "data/mp_20/test.lmdb"
      niggli: False
  batch_size: 32
  num_workers: 4
  pin_memory: True
  persistent_workers: True

Every record in the lmdb files should contain a crystal structure. The key of each record is assumed to be an (arbitrary) encoded string, while the value is assumed to be a pickled dictionary with, at least, the following keys:

• pos: A torch.Tensor of shape (N, 3) containing the fractional coordinates of the atoms in the crystal structure.
• cell: A torch.Tensor of shape (3, 3) containing the lattice vectors of the crystal structure.
• atomic_numbers: A torch.Tensor of shape (N,) containing the atomic numbers of the atoms in the crystal structure.

The data section can also contain additional parameters for the data loading (such as batch_size, num_workers, pin_memory, and persistent_workers in the above example). These parameters are passed to the underlying PyTorch DataLoader instances.

Within OMatG, the data is passed around as torch_geometric.data.Data instances. For a batch size of batch_size, these instances contain the following attributes:

• n_atoms: torch.Tensor of shape (batch_size, ) containing the number of atoms in each configuration.
• batch: torch.Tensor of shape (sum(n_atoms),) containing the index of the configuration to which each atom belongs.
• species: torch.Tensor of shape (sum(n_atoms),) containing the atomic numbers of the atoms in the configurations.
• pos: torch.Tensor of shape (sum(n_atoms), 3) containing the atomic positions of the atoms in the configurations.
• cell: torch.Tensor of shape (batch_size, 3, 3) containing the cell vectors of the configurations.
• ptr: torch.Tensor of shape (batch_size + 1,) containing the indices of the first atom of each configuration in the species and pos tensors.
• property: dict containing the properties of the configurations.

LightningModule

The model section of the configuration file constructs the OMGLightningModule (see omg/omg_lightning.py). Its arguments are documented in the class docstring. An example model section looks as follows:

model:
  si:  # Collection of stochastic interpolants.
    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
    init_args:
      stochastic_interpolants:
        # Chemical species.
        # The SingleStochasticInterpolantIdentity keeps the species unchanged during interpolation (CSP task).
        # For DNG, use, e.g., omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask.
        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
        # Fractional coordinates.
        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
          init_args:
            # Use a periodic interpolant for fractional coordinates.
            interpolant: omg.si.interpolants.PeriodicLinearInterpolant
            gamma: null
            epsilon: null
            differential_equation_type: "ODE"
            integrator_kwargs:
              method: "euler"
            velocity_annealing_factor: 10.182659004291072
            correct_center_of_mass_motion: true
        # Lattice vectors.
        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
          init_args:
            # Use a non-periodic interpolant for lattice vectors.
            interpolant: omg.si.interpolants.LinearInterpolant
            gamma: null
            epsilon: null
            differential_equation_type: "ODE"
            integrator_kwargs:
              method: "euler"
            velocity_annealing_factor: 1.824475401606087
            correct_center_of_mass_motion: false
      data_fields:
        # If the order of the data_fields changes,
        # the order of the above StochasticInterpolant inputs must also change.
        - "species"
        - "pos"
        - "cell"
      integration_time_steps: 1000
  relative_si_costs:
    species_loss: 0.0
    pos_loss_b: 0.999
    cell_loss_b: 0.001
  sampler:
    class_path: omg.sampler.sample_from_rng.SampleFromRNG
    init_args:
      # Uniform distribution for fractional coordinates.
      pos_distribution: null
      cell_distribution:
        class_path: omg.sampler.distributions.InformedLatticeDistribution
        init_args:
          dataset_name: mp_20
      species_distribution:
        # For DNG, use omg.sampler.distributions.MaskDistribution.
        class_path: omg.sampler.distributions.MirrorData
  model:
    class_path: omg.model.model.Model
    init_args:
      encoder:
        class_path: omg.model.encoders.cspnet_full.CSPNetFull
      head:
        class_path: omg.model.heads.pass_through.PassThrough
      time_embedder:
        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
        init_args:
          dim: 256

The si section combines the stochastic interpolants for the species, pos, and cell data fields of the crystal structures in the StochasticInterpolants class. This class is documented in its docstring but, in a nutshell, it is a container for multiple StochasticInterpolant instances. The typically used implementations of this abstract class are:

• SingleStochasticInterpolant: For continuous data fields such as fractional coordinates and lattice vectors with arbitrary base distributions. The specific interpolant and its parameters are specified on initialization of this class. Every interpolant has a periodic (for fractional coordinates) and a non-periodic (for lattice vectors) version.
• SingleStochasticInterpolantOS: For continuous data fields such as fractional coordinates and lattice vectors, but explicitly assuming a Gaussian base distribution as it implements one-sided stochastic interpolants.
• SingleStochasticInterpolantIdentity: For keeping the corresponding data field unchanged during interpolation and generation.
• DiscreteFlowMatchingMask: For discrete data fields such as atomic species with a completely masked base distribution.
• DiscreteFlowMatchingUniform: For discrete data fields such as atomic species with a uniform base distribution.

Every StochasticInterpolant in the StochasticInterpolants class computes losses and returns them in a dictionary (see the loss_keys method in the respective class). The StochasticInterpolants class prefixes these keys with the name of the corresponding data field so that the losses can be identified. The relative_si_costs section specifies the relative weights of these losses when they are added up during training.

The sampler section specifies the base distributions for the positions, lattice vectors, and atomic species. Depending on the choice of the stochastic interpolant, one should choose the matching base distribution:

• SingleStochasticInterpolant: The choice of the base distribution is arbitrary. As in the example above, we typically use a uniform distribution for the fractional coordinates and an informed base distribution for the lattice vectors.
• SingleStochasticInterpolantOS: Explicitly assumes a NormalDistribution.
• SingleStochasticInterpolantIdentity: Explicitly assumes that the training data is just taken over in the "random" sample as implemented by the MirrorData distribution.
• DiscreteFlowMatchingMask: Explicitly assumes fully masked samples as the base distribution as implemented in the MaskDistribution.
• DiscreteFlowMatchingUniform: Explicitly assumes uniformly distributed atomic species as the base distribution which can achieved by using species_distribution: null.

The model section specifies the model architecture. In the above example, we just use DiffCSPNet.