MatterTune: A Modular AFM Fine-Tuning Platform

• MatterTune is a modular framework that standardizes the fine-tuning of pre-trained atomistic foundation models, enhancing reproducibility in materials simulations.
• It organizes workflows into clear subsystems for model adaptation, data formatting, training integration with PyTorch Lightning, and deployment via ASE-compatible interfaces.
• The platform demonstrates robust few-shot transfer, scalable molecular dynamics, and superior benchmark performance in property prediction tasks.

MatterTune is an integrated, user-friendly platform for fine-tuning atomistic foundation models (AFMs) and embedding them into downstream materials informatics and simulation workflows. It was introduced to address a practical gap in chemical and materials machine learning: pre-trained geometric models had become increasingly capable, but standard software infrastructure for fine-tuning, reproducibility, and workflow integration remained limited. In its stated scope, MatterTune is a modular and extensible framework rather than a single model, and it natively supports state-of-the-art backbones including ORB, MatterSim, JMP, and EquformerV2 (Kong et al., 14 Apr 2025).

1. Origin, motivation, and problem setting

MatterTune emerges from the convergence of two developments in atomistic machine learning. First, geometric machine learning models such as graph neural networks achieved strong performance in property prediction, high-throughput virtual screening, and atomistic simulations. Second, these models remained constrained by high data requirements, a particularly acute limitation in materials science, where many relevant datasets contain only hundreds or fewer labeled examples. The resulting emphasis on pre-trained models led to the rise of atomistic foundation models: large models, often equivariant GNNs or transformers, pre-trained on diverse large-scale atomistic datasets and then fine-tuned for downstream applications (Kong et al., 14 Apr 2025).

The platform’s motivating claim is therefore infrastructural rather than architectural. MatterTune was designed to lower the barriers to adoption of AFMs by standardizing fine-tuning and by connecting these models to established simulation and informatics pipelines. A common misconception is to treat MatterTune itself as an interatomic potential or a new AFM backbone. In the source formulation, it is instead a framework that loads, adapts, fine-tunes, and deploys multiple pre-trained backbones in data-scarce, application-specific settings (Kong et al., 14 Apr 2025).

2. Architectural organization and core abstractions

MatterTune is organized as a modular, separated-subsystem architecture with four principal subsystems: model, data, trainer, and application. The model subsystem introduces a Backbone Abstraction that provides standardized APIs such as model_forward and input conversion, irrespective of whether the underlying foundation model is a GNN or a transformer. It also manages output heads for property prediction and allows those heads to be added or customized even when the original backbone does not provide them (Kong et al., 14 Apr 2025).

The data subsystem is built around a Data Abstraction that converts heterogeneous sources, including XYZ, JSON, and ASE DB, into a standardized internal representation based on ase.Atoms objects. In the platform’s formalism, the dataset is represented as

$D: \mathbb{N} \rightarrow S$

where $S$ is the standardized atomic structure space. This standardization is central to MatterTune’s claim of interoperability, because it enables FM-specific input formatting to be treated as a downstream transformation rather than as a property of the raw dataset itself (Kong et al., 14 Apr 2025).

The trainer subsystem is integrated with PyTorch Lightning and handles training, validation, checkpointing, distributed computing, and user-injected custom logic through Lightning callbacks. The application subsystem then exposes deployment-facing interfaces, notably MatterTuneCalculator, which inherits from the ASE calculator abstraction, and MatterTunePropertyPredictor, which supports batch and parallel inference for screening workflows. Across these layers, a separate Property Abstraction decouples the definition of physical targets—such as energies, forces, band gaps, and stresses—from their implementation, with an emphasis on type safety and extensibility (Kong et al., 14 Apr 2025).

3. Supported foundation models and fine-tuning methodology

In its reported state, MatterTune supports ORB, MatterSim, JMP, and EquformerV2. The platform summary specifies ORB-v1 as a denoising pre-trained foundation model with 25.2M parameters and 32.1M samples, MatterSim-v1 as an equivariant model with 4.55M parameters and 17M samples, JMP-S as a 30M-parameter model trained on 120M samples, JMP-L as a 235M-parameter model, and EquformerV2 variants including eqV2-S at 31.2M parameters and eqV2-M at 86.6M parameters. The architecture is described as plug-and-play, so future AFMs can be incorporated without redesigning the core framework (Kong et al., 14 Apr 2025).

Fine-tuning in MatterTune is deliberately configurable. The platform supports Adam, AdamW, and SGD; schedulers including linear, step, exponential, cosine, and reduce-on-plateau; parameter-group customization for assigning different learning rates to different model components; combinable schedulers such as cosine warm-up; optional exponential moving average; and multiple normalization strategies, including mean/std, RMS, and domain-specific normalization such as element-wise regression. These controls are presented as especially important for low-data regimes and multi-target objectives involving combinations such as energy, force, and stress prediction (Kong et al., 14 Apr 2025).

Few-shot transfer is a central design target. MatterTune explicitly reports robust fine-tuning with as few as 30 data points. This emphasis is not incidental: it encodes the platform’s intended role as an interface between large pre-trained atomistic representations and the small, specialized datasets that dominate many experimental and computational materials settings (Kong et al., 14 Apr 2025).

4. Downstream tasks and workflow integration

MatterTune is designed to bridge AFMs with mainstream materials simulation and informatics workflows. The supported downstream tasks listed in the platform summary include materials property prediction, molecular dynamics simulations, structure prediction and geometry optimization, high-throughput screening for materials discovery, zero-shot and few-shot transfer learning, and active learning and representation-space analysis such as t-SNE. The framework’s role is therefore not limited to supervised regression; it also targets deployment contexts in which a fine-tuned AFM becomes an executable component in screening or simulation loops (Kong et al., 14 Apr 2025).

ASE integration is a particularly important element of this strategy. By exposing an ASE-compatible calculator, MatterTune allows fine-tuned models to be used directly in molecular dynamics and relaxation pipelines that already rely on the Atomic Simulation Environment. The same abstraction logic extends to high-throughput inference through MatterTunePropertyPredictor, which is intended for database-scale evaluation of candidate structures. The summary further states that the framework is future-proofed for LAMMPS and other workflow tools, suggesting an explicit orientation toward interoperability with established atomistic software ecosystems (Kong et al., 14 Apr 2025).

This integration layer also clarifies the platform’s position relative to conventional MLIPs. MatterTune does not replace downstream simulation codes; it standardizes how pre-trained AFMs are adapted and inserted into those codes. A plausible implication is that its main contribution lies in workflow continuity: data ingestion, model wrapping, fine-tuning, and deployment are treated as a single configurable pipeline rather than as separate bespoke scripts (Kong et al., 14 Apr 2025).

5. Benchmarks, case studies, and empirical behavior

The platform summary reports several benchmark and case-study results. On MatBench property-prediction tasks, MatterTune fine-tuned JMP-S, ORB-V2, and Equiformer-31M-mp and is reported to reproduce or exceed original FM-reported benchmarks in most cases without task-specific tuning. It further states that the best fine-tuned models via MatterTune outperform leading from-scratch non-foundation-model baselines in most tasks. Representative values given in the summary table include an MP Gap result of 0.093 eV for ORB-V2 (MT) and a dielectric result of 0.111 for EqV2-31M-mp (MT) (Kong et al., 14 Apr 2025).

In out-of-distribution high-throughput screening, MatterTune was used to predict band gaps for 381,000 new structures from GNoME data. The summary attributes to ORB-V2 an MAE of 0.039 eV and an accuracy of 98.80% on this task. In a few-shot liquid-water study, the framework fine-tuned AFMs using 30 samples versus a 900-sample training condition and reports that MatterSim achieved an energy MAE of 1.20 meV/atom with 900 samples and 1.21 meV/atom with 30 samples. For molecular dynamics, MatterSim and EqV2 are reported to reproduce the experimental oxygen–oxygen radial distribution function, while JMP-S and ORB-V2 did not, which the summary associates with the importance of energy-conserving architectures for simulation stability (Kong et al., 14 Apr 2025).

Additional cases include zero-shot structure prediction on Matbench-Discovery, where MatterTune is reported to reproduce leaderboard-level results on crystal relaxation tasks, and t-SNE-based latent-space visualization, which revealed distinct clustering by element group for MatterSim and JMP, with more diffuse clustering for ORB and Equformer. Taken together, these studies position MatterTune as both an engineering framework and an experimental substrate for comparing representation quality, transfer behavior, and deployment stability across AFM families (Kong et al., 14 Apr 2025).

6. Scalability extensions and large-scale molecular dynamics

A subsequent development explicitly implemented large-scale acceleration strategies within MatterTune. The later study introduces a two-part workflow: Message-Passing Depth Pruning (MPDP), which removes low-contribution message-passing layers from AFM backbones, and a graph-partitioned, GPU-distributed strategy for scalable inference. The paper states that these methods were implemented and demonstrated within the AFM fine-tuning platform MatterTune, thereby extending the platform from fine-tuning and deployment infrastructure into a vehicle for scalable AFM molecular dynamics (Kong et al., 25 Sep 2025).

The pruning argument is based on the observation that deeper message-passing layers often contribute little additional representational value because of over-smoothing. For a backbone with $M$ message-passing layers, MPDP truncates the model after the first $K$ layers, with $K < M$, often $K=1$ or $K=2$. The study characterizes redundancy using Mean Feature Shift,

$$\text{MFS}^{(i \to i+1)} = \frac{1}{N} \left\| H^{(i+1)} - H^{(i)} \right\|_F,$$

and Centered Kernel Alignment, and then couples the pruned model to a partitioned inference pipeline. Because the effective receptive field becomes

$r_\text{eff} = K \cdot r_c,$

smaller $K$ directly reduces halo size during graph partitioning and lowers memory duplication across GPUs (Kong et al., 25 Sep 2025).

The partitioning workflow first assigns root nodes to partitions, then expands each subgraph via $S$0-hop breadth-first search to preserve local neighborhoods. Spatial partitioning is used for bulk systems, while METIS graph partitioning is used for heterogeneous systems. The study emphasizes that no inter-GPU communication is required at inference: each GPU processes its own expanded subgraph and the outputs are subsequently aggregated. The implementation uses PyTorch Lightning Distributed Data Parallel and retains an ASE-compatible calculator path, preserving continuity with MatterTune’s original workflow orientation (Kong et al., 25 Sep 2025).

The reported performance figures are substantial. On Nvidia A100-40GB GPUs, the pruned MatterSimV1-1M model with one message-passing layer achieved 58.63 $S$1s per atom per step for 421,824 atoms on 1 GPU, 33.18 on 2 GPUs, and 23.22 on 4 GPUs; 53.83 for 943,296 atoms on 1 GPU; and 68.65 for 5,184,000 atoms on 1 GPU. The study states that the combined pruning and graph-partitioning approach supports million-atom simulations on both single and multiple GPUs and enables task-specific large-scale simulations at nanosecond timescales with AFM-level accuracy. Physical fidelity is illustrated by Li$S$2PO$S$3 quenching, silicon–silica segregation, and nanoparticle encapsulation on oxide surfaces involving over 42,000 atoms (Kong et al., 25 Sep 2025).