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

TL;DR

OmniMol tackles cross-domain transfer by adapting a jet-pretrained Point-Edge Transformer backbone to molecular dynamics, repurposing the generator head for per-atom force prediction while enforcing energy conservation and partial equivariance. It leverages a physics-informed global attention bias and MD-specific local features, trained with a joint energy/force objective on the oMoL dataset, achieving state-of-the-art performance with a compact transformer and demonstrating strong gains in low-data regimes via pre-training and LoRA fine-tuning. The work shows that cross-domain, point-cloud foundation models can deliver accurate MLIPs with far fewer parameters than large-domain baselines, enabling scalable, interpretable MD simulations. It also outlines future directions, including larger backbones, complete energy-conservation in bigger models, and broader evaluations across diverse point-cloud domains.

Abstract

We present OmniMol, a state-of-the-art 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. This study lays the foundation for building interdisciplinary connections, given datasets represented as collections of point clouds.

Figures (7)

• Figure 1: Local Embedding Block, where K-Nearest Neighbors for either Molecule or Jet point-clouds are computed along with pairwise features, these are fed to an MLP, and a small local transformer resulting in $\vec{x}_{embed}^{local}$
• Figure 2: OmniLearned's Architecture
• Figure 3: OmniMol's Architecture
• Figure 4: Conservative and Equivariant OmniMol's Architecture
• Figure 5: Scaling behavior for (left) energy and (right) forces of OmniMol direct small and medium pre-trained and from scratch. Finetuning with ten and one hundred thousand molecules on OmniMol small proceeds with LoRA, 4 million and OmniMol medium with full finetuning.