Source-linked AI summary
MOFormer: Self-Supervised Transformer model for Metal-Organic Framework Property Prediction
Zhonglin Cao, Rishikesh Magar, Yuyang Wang, Amir Barati Farimani
TL;DR
MOF discovery requires efficient prediction across an enormous chemical space, while existing approaches can depend on costly simulations or optimized 3D structures. The paper introduces MOFormer, a Transformer using MOFid strings with self-supervised pretraining alongside CGCNN. Pretraining improves prediction accuracy, and MOFormer outperforms another structure-agnostic method on several tasks while offering data-efficiency advantages when training data is limited.
Problem
MOF property screening must search a vast design space, but existing approaches can require optimized 3D structures and extensive feature-selection expertise.
Method
MOFormer predicts MOF properties from MOFid text and is jointly pretrained with CGCNN through a self-supervised framework on more than 400k MOF structures.
Results
21.4% higher accuracy on band gap prediction and 35%-48% higher accuracy on various gas predictions are achieved by MOFormer versus Stoichiometric-120; pretraining also improves model accuracy.
Takeaways & Limitations
MOFormer provides a structure-agnostic approach for efficient MOF property prediction and can be more data-efficient than CGCNN when training data is limited.
Takeaways & Limitations
MOFormer does not directly include neighborhood-of-atoms information that is critical in some property-prediction tasks.
Abstract
from arXiv · showhide
Metal-Organic Frameworks (MOFs) are materials with a high degree of porosity that can be used for applications in energy storage, water desalination, gas storage, and gas separation. However, the chemical space of MOFs is close to an infinite size due to the large variety of possible combinations of building blocks and topology. Discovering the optimal MOFs for specific applications requires an efficient and accurate search over an enormous number of potential candidates. Previous high-throughput screening methods using computational simulations like DFT can be time-consuming. Such methods also require optimizing 3D atomic structure of MOFs, which adds one extra step when evaluating hypothetical MOFs. In this work, we propose a structure-agnostic deep learning method based on the Transformer model, named as MOFormer, for property predictions of MOFs. The MOFormer takes a text string representation of MOF (MOFid) as input, thus circumventing the need of obtaining the 3D structure of hypothetical MOF and accelerating the screening process. Furthermore, we introduce a self-supervised learning framework that pretrains the MOFormer via maximizing the cross-correlation between its structure-agnostic representations and structure-based representations of crystal graph convolutional neural network (CGCNN) on >400k publicly available MOF data. Using self-supervised learning allows the MOFormer to intrinsically learn 3D structural information though it is not included in the input. Experiments show that pretraining improved the prediction accuracy of both models on various downstream prediction tasks. Furthermore, we revealed that MOFormer can be more data-efficient on quantum-chemical property prediction than structure-based CGCNN when training data is limited. Overall, MOFormer provides a novel perspective on efficient MOF design using deep learning.
Introduction
MOFs span an enormous design space, making accurate and efficient property screening difficult. MOFormer addresses this challenge with a Transformer that uses MOFid strings instead of optimized 3D structures, while self-supervised pretraining incorporates structural information.
- Motivation: MOFs combine diverse metal nodes, organic linkers, and topologies, creating a vast space of structures and properties to explore.Their applications include gas adsorption, water harvesting and desalination, and energy storage.
- Existing approaches: Simulation-based screening can evaluate many hypothetical MOFs but requires computationally expensive calculations for each new structure.Large datasets contain atomic structures and properties suitable for training data-driven prediction models.
- Existing approaches: Handcrafted-feature MLPs can be fast but have limited accuracy and depend on extensive domain knowledge and optimized 3D structures.These requirements make the approach less generic.
- Existing approaches: CGCNN learns chemical information from crystal graphs but requires optimized 3D atomic structures and can be memory-inefficient for very large MOFs.Some MOFs contain hundreds or thousands of atoms.
- MOFormer: MOFormer uses a modified MOFid text representation, preserving chemical and topological information without requiring a MOF's 3D atomic structure.MOFid encodes building-block chemistry through SMILES and topology and catenation through RCSR-derived codes.
- MOFormer: The model enables faster and more flexible exploration of hypothetical MOFs while predicting diverse MOF properties from structure-agnostic inputs.Its Transformer architecture processes the MOFid sequence directly.
- Self-supervised learning: Self-supervised pretraining jointly trains MOFormer and CGCNN on more than 400k MOF structures to improve downstream representations and prediction performance.The study also visualizes latent representations and attention weights and compares data efficiency when training data is limited.
Methods
MOFormer encodes MOFid text with a Transformer and uses self-supervised alignment with CGCNN representations to incorporate structural information. The pretrained models are then fine-tuned for downstream MOF property prediction across benchmark datasets.
- MOFid and Transformer encoding: MOFormer tokenizes MOFid strings, combines token embeddings with positional encoding, and processes them through Transformer encoder layers.MOFid includes SMILES-based building-block information and separately encoded topology and catenation sections joined by a separator token.
- MOFid and Transformer encoding: The Transformer encoder uses multi-head scaled dot-product attention, an MLP, residual connections, and layer normalization.Learnable query, key, and value projections allow attention heads to capture contextual information across representation subspaces.
- MOFid and Transformer encoding: The first-token [CLS] embedding represents the full sequence and is passed to an MLP regression head for downstream prediction.Self-attention incorporates whole-sequence information into token embeddings, supporting use of the first token as a sequence representation.
- Self-supervised pretraining: Self-supervised pretraining aligns MOFormer text representations with CGCNN structure representations for the same MOF using Barlow Twins loss.The two encoders produce projected embeddings, whose cross-correlation matrix is optimized toward the identity matrix.
- Scope and limitation: Text-string inputs lack geometric and atomic-neighborhood information, creating an information bottleneck that can limit Transformer performance relative to graph-based models.The paper identifies this structure-agnostic input limitation as a motivation for incorporating structural information during pretraining.
- Self-supervised pretraining: The pretraining dataset contains 413535 unique MOFs with both 3D structures and MOFids, after duplicate removal.The models are subsequently initialized with pretrained weights and fine-tuned separately for downstream prediction tasks.
Results and discussion
Across QMOF and hMOF benchmarks, pretrained MOFormer performs strongly on band-gap prediction and becomes more data-efficient, while structure-based CGCNN remains better for CO2 adsorption. Learned representations and attention highlight complementary emphasis on topology, adsorption-relevant structure, and key MOF tokens.
- QMOF band-gap prediction: CGCNN > MOFormer > SOAP > Stoichiometric-120 for QMOF band-gap prediction, while MOFormer has 21.2% lower MAE than Stoichiometric-120.MOFormer also outperforms structure-based SOAP with a smaller training set, indicating useful feature extraction from MOFid.
- QMOF band-gap prediction: Pretraining reduces QMOF band-gap MAE by 6.79% for CGCNN and 5.34% for MOFormer.The reported reductions are attributed to the self-supervised pretraining framework.
- QMOF band-gap prediction: MOFormer and CGCNN predictions are closer to DFT values than SOAP and Stoichiometric-120, especially for band gaps ≤2eV.SOAP and Stoichiometric-120 more often overpredict lower band gaps, whereas MOFormer and CGCNN better estimate the two lowest-band-gap MOFs.
- hMOF gas adsorption: For hMOF gas adsorption, pretrained CGCNN achieves the strongest overall performance, with 48% lower MAE than the compared alternatives.Pretraining improves accuracy by 4.3% for MOFormer and 16.5% for CGCNN across gas-adsorption predictions.
- Learned representations: CGCNN representations cluster high-CO2-adsorption MOFs more closely, while MOFormer more strongly groups MOFs sharing the same topology.MOFormer attention emphasizes key atoms such as Y, O, and C together with topology encodings, including strong attention to Y and pcu.
- Data efficiency: MOFormer outperforms CGCNN for QMOF band-gap prediction when the training set size is ≤1000, whereas CGCNN consistently performs better for CO2 adsorption.Both models outperform SOAP on QMOF, and pretraining consistently improves their data efficiency.
Conclusion
MOFormer is a structure-agnostic Transformer that uses MOFid strings for rapid MOF property prediction and is jointly pretrained with CGCNN using unlabeled MOF data. It achieves strong accuracy across band-gap and gas-adsorption tasks, with data-efficiency benefits and interpretable attention to atoms and topology.
- Model and purpose: MOFormer predicts MOF properties from MOFid strings, enabling rapid inference without requiring optimized 3D structures.The model is intended to explore the vast chemical space of hypothetical MOFs using arbitrarily constructed MOFid inputs.
- Model and purpose: A self-supervised framework jointly pretrains MOFormer and CGCNN on a large unlabeled MOF dataset to improve downstream prediction accuracy.The framework connects structure-agnostic and structure-based representations during pretraining.
- Prediction results: MOFormer achieves 21.4% higher band-gap accuracy and 35%-48% higher accuracy on various gas-adsorption tasks than Stoichiometric-120.These comparisons are against another structure-agnostic method.
- Prediction results: Pretraining improves average accuracy by 5.34% and 4.3% for MOFormer and by 6.79% and 16.5% for CGCNN across band-gap and gas-adsorption prediction, respectively.The reported improvements are ordered as band gap and gas adsorption.
- Interpretation: MOFormer and CGCNN are less likely than SOAP and Stoichiometric-120 to overpredict band gaps, supporting conductive-MOF prescreening for energy applications.For gas adsorption, MOFormer relies more on topology information than CGCNN because topology is strongly correlated with the label.
Supporting Information Available
The supporting information covers model architecture, self-supervised pretraining, dataset distributions, SOAP feature construction, and band-gap prediction analysis.
- Contents: The supporting information details the Transformer and self-attention mechanism, along with the CGCNN and MOFormer models.
- Contents: It describes the self-supervised pretraining procedure and the distributions of the QMOF and hMOF datasets.
- Contents: It reports SOAP feature-vector parameters, the effect of SOAP vector length on model accuracy, and kernel-density estimates of band-gap predictions.
Self-attention and the Transformer model
The Transformer encoder learns input-sequence representations through self-attention rather than recurrent architecture. MOFormer uses this encoder, with multihead attention enabling parallel processing and representations from multiple subspaces.
- MOFormer uses the encoder part of the Transformer model as its base.
- Transformers replace recurrent architectures with self-attention to learn representations of input sequences.Self-attention also supports learning long-range dependencies.
- Self-attention is parallelizable through matrix multiplication, allowing faster training than sequential processing.
- Each token is linearly projected into query, key, and value vectors before scaled dot-product attention is calculated.The projection matrices are Wq, Wk, and Wv.
- The Transformer encoder contains 6 self-attention layers.
Training Details
The training framework combines MOFormer, CGCNN, and self-supervised pretraining, then attaches regression heads for downstream prediction. Pretraining aligns the two models’ representations and improves their subsequent performance.
- The framework consists of MOFormer, CGCNN, and self-supervised pretraining.
- During finetuning, a multilayer perceptron regression head is attached to either encoder for prediction.The regression head has 4 layers with 512, 256, 128, and 64 neurons.
- MOFormer finetuning uses different learning rates for the Transformer encoder and the MLP regression head.
- MOFormerScratch is trained to measure performance improvement after self-supervised pretraining.
- CGCNN uses crystal structure files as input, whereas MOFormer uses a structure-agnostic representation that can lose atom-neighborhood and crystal-geometry information.
- Pretrained CGCNN performs better than CGCNNScratch, the model trained without pretraining.
- Self-supervised pretraining aligns MOFormer and CGCNN representations so each model can leverage embeddings learned by the other.Performance enhancements for both models are reported during finetuning.
Distribution of the benchmark datasets
The QMOF and hMOF subsets with available MOFid have distributions similar to their full datasets. The authors therefore treat benchmark comparisons involving these subsets as fair.
- Only part of the QMOF and hMOF datasets has an available MOFid.
- The QMOF full dataset and MOFid subset are compared using their value distributions and reported means and standard deviations.
- 0.135 is the reported difference between the average band-gap values of the full QMOF dataset and its MOFid subset.The passage also states that their distributions are similar.
- The average CO2 and CH4 adsorption values differ very little between the full hMOF dataset and its MOFid subset, whose distributions are also similar.
- The authors regard comparisons between MOFormer and other models on QMOF and hMOF as fair.
Parameters to create SOAP feature vector
SOAP feature-vector size and usefulness depend strongly on the number and variation of elements represented in the dataset. Its performance is high on hMOF but declines when the elemental space is expanded, limiting its use for diverse MOF exploration.
- The SOAP feature vector is created with the DScribe package, using parameters recorded in Table S4.
- SOAP vector length varies with the number of different atom elements, becoming longer and sparser for datasets with broader elemental diversity.
- The hMOF SOAP vector has length 2772, whereas the QMOF vector using 79 elements has a much larger length.
- High SOAP accuracy on hMOF is attributed to the dataset’s small variation in element type.
- Doubling the original element set for hMOF substantially reduces XGBoost performance using the resulting SOAP feature.
- SOAP is not considered suitable for exploring MOF chemical spaces with diverse elements.
Comparison between models’ performance in QMOF
Figure S4 compares predicted MOF band-gap distributions from different models against DFT ground truth on a randomly selected QMOF test set. SOAP and Stoichiometric-120 with XGBoost underpredict low-band-gap values more frequently than CGCNN and MOFormer.
- Figure S4 uses KDE plots to compare predicted MOF band-gap distributions from different models on a randomly selected QMOF test set.
- The DFT-calculated band gap is the ground-truth reference, shown by the blue curve.
- SOAP and Stoichiometric-120 features with XGBoost underpredict low-band-gap values more frequently than CGCNN and MOFormer.
- The figure includes a zoomed-in view of MOFs with band gaps ≤1eV.