Technical Summary: Crystal Structure Based Graph Neural Network and LLM Multi-Modal Approach for Material Magnetic Ordering Prediction

Problem Statement

The determination of material properties, specifically magnetic ordering (Non-magnetic, Ferromagnetic, Antiferromagnetic, and Ferrimagnetic), is critical for material engineering applications such as shielding, sensors, and actuators. Traditional experimental determination is accurate but time-consuming and costly. While Density Functional Theory (DFT) offers a mathematical alternative, its accuracy is limited by approximation functions and struggles with complex compound structures. Although recent AI/ML approaches, particularly Graph Neural Networks (GNNs) and Large Language Models (LLMs), have shown promise in predicting material properties, there is a need to effectively integrate structural data with semantic descriptions to improve prediction accuracy, especially in the context of imbalanced datasets where certain magnetic classes are underrepresented.

Methodology

The study proposes a novel multi-modal architecture that fuses crystal structure information with semantic text descriptions to predict magnetic ordering. The model consists of four integrated blocks: a GNN, an LLM, an information fusion module, and a classification head. All blocks are trained jointly.

1. Model Architecture

• GNN Block (ALIGNN): The model utilizes the Atomistic Line Graph Neural Network (ALIGNN). Material structures are extracted in Pymatgen JSON format and encoded using Jarvis Atom functionality. ALIGNN employs two distinct graph structures:
  • A standard Graph Convolutional Network (GCN) where atoms are nodes and bonds are edges, utilizing CGCNN embeddings (92-dimensional vectors based on periodic table properties).
  • A line GCN where bonds are nodes and bond angle triplets are edges, allowing the model to capture bond angle data which previous models often missed.
• LLM Block (GPT-2): Material compounds are converted into natural language text descriptions using the ChemNLP implementation. These descriptions include atomic formulas, molecular weights, lattice parameters, XRD peaks, density, and bond distances. The GPT-2 base model (124.64M parameters) processes these semantic descriptors.
• Fusion Strategies: Three fusion strategies were evaluated to combine the outputs of the GNN and LLM:
  1. Early Fusion: Serialized GNN graph data and tokenized text descriptors are fused as input to the LLM.
  2. Intermediate Fusion: Output features from the GNN and LLM are concatenated before being fed to the classification head.
  3. Gated Late Fusion: A learnable gate balances the logits from both branches for each class before the final prediction.
• Classification: A Multi-Layer Perceptron (MLP) based head performs the final classification into the four magnetic ordering categories.

2. Data Preparation

The dataset comprises 200,431 compounds sourced from the Materials Project, Magndata, and a specific set of uranium compounds from Broyles et al. [13]. The data was split into 80% training, 10% validation, and 10% testing.

• Class Imbalance: The dataset is highly imbalanced: Non-magnetic (59.48%), Ferromagnetic (32.06%), Ferrimagnetic (6.61%), and Antiferromagnetic (1.85%).
• Training Configuration: To address the imbalance, Class-weighted Cross Entropy Loss was used with weights of 1.0 (NM), 1.5 (FM), 20.0 (AFM), and 5.0 (FiM). The models were trained for 30 epochs using the AdamW optimizer on dual Nvidia 3090 GTX GPUs.

Key Contributions

1. Multi-Modal Integration: The paper introduces a framework that simultaneously leverages the geometric precision of ALIGNN (for crystal structures) and the semantic understanding of GPT-2 (for text descriptions) to predict magnetic ordering.
2. Joint Training: Unlike hybrid approaches that use pre-trained models solely for embedding, this study trains the GNN, LLM, and fusion layers together.
3. Fusion Strategy Evaluation: The study systematically compares Early, Intermediate, and Gated Late fusion strategies, demonstrating that information fusion yields superior accuracy compared to standalone models.
4. Handling Imbalanced Data: The approach demonstrates improved performance on a highly imbalanced dataset through specific loss weighting and multi-modal feature integration.

Results and Analysis

The proposed ALIGNN-LLM multi-modal approach was benchmarked against standalone ALIGNN, standalone GPT-2, and other state-of-the-art models (CGCNN, coGN, LLM-Prop, MatBERT, AMD-LLAMA).

• Overall Accuracy: The intermediate-fusion multi-modal approach achieved the highest overall accuracy of 88.6%. This outperformed the standalone ALIGNN model (83.6%) and the standalone LLM model (86.6%).
• Per-Class Performance:
  • Non-magnetic & Ferromagnetic: The multi-modal approach achieved high precision and recall (e.g., 95.4% F1 for Non-magnetic, 84.2% F1 for Ferromagnetic), generally outperforming or matching the standalone models.
  • Antiferromagnetic (AFM): This class remained the most challenging due to data scarcity. While the multi-modal approach (41.1% F1) improved upon ALIGNN (36.3% F1), it slightly underperformed compared to the standalone GPT-2 (46.8% F1) for this specific class.
  • Ferrimagnetic: The multi-modal approach achieved the best F1 score (58.3%) compared to ALIGNN (54.5%) and GPT-2 (54.4%).
• Model Complexity: The ALIGNN model (4.3M parameters) offers a favorable trade-off between computational cost and accuracy, outperforming larger LLMs like LLM-Prop (36M) and MatBERT (110M) in specific contexts, making it suitable for edge devices.
• Fusion Comparison: All three fusion strategies yielded similar overall accuracies (ranging from 88.6% to 88.77%), with the intermediate feature fusion showing robust performance across most classes.

Significance and Claims

The paper claims that the proposed multi-modal approach represents a significant advancement in material property prediction by effectively combining structural and semantic data. The results demonstrate that fusing information from GNNs and LLMs provides better overall accuracy than using either modality in isolation.

The authors modestly note that while the multi-modal approach improves accuracy, it comes at the cost of increased memory usage and computational complexity. They acknowledge that the Antiferromagnetic class remains difficult to predict due to data imbalance and suggest that future work should focus on balancing the dataset with more samples. Additionally, the authors propose that this architecture could be adapted for regression problems, multi-property prediction, and potentially material generation by utilizing the LLM's decoding capabilities to generate new materials with desired properties. The study concludes that the approach highlights the potential of multi-modal AI in accelerating materials discovery.