TriForces: Augmenting Atomistic GNNs for Transferable Representations

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.

Discovering new materials often requires predicting how atoms will arrange and move, but the most accurate physics calculations are too slow to use at large scale. Machine learning models can learn from these calculations and make predictions much faster, but they often struggle when used on new kinds of materials or when only a small amount of new data is available. We introduce TriForces, a way to improve these models by helping them keep track of three complementary kinds of information: what atoms are present, how they are arranged in space, and how these two aspects interact. We also train the model with simple self-supervised tasks, so it can learn useful patterns from atomic structures even when expensive physics labels are not available. Across several materials and molecule benchmarks, TriForces improves prediction accuracy, especially when little training data is available. It also makes the model’s internal representations more useful for comparing materials, for example by searching for structures with similar chemistry or similar geometry. This can help researchers adapt atomistic machine learning models more efficiently and explore candidate materials more easily.