LOREM (Experimental)

Maintained by @EricBoittier.

LOREM (Learning Long-Range Representations with Equivariant Messages) [1] predicts energies (forces and stress by autograd) and per-atom Born effective charges. Short-range features are a spherical neighbor density with species-dependent radial filters. Long-range features are Coulomb potentials of learned charges, evaluated with torch-pme.

Born effective charges need max_degree >= 2 and a target whose quantity is bec.

Installation

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

pip install metatrain[lorem]

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

For conda installation instructions, see the Installation page. Note that conda may not provide all optional architecture-specific dependencies.

Default Hyperparameters

The description of all the hyperparameters used in lorem 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.lorem
  model:
    cutoff: 5.0
    max_degree: 6
    max_degree_lr: 2
    num_features: 128
    num_spherical_features: 8
    num_radial: 32
    num_species: 8
    num_message_passing: 0
    equivariant_message_passing: true
    initialize_node_features: true
  training:
    distributed_port: 39591
    batch_size: 8
    max_atoms_per_batch: null
    min_atoms_per_batch: 0
    num_epochs: 100
    warmup_fraction: 0.01
    learning_rate: 0.001
    weight_decay: null
    log_interval: 1
    checkpoint_interval: 100
    scale_targets: true
    atomic_baseline: {}
    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

Tuning hyperparameters

The most impactful hyperparameters (roughly in decreasing order of importance):

ModelHypers.cutoff: float = 5.0

Short-range cutoff radius, in the length unit of the dataset.

The cosine cutoff goes to zero at this radius. Ewald smearing is cutoff / 4 and the reciprocal-space wavelength is cutoff / 8.

ModelHypers.num_features: int = 128

Width of the scalar node features.

ModelHypers.max_degree: int = 6

Maximum angular momentum of the short-range spherical features. Born effective charges need at least 2.

TrainerHypers.learning_rate: float = 0.001

Learning rate.

Default values follow the LOREM paper (5 Å cutoff, 128 features). Reduce max_degree, num_radial and num_features for faster iteration.

Model hyperparameters

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

ModelHypers.cutoff: float = 5.0

Short-range cutoff radius, in the length unit of the dataset.

The cosine cutoff goes to zero at this radius. Ewald smearing is cutoff / 4 and the reciprocal-space wavelength is cutoff / 8.

ModelHypers.max_degree: int = 6

Maximum angular momentum of the short-range spherical features. Born effective charges need at least 2.

ModelHypers.max_degree_lr: int = 2

Maximum angular momentum of the long-range charges. Must not exceed max_degree. The paper default is 2.

ModelHypers.num_features: int = 128

Width of the scalar node features.

ModelHypers.num_spherical_features: int = 8

Width of the spherical node features.

ModelHypers.num_radial: int = 32

Number of Bernstein radial basis functions.

ModelHypers.num_species: int = 8

Width of the learned species embedding.

ModelHypers.num_message_passing: int = 0

Extra message-passing steps after the initial density. The paper default is 0.

ModelHypers.equivariant_message_passing: bool = True

When message passing is on, also update the spherical features. Ignored when num_message_passing is 0.

ModelHypers.initialize_node_features: bool = True

Start the scalar node features from a projection of the species embedding. When false they start from zero, as in lorem-jax’s LoremBEC.

Trainer hyperparameters

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

TrainerHypers.distributed_port: int = 39591

Port for DDP communication.

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.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 = 100

Number of epochs.

TrainerHypers.warmup_fraction: float = 0.01

Fraction of training steps used for learning rate warmup, before the cosine decay.

TrainerHypers.learning_rate: float = 0.001

Learning rate.

TrainerHypers.weight_decay: float | None = None

Weight decay. If set, AdamW is used instead of Adam.

TrainerHypers.log_interval: int = 1

Interval to log metrics.

TrainerHypers.checkpoint_interval: int = 100

Interval to save 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 fits a linear model (CompositionModel) to compute the least-squares baseline for each atomic species for each target.

This hyperparameter instead lets you provide those baselines, either as a dictionary or as a path to a pre-trained composition model checkpoint:

• a dictionary whose keys are target names, and whose values are either one baseline for every atomic type, or a dictionary mapping atomic types to 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}} fixes the energy baseline for hydrogen (Z=1) to -0.5 and for carbon (Z=6) to -10.0, and fits the rest.

• atomic_baseline: {"energy": -5.0} fixes the energy baseline for every atomic type to -5.0.

• atomic_baseline: {"mtt:dos": 0.0} sets that target’s baseline to 0.0, which disables the atomic baseline for it.

• atomic_baseline: "/path/to/model.ckpt" loads a pre-trained composition model checkpoint instead of fitting one.

The baseline is subtracted from the targets during training and added back at evaluation. It is a per-atom contribution, so for a per-structure target such as the total energy it is multiplied by the number of atoms of that type.

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.

A path to a model checkpoint is also accepted. If that checkpoint is a Scaler model, the pre-trained scaler is loaded. When passing a checkpoint for the scaler, atomic_baseline must also be a checkpoint for a composition model.

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.grad_clip_norm: float = 1.0

Maximum gradient norm value.

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]

Egor Rumiantsev, Marcel F. Langer, Tulga-Erdene Sodjargal, Michele Ceriotti, and Philip Loche. Learning long-range representations with equivariant messages. 2025. arXiv:2507.19382, doi:10.48550/arXiv.2507.19382.