Latent Diffusion Pretraining for Crystal Property Prediction
Abstract
Fast and accurate prediction of crystal properties is a central challenge in new materials design. Graph neural networks and Transformer-based models have emerged as powerful tools for this task due to their ability to encode the local structural environment of atoms within a crystal. However, these models are data-hungry and in practice labeled data for crystal properties are very scarce. Pretraining–finetuning strategies, particularly those based on diffusion models, have shown promise in addressing these limitations. In this work, we introduce a novel latent-diffusion based pretraining framework CrysLDNet, designed to mitigate the data scarcity. Our approach integrates a Variational Autoencoder (VAE) with a diffusion model during the pretraining stage. The VAE encoder maps 3D crystal structures into a smooth latent space, within which the diffusion process is applied. This latent diffusion pretraining enables the graph encoder to effectively capture structural and chemical semantics from large-scale unlabeled data, which can then be finetuned for specific property prediction tasks. Comprehensive experiments on popular DFT datasets for property prediction reveal that CrysLDNet significantly outperforms both training-from-scratch and pretrained baselines, with improvements of 4.26% and 4.90% on the JARVIS and MP datasets. Additionally, the learned representations remain robust in sparse-data conditions and are expressive enough to correct DFT errors when finetuned with limited experimental data. Code is available at https://github.com/shrimonmuke0202/CrysLDNet.git.
Lay Summary
Modern artificial intelligence can help scientists discover new materials much faster than traditional trial-and-error experiments. One important challenge is predicting the properties of crystal materials — such as formation energy or bandgap — from their atomic structure. However, collecting labeled training data for these properties is expensive and time-consuming, which limits the performance of many existing AI models. In this work, we introduce CrysLDNet, a new AI framework that learns from large collections of unlabeled crystal structures before being adapted to specific prediction tasks. Instead of directly learning from complicated raw crystal data, the model first compresses crystal structures into a simpler hidden representation and then improves this representation using a diffusion-based learning process inspired by recent advances in generative AI. We show that this approach helps the model better understand the relationship between atomic arrangement and material properties. Across several benchmark datasets, CrysLDNet consistently predicts crystal properties more accurately than existing methods, especially when only a small amount of labeled data is available. The framework is also flexible and can easily work with future AI architectures, making it a promising tool for accelerating materials discovery and design.