Notes on Orb

The paper is somewhat basic (and probably still in preprint), but this contribution is nonetheless great!

Repo and Paper

Combining the ideas and implementations we get from the MatterSim paper with the implementation and open source nature of Orb and we have a really solid foundation for building an end-to-end Tc prediction model.

Currently working on training a prediction head on the 3DSC dataset using Orb as the base model. More on that coming soon.

Here is a quick summary of the paper and my thoughts on Orb.

Orb is a neural network designed to predict how atoms interact - the forces between them, the energy of the system, and how the system responds to stress. It performs these tasks well (#2 on the Matbench leaderboard, beating MatterSim and GNoME), and it does so simply.

Model Architecture

Orb uses a three-stage pipeline: an Encoder that processes the raw atomic information, a Processor that performs multiple rounds of message passing between atoms to understand their interactions, and a Decoder that produces the final predictions.

Orb is unique in how it handles physical constraints. Rather than building physical laws directly into the architecture (as many previous models have done), Orb learns these properties from data. For example, instead of using an energy-based model where forces are computed as gradients, it predicts forces directly and uses clever correction terms to ensure physical consistency. This turns out to be both faster and more accurate.

Training

The training process is interesting. Orb uses a two-phase approach:

1. First, it's pretrained using a diffusion model on a broad dataset of ground state materials. This teaches it basic chemistry - what atomic configurations are reasonable and how atoms typically arrange themselves.

2. Then, it's fine-tuned on specific trajectories from density functional theory (DFT) calculations to learn precise energetics and forces.

This approach allows it to combine data from different sources and different levels of theory during pretraining, something that's typically very difficult to do.

Practical Implementation

The codebase is clean, well-documented PyTorch, released under the Apache 2.0 license. It's designed with transfer learning in mind - you can easily take the pretrained model and adapt it for new properties while leveraging all the chemical understanding it has already learned. Digging deeper into MatterSim, I did not find this to be the case, so it makes Orb a great candidate for building on.

The model is fast - 3-6x faster than existing methods on commodity hardware. This speed, combined with its accuracy and stability, makes it practical for applications like molecular dynamics simulations.

Finetuning on specific levels of theory or using the pretrained D3 model allows us to get to specific levels of theory without computationally expensive additions of corrections.