TriForces: Augmenting Atomistic GNNs for Transferable Representations

TL;DR

TriForces introduces a three-stream framework that enhances the transferability of atomistic GNNs by separating composition and structure information and employing self-supervised learning, leading to improved performance on multiple benchmarks.

Contribution

The paper presents TriForces, a novel model-agnostic framework that improves transferability of MLIPs by separating composition and structure info and using self-supervised learning.

Findings

Reduces energy MAE by 57% on OMat24 with only 20K samples

Improves force MAE across various sample sizes

Enables efficient retrieval of similar structures in learned latent space

Abstract

Machine learning interatomic potentials (MLIPs) achieve excellent accuracy when trained on large Density Functional Theory (DFT) data. To be useful in practice, they must often be adapted to target chemistries using small and expensive task-specific datasets. However, MLIPs transfer inconsistently across domains, with representations that often loose accessible composition and structure information. To address this, we present TriForces, a model-agnostic three-stream framework that separates composition and structure information, combined with self-supervised learning to preserve transferable representations. TriForces improves performance on MatBench and QM9 over baselines without needing DFT labels and enables efficient similar structure retrieval through its learned latent space. On OMat24, in limited-data training regime, TriForces reduces energy MAE by 57% at 20K samples only and improves force MAE across sample sizes. We release pretrained TriForces variants across multiple MLIP architectures with code at https://github.com/Ramlaoui/triforces.