LTAU

LTAU (Loss Trajectory Analysis for Uncertainty) is a technique for estimating the uncertainty in a model's predictions for any regression task by approximating the probability distribution functions (PDFs) of the model's per-sample errors over the course of training. The approximated PDFs, combined with a similarity search algorithm in the model's descriptor space, can then be used to estimate the model's uncertainty on any given test point.

LTAU-FF

LTAU-FF (LTAU for atomistic Force Fields) is the application of LTAU spefically for use with atomistic force fields. LTAU-FF has a number of advantages including:

• Simplicity: just log per-sample errors at every epoch!
• Speed: leveraging the FAISS library for fast similarity search, the computational cost of LTAU-FF can be made negligible compared to a forward pass of the model.
• Utility: in our published work, we show that LTAU-FF can be used in practical applications for:
  • Generating a well-calibrated UQ metric
  • Detecting out-of-domain data
  • Predicting failure in the OC20 IS2RS task

When publishing results using this package, please cite:

@misc{2402.00853,
    Author = {Joshua A. Vita and Amit Samanta and Fei Zhou and Vincenzo Lordi},
    Title = {LTAU-FF: Loss Trajectory Analysis for Uncertainty in Atomistic Force Fields},
    Year = {2024},
    Eprint = {arXiv:2402.00853},
}

Getting started

Installation

The only dependency of LTAU-FF is the FAISS package. We recommend first trying:

git clone https://github.com/LLNL/ltau-ff/ltau-ff.git
cd ltau-ff
pip install -e .

which will attempt to install faiss-cpu==1.8.0. If this fails (which it will on POWER9 architectures), then instead try:

conda install faiss-cpu=1.8.0

If you encounter further installation difficulties with FAISS, please refer to the instructions provided by the FAISS documentation.

Additional software patches are provided for external packages. These do not need to be installed unless you want to use the corresponding functionality.

• ase_trajectory_patch.py: for allowing ASE Trajectory objects to write non-standard properties (i.e., 'uq', and 'descriptors'). Replaces ase.io.trajectory. Also requires commenting out this line in ase.calculators.singlepoint.
• nequip
  • (only needed for nequip < v0.7.0) nequip_trainer_patch.py: for logging per-sample energy/force errors at every epoch while training a NequIP model. Replaces nequip.train.trainer.
  • nequip_modules.py: custom trainer and dataset modules which allow for logging per-atom errors at each epoch. Currently only tested on 1 GPU. You can use these when training your nequip-style model by using these blocks in your config file:

  training_module:
      _target_: ltau_nequip_modules.LTAULightningModule
      error_save_dir: ${hydra:runtime.output_dir}/${trainer.logger.name}
  data:
      _target_: ltau_nequip_modules.ASEDataModuleWithID 
  

Depending on demand, these patches may eventually be opened as pull requests on their respective repositories.

Examples

• Basic tutorial

Supported models

The UQEstimator class is model-agnostic, and will work immediately as long as you can provide the training loss trajectories and the per-sample descriptors.

Some additionally functionality is provided specifically for the NequIP and MACE models; in particular, these include an ASE wrapper for running simulations with UQ, and helper scripts for extracting descriptors and performing energy minimization using a trained model. If you would like to add similar suppport for another model, we recommend taking a look at the following files for reference:

• ASE wrappers:
  • NequIPUQWrapper
  • MACEUQWrapper
• Descriptor extraction:
  • ltau-ff-nequip-descriptors
  • ltau-ff-nequip-v0.7-descriptors (for use with nequip >= 0.7.0)
  • ltau-ff-mace-descriptors
• Energy minimization
  • ltau-ff-nequip-minimizer
  • (not implemented for MACE yet)

FAQ

• Can LTAU be used with pre-trained models?
  • Yes, but you still need to have a method for sampling an ensemble of model parameters. One possibility would be to run Monte Carlo sampling, perturbing the model parameters and logging the per-atom errors. This is the approach taken in "Bayesian Ensemble Approach to Error Estimation of Interatomic Potentials" by Frederiksen et. al.

Contact

If you have any questions or comments, please either open an issue or email Josh Vita ([email protected])