In context learning Foundation models for Materials Property Prediction with Small datasets

TL;DR

The paper tackles data-efficient materials property prediction by introducing a unified in-context learning foundation model (ICL-FM) that integrates both composition-based and structure-aware representations. It combines the TabPFN foundation model with Magpie and MagpieEX descriptors and structural embeddings from graph neural networks (ALIGNN and CGCNN) to achieve training-free, data-efficient property predictions on small datasets. Empirical evaluations across the MatBench benchmark suite and a standalone lattice thermal conductivity dataset show competitive or superior performance relative to state-of-the-art methods, including a 9.93% improvement on phonon frequency prediction when using MagpieEX. Interpretability analyses (t-SNE and SHAP) reveal that the ICL-FM refines disparate feature spaces into physics-consistent manifolds and distributes attributions across a broad set of chemical and bonding descriptors, aligning with known physical principles such as lattice stiffness and bond ionicity. Overall, the work establishes ICL-FM as a generalizable, data-efficient paradigm for materials informatics, with strong potential for rapid, plug-and-play integration across composition- and structure-based property prediction tasks, while acknowledging limitations on structure-dominated properties and suggesting future extensions to tailor coupling strategies to physical complexity.

Abstract

Foundation models (FMs) have recently shown remarkable in-context learning (ICL) capabilities across diverse scientific domains. In this work, we introduce a unified in-context learning foundation model (ICL-FM) framework for materials property prediction that integrates both composition-based and structure-aware representations. The proposed approach couples the pretrained TabPFN transformer with graph neural network (GNN)-derived embeddings and our novel MagpieEX descriptors. MagpieEX augments traditional features with cation-anion interaction data to explicitly measure bond ionicity and charge-transfer asymmetry, capturing interatomic bonding characteristics that influence vibrational and thermal transport properties. Comprehensive experiments on the MatBench benchmark suite and a standalone lattice thermal conductivity (LTC) dataset demonstrate that ICL-FM achieves competitive or superior performance to state-of-the-art (SOTA) models with significantly reduced training costs. Remarkably, the training-free ICL-FM outperformed sophisticated SOTA GNN models in five out of six representative composition-based tasks, including a significant 9.93\% improvement in phonon frequency prediction. On the LTC dataset, the FM effectively models complex phenomena such as phonon-phonon scattering and atomic mass contrast. t-SNE analysis reveals that the FM acts as a physics-aware feature refiner, transforming raw, disjoint feature clusters into continuous manifolds with gradual property transitions. This restructured latent space enhances interpolative prediction accuracy while aligning learned representations with underlying physical laws. This study establishes ICL-FM as a generalizable, data-efficient paradigm for materials informatics.

Figures (9)

• Figure 1: TabPFN based ICL-FM framework for materials property prediction: (left) The transformer based TabPFN model is pretrained using millions of synthetic regression datasets; (top/bottom): material representations can be composition and structural descriptors or embeddings extracted from GNN models; (right): material feature-property pairs are fed to the FM model for training-free in-context property prediction.
• Figure 2: t-SNE visualization comparing the raw input feature spaces and foundation model (FM) learned embeddings for both Magpie and MagpieEX representations over phonons dataset fold 1. The colors of points indicate the material property values (phonon frequency). (a) The original Magpie feature space shows scattered and less distinct clusters. (b) The FM embedding space learned from Magpie inputs exhibits more structured separation, indicating improved representation learning. (c) The MagpieEX feature space provides a richer compositional basis. (d) The FM embedding space using MagpieEX inputs produces clearer and more distinct material clusters, highlighting the benefit of extended descriptors.
• Figure 3: Comparison of SHAP analyses for Magpie and MagpieEX feature inputs of their aggregate feature importance. (a) SHAP summary plot showing the contribution of individual Magpie features to the FM’s predictions. (b) SHAP summary plot for MagpieEX features, illustrating the additional influence of extended compositional descriptors on model interpretability.
• Figure 5: t-SNE visualization of latent representations with LTC values (fold 4). Each point corresponds to a material, colored by its measured $\kappa$. (a) The raw compositional and structural features form disjoint clusters with weak separation by conductivity. (b) After FM processing, the embeddings become smoother and more continuous, exhibiting clearer gradients with respect to $\kappa$. This transformation indicates that the FM extracts latent physical correlations that align compositional–structural descriptors with the underlying thermal transport physics.
• Figure 6: t-SNE visualization of sample distributions in the structural representation spaces for lattice thermal conductivity prediction (fold 4 of our 5-fold cross-validation experiment). (a) Sample distribution in the Last-layer ALIGNN feature space. This feature set capture certain correlations of crystal structures and LTC with overall continuous space. But the continuity of LTC values are not high. (b) Sample distribution in the FM-embedding space with Last-layer ALIGNN features: the embedding space become smoother and LTC conductivity is much better than (a). (c) Sample distribution in the Multi-layer fusion feature space, obtained by combining intermediate ALIGNN embeddings. The sample continuity and LTC value continuity are all low. (d) Sample distribution in the FM-embedding space with the fused features. The FM significantly enhances the cluster continuity and alignment with thermal conductivity values. Together, these panels illustrate that the FM acts as a physics-aware feature refiner, improving the structural–property relationship within learned manifolds.