Robust Lightweight Graph Neural Network Framework for Accelerating Crystal Structure Prediction

This work presents a crystal structure prediction framework that employs a structural search using a derivative-free optimization method, with a supervised Graph Neural Network (GNN) model as the energy evaluator. We address the limitations of existing GNN-based crystal structure prediction (CSP) frameworks and propose methods for designing a robust and computationally efficient predictor. In particular, we first highlight the often-overlooked sensitivity of GNN models to weight initialization in crystal structure prediction, and to address this, we introduce a model selection framework that consistently identifies an appropriate GNN model for downstream crystal structure prediction tasks. Using this framework, we conduct a meaningful comparison of multiple GNN architectures for CSP involving a Bayesian optimization approach. Furthermore, we propose a data augmentation strategy that incorporates unrelaxed structures in the supervised training process, and additionally explore the impact of unsupervised GNN pretraining with and without augmentation on crystal structure prediction. Finally, we demonstrate that our proposed crystal structure prediction framework, in conjunction with the lightweight GNN architecture CGCNN, can achieve a level of performance comparable to that of more complex GNN architectures, which are typically computationally expensive to train and infer. The approaches introduced in this work are generic and can be extended to any GNN-based crystal structure prediction framework, paving the way for developing novel and high-throughput crystal structure predictors in the future.