OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

TL;DR

OmniMol is a transformer-based model for molecular dynamics that adapts techniques from high-energy physics, enabling accurate predictions with fast inference and interdisciplinary potential.

Contribution

This work introduces OmniMol, a novel molecular dynamics potential leveraging transfer from particle physics models with physics-informed attention mechanisms.

Findings

OmniMol achieves high accuracy on the oMol dataset.

The model demonstrates fast inference due to architectural transfer.

OmniMol performs well with limited fine-tuning data.

Abstract

We present OmniMol, a state-of-the-art all-to-all transformer-based small molecule machine-learned interatomic potential (MLIP). OmniMol is built by adapting Omnilearned, a foundation model for particle jets found in high-energy physics (HEP) experiments such as at the Large Hadron Collider (LHC). Omnilearned is built with a Point-Edge-Transformer (PET) and pre-trained using a diverse set of one billion particle jets. It includes an interaction-matrix attention bias that injects pairwise sub-nuclear (HEP) or atomic (molecular-dynamics) physics directly into the transformer's attention logits, steering the network toward physically meaningful neighborhoods without sacrificing expressivity. We demonstrate OmniMol using the oMol dataset and find excellent performance even with relatively few examples for fine-tuning. Further, due to architectural transfer from Omnilearned, we demonstrate uniquely fast inference. This study lays the foundation for building interdisciplinary connections given datasets represented as collections of point clouds.