TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces is a model-agnostic, three-stream framework that combines self-supervised learning with separated composition and structure representations to significantly enhance the transferability and data efficiency of atomistic graph neural networks for machine learning interatomic potentials.

The Gist

Imagine you are trying to teach a robot chef how to cook.

The Problem: The "One-Size-Fits-All" Chef
Currently, scientists use powerful AI models (called MLIPs) to predict how atoms behave, like how much energy a material has or how hard it is to push atoms around. These models are trained on massive amounts of data from super-computers (DFT).

However, these models have a flaw: they are like a chef who memorized the exact taste of a specific dish but forgot why it tasted that way. If you ask them to cook a slightly different dish (a new type of material), they struggle. They mix up the ingredients (composition) with the shape of the pot (structure). If you change the ingredients, they get confused about the shape, and vice versa. This makes them bad at learning new tasks quickly, especially when you don't have a lot of data to teach them.

The Solution: TriForces (The Three-Stream Kitchen)
The authors introduce TriForces, a new way to build these AI chefs. Instead of one giant brain trying to remember everything at once, they split the brain into three specialized "streams" or departments:

1. The Ingredient Stream (Composition): This department only looks at what is in the pot (e.g., "We have 2 Hydrogens and 1 Oxygen"). It ignores the shape entirely. It learns the chemistry.
2. The Shape Stream (Structure): This department only looks at how the atoms are arranged in space (e.g., "They are in a triangle"). It ignores what the atoms actually are. It learns the geometry.
3. The Interaction Stream: This is the main chef who takes the notes from the Ingredient and Shape departments and combines them to predict the final result (energy or force).

The Secret Sauce: Self-Supervised Learning
Before the model is ever asked to predict a specific property, the authors train it using a game called "Self-Supervised Learning." Think of this as a practice session where the AI has to:

• Denoise: Look at a slightly broken or noisy picture of a molecule and fix it.
• Masking: Cover up an ingredient and guess what it was based on the neighbors.
• Matching: Look at two slightly different versions of the same molecule and realize they are the same thing.

This training forces the AI to organize its knowledge neatly. It learns that "ingredients" belong in one folder and "shapes" in another, rather than jumbling them together.

Why This Matters (The Results)
The paper shows that this new "Three-Stream" kitchen works much better than the old "One-Brain" kitchens:

• Faster Learning: When given a small amount of new data (like 20,000 examples instead of millions), TriForces learns much faster and makes fewer mistakes. It's like a chef who can learn a new recipe after tasting it once, rather than needing to cook it a thousand times.
• Better Memory: The AI doesn't forget what it learned. It can transfer its knowledge from one type of material to another without getting confused.
• Searchable Knowledge: Because the AI keeps "ingredients" and "shapes" separate, you can ask it to find materials that look the same but have different ingredients, or materials with the same ingredients but different shapes. The old models couldn't do this because their knowledge was too mixed up.

In Summary
TriForces is a framework that takes apart the complex job of understanding atoms into three simpler jobs: knowing the ingredients, knowing the shape, and knowing how they work together. By training the AI to keep these jobs separate and practicing with "guessing games" (self-supervised learning), the model becomes a much more flexible, efficient, and accurate tool for discovering new materials.

The authors have released their code and pre-trained models so other scientists can use this "three-stream kitchen" to build better AI for materials science.

Technical Summary

Technical Summary: TriForces

Problem Statement

Machine Learning Interatomic Potentials (MLIPs) based on geometric Graph Neural Networks (GNNs) have achieved high accuracy in predicting atomistic properties when trained on large-scale Density Functional Theory (DFT) datasets. However, their practical utility is hindered by two primary limitations:

1. Inconsistent Transferability: MLIPs often fail to transfer effectively to new chemistries or tasks, particularly when fine-tuning on small, expensive, task-specific datasets. Pre-trained models frequently struggle with simple diagnostic tasks (e.g., identifying crystal systems or majority elements) and exhibit "catastrophic forgetting" or unstable transfer performance across domains.
2. Poor Representation Reusability: Current MLIPs optimize representations for specific prediction targets (energy and forces) rather than for general reuse. Consequently, these representations often entangle composition (chemistry) and structure (geometry), making them unsuitable for exploratory analysis, nearest-neighbor retrieval, or decomposed similarity search. Standard supervised objectives encourage representations sufficient for regression but not organized to preserve accessible composition and structural information.

Methodology: The TriForces Framework

The authors propose TriForces, a model-agnostic framework that augments existing geometric GNNs with a three-stream architecture and multi-objective self-supervised learning (SSL). The core innovation is the explicit factorization of atomistic representations into three distinct components:

1. Three-Stream Architecture

Instead of a single latent vector, TriForces decomposes the node-level representation $h_i$ into three concatenated streams:

• Composition Stream ($h^{comp}$): Encodes chemical information without coordinates. It processes the set of unique atomic elements and their stoichiometric counts using a Transformer with count-weighted attention. This stream preserves absolute element counts to encode system size and energy scale, independent of geometry.
• Structure Stream ($h^{struct}$): Encodes geometry without element identity (type-agnostic). It constructs rotation-invariant local descriptors inspired by SOAP (Smooth Overlap of Atomic Positions), utilizing radial basis functions, spherical harmonics, and multi-scale cutoffs. This stream captures reusable geometric motifs and topological patterns, enhanced by invariant message passing.
• Interaction Stream ($h^{int}$): A standard base geometric GNN (e.g., MACE, eSEN, Orb-v3) that captures the coupling between composition and geometry, preserving the expressiveness of the original architecture.

2. Self-Supervised Pretraining

To organize the latent space and improve transferability, TriForces employs a multi-objective SSL pretraining strategy using stochastic augmentations (position noise, element masking, graph variation, and rotations). The framework combines three complementary objectives:

• Non-Reconstruction (LeJEPA): Aligns embeddings from two augmented views of the same structure at both node and graph levels. This enforces invariance to augmentations and organizes the global latent space without requiring stop-gradients or momentum encoders.
• Denoising: Trains the model to recover clean atomic positions from noisy inputs. This stabilizes geometric representations and implicitly provides rotation augmentation.
• Masking: Predicts masked atomic types based on surrounding geometry and composition. This encourages the model to learn compositional patterns and context.

The final pretraining loss is a weighted sum of these three objectives.

Key Contributions

1. Architectural Decomposition: A three-stream design that explicitly separates composition, structure, and interaction, ensuring both factors are preserved by design rather than entangled.
2. Hybrid Pretraining Strategy: A self-supervised approach combining reconstruction-based objectives (denoising, masking) with latent-prediction learning (LeJEPA) to structure the embedding space for better downstream transfer.
3. Interpretable Retrieval: The ability to perform targeted similarity search in the compositional, structural, or joint embedding spaces, enabling materials comparison based on specific criteria (e.g., chemistry-only or structure-only).
4. Empirical Validation: Extensive experiments across multiple architectures (Orb-v3, eSEN, MACE) and benchmarks (OMat24, MatBench, QM9) demonstrating improved data efficiency, transfer performance, and representation quality.

Results

• Transfer Performance (OMat24): In limited-data regimes, TriForces significantly outperforms baselines. At 20K samples, it reduces energy Mean Absolute Error (MAE) by 57% compared to the base model. It improves force MAE across all sample sizes and reduces stress errors.
• Data Efficiency: TriForces achieves lower errors at every dataset size (20K to 2M samples), with the most significant gains observed in low-data settings.
• Benchmark Performance:
  • MatBench: TriForces variants achieve the best overall results on 6 out of 8 tasks, outperforming both self-supervised and DFT-labeled pretraining baselines. For instance, Phonons MAE improved from 57.8 to 19.5 cm$^{-1}$.
  • MatBench Discovery: TriForces eSEN-sm achieves comparable energy MAE to a much larger eSEN-30M-OAM model while using 60% fewer parameters and training up to 5$\times$ faster.
  • QM9: Pretraining on diverse chemical inputs (bulk + molecules) consistently lowers MAE compared to bulk-only or no-SSL baselines.
• Representation Quality: Linear probing on frozen embeddings shows that TriForces preserves fundamental information (crystal system, majority element, coordination number) that standard MLIPs lose. TriForces achieves 96–100% accuracy on crystal system and majority element classification, whereas baselines struggle (55–73%).
• Retrieval: The framework enables effective k-NN retrieval where the composition stream excels at element-set recall and the structure stream excels at space-group recall, a capability absent in single-stream models.

Significance and Claims

The paper positions TriForces not merely as a self-supervised method, but as an architectural framework whose representations are further enhanced by SSL.

• Regime Dependence: The authors claim that stream factorization provides the dominant gains in large-scale supervised settings, while SSL is most valuable for low-data transfer, representation organization, and retrieval tasks.
• Decoupling: By separating composition and structure, TriForces addresses the "brittle transfer" and "hard-to-reuse" problems of current foundation models. It allows models to learn representations that are organized for analysis (e.g., probing, retrieval) in addition to prediction.
• Practicality: The framework is model-agnostic and plug-in, making it immediately applicable to existing or new atomistic architectures. The authors release pretrained checkpoints and code to facilitate reuse in downstream atomistic modeling.

The work suggests that future atomistic foundation models should move beyond single-stream prediction architectures toward factorized representations that explicitly preserve the distinct physical factors of chemical systems.