Multi-Task Fine-Tuning Enables Robust Out-of-Distribution Generalization in Atomistic Models

TL;DR

This work tackles the problem of unreliable out-of-distribution generalization in pretrained atomistic models when adapting to downstream properties. It reveals that standard fine-tuning induces representation collapse, erasing pretrained chemical priors and harming OOD performance. The authors propose multi-task fine-tuning (MFT), which jointly optimizes the downstream property objective with a force-field objective inherited from pretraining, preserving priors while enabling task-specific updates. Across molecular and materials benchmarks, MFT closes the ID–OOD gap, achieves state-of-the-art OOD performance without sacrificing ID accuracy, and demonstrates strong data efficiency, making it a practical approach for robust molecular and materials discovery. The results offer mechanistic insight into the balance between plasticity and stability during adaptation and suggest MFT as a general strategy for safe, data-efficient fine-tuning of atomistic foundation models.

Abstract

Accurate de novo molecular and materials design requires structure-property models that generalize beyond known regimes. Although pretrained atomistic models achieve strong in-distribution accuracy after fine-tuning, their reliability under out-of-distribution (OOD) conditions remains unclear. We identify a critical failure mode in downstream adaptation: standard fine-tuning induces representation collapse, erasing pretrained chemical and structural priors and severely degrading OOD performance. To address this limitation, we propose multi-task fine-tuning (MFT), which jointly optimizes downstream property prediction with a physically grounded force-field objective inherited from pretraining. This approach preserves essential chemical priors while enabling task-specific adaptation. Across molecular and materials benchmarks, MFT consistently improves OOD generalization, approaching the theoretical limit set by in-distribution accuracy, while outperforming standard fine-tuning, training from scratch, and state-of-the-art task-specific models. These results establish safe adaptation as a central requirement for large atomistic models and position MFT as a practical and data-efficient pathway toward robust molecular and materials discovery.

Figures (8)

• Figure 1: Schematic illustration of pretraining and downstream adaptation strategies. The left panel depicts the pretraining phase. The right panel compares four downstream adaptation strategies: (a) training from scratch (Scratch), in which both the backbone and the task-specific prediction head are randomly initialized and optimized end-to-end; (b) fine-tuning (FT), where the backbone is initialized from the pretrained model and all parameters are jointly updated; (c) linear probing (LP), which freezes the pretrained backbone and optimizes only a randomly initialized task-specific head; and (d) multi-task fine-tuning (MFT), which jointly optimizes the pretrained backbone together with two heads: a task-specific head for property prediction and an auxiliary force-field head inherited from the pretraining stage.
• Figure 2: Analysis of representation collapse across downstream adaptation strategies. (a–c) Relative performance improvements of fine-tuning (FT), linear probing (LP), and multi-task fine-tuning (MFT) compared with training from scratch (Scratch) baseline. (d–f) and (g–i) t-SNE visualizations of atomic (node) and edge representations, respectively, projected into two-dimensional spaces. The clarity and separation of atomic and edge representation clusters are quantified using the Davies–Bouldin (DB) index, normalized with respect to the pretrained model. Larger relative DB index values indicate more severe degradation in cluster cohesion and separability, providing quantitative evidence of collapse in the pretrained representation space.
• Figure 3: Data efficiency and ablation study of multi-task fine-tuning (MFT). (a) Data efficiency on the QM9 band-gap (GAP) prediction task. (b) Data efficiency on the MatBench dielectric prediction task. In (a,b), MFT is compared with training from scratch (Scratch) and standard fine-tuning (FT); in-distribution (ID) and out-of-distribution (OOD) test RMSEs or MAEs are shown as functions of training set size. (c–e) Ablation study of MFT on the QM9 HOMO, LUMO, and GAP prediction tasks. (c) FT baseline. (d) Ablation of transferred pretrained components. (e) Ablation of auxiliary force-field task selection. "Avg. $\uparrow$RMSE" denotes the average RMSE increase across the HOMO, LUMO, and GAP prediction tasks relative to the baseline MFT configuration.
• Figure 4: Analysis of representation collapse of JMP-L. Relative performance improvements of fine-tuning (FT) over the training from Scratch baseline.
• Figure 5: t-SNE visualizations of atomic (node) and edge representations representations of the last backbone layer of JMP-L on QM9 LUMO task, reduced to 2-dimension space. The Davies–Bouldin (DB) index values are normalized to the pretrained model baseline.