arXiv:2605.25866v1 [cs.LG] 25 May 2026
UNATE: UNsupervised ATomic Embedding for crystal structures property prediction Laura Solà-Garcia
Àlex Solé
Javier Ruiz-Hidalgo
Signal Theory and Comms. Dept. UPC-Univ. Politècnica de Catalunya Barcelona, Spain 0009-0003-2078-093X
Signal Theory and Comms. Dept. UPC-Univ. Politècnica de Catalunya Barcelona, Spain 0000-0002-2071-9317
Signal Theory and Comms. Dept. UPC-Univ. Politècnica de Catalunya Barcelona, Spain 0000-0001-6774-685X
Abstract—Accurately predicting crystal properties is critical for accelerating materials discovery, but it is often limited by scarce labeled data and costly theoretical calculations. To alleviate this, we propose UNATE (Unsupervised Atomic Embedding), a framework that leverages structural information extracted from unlabeled crystal structures. UNATE integrates an unsupervised denoising autoencoder with self-supervised contrastive learning to learn robust atomic representations, which are then used as input features for downstream property prediction. Experimental results show that replacing raw atomic numbers with UNATE-pretrained node embeddings yields a 2.7% improvement over the full-data baseline. Notably, the benefits become more pronounced in scenarios with limited labeled data, reaching improvements of up to 10% when only 25% of the labeled data is used. The code is available at https://github.com/LauraSola/ Unsupervised-Atom-Embedding-Gen. Index Terms—Unsupervised Learning, Graph Neural Networks, Embeddings, Pretraining, Multitasking, Materials Science.
I. I NTRODUCTION Predicting the properties of crystalline materials is crucial for the discovery of new materials, with applications in electronics, energy storage, and renewable energy [4], [10], [13]. Properties such as band gap, formation energy, and elastic constants determine a material’s functionality, but experimental characterization is often slow, costly, and inefficient. To overcome these limitations, computational models based on theoretical estimations [12] have become widely used, though they remain expensive at scale due to the complexity of atomic structures [5]. In recent years, machine learning has emerged as a promising alternative, achieving near-theoretical accuracy while enabling large-scale screening [1], [8], [18]. A major challenge in developing effective machine learning models is the limited availability of high-quality labeled data, which requires expensive theoretical calculations or experimental measurements. In contrast, large public repositories contain crystal structures without associated property labels, motivating the use of techniques that exploit unlabeled data to learn meaningful structural representations [2], [9]. The contributions for this work are: • A proposed dual-branch self-supervised framework, UNATE, that learns transferable atomic embeddings from unlabeled crystal graphs, and transfer them by replacing
raw atomic numbers without modifying downstream architectures. • We demonstrate consistent improvements in supervised band gap prediction, including improved label efficiency (up to 10% MAE improvement with 25% labeled data, and +2.7% with 100% labeled data training). • We provide qualitative evidence via t-SNE visualizations that UNATE embeddings recover meaningful chemicalspace structure directly from data, without encoding explicit chemistry theory. II. R ELATED W ORK Graph Neural Networks (GNNs) are widely used for molecular and crystal property prediction due to their ability to model relational structures. Early models such as CGCNN [18] represented crystals as periodic graphs with periodic boundary conditions (PBC). More recent architectures improve geometric reasoning and long-range interaction modeling, including Matformer [20], which integrates lattice information into a transformer framework; PotNet [7], which captures local and global interactions using physically motivated potentials; iComformer/eComformer [19], which enforces rotation invariance through refined unit-cell representations or rotation equivariance via tensor products; CartNet [15], which achieves soft rotation-equivariance via rotational data augmentation, improving the speed and accuracy; and PRISM [14], which explicitly integrates multiscale representations and periodic feature encoding through a mixture of expert modules. Because labeled crystal datasets are scarce due to the high cost of experiments and simulations, recent work has focused on pretraining GNNs on unlabeled crystal structures to learn transferable representations. DualSSL [3] proposes a multi-task self-supervised framework combining atom masking, contrastive learning between perturbed views, and microproperty prediction (like atomic stiffness) to capture both local and global structure, followed by supervised fine tuning of the pretrained backbone. CrysGNN [2] trains a teacher model using both self-supervised (like node reconstruction) and supervised objectives (like space-group classification), and transfers knowledge to downstream models via distillation, enabling architecture-independent reuse of learned representations by aligning intermediate features between the teacher
and student models. Finally, CrysAtom [9] learns transferable atom-level embeddings using an autoencoder with a CGCNN backbone that reconstructs node and edge attributes and applies a Barlow Twins loss [21] to promote invariant and non-redundant representations. These embeddings are then reused in downstream models as input node features. Previous pretraining approaches are often evaluated on relatively simple backbones, where gains are easier to obtain but tend to diminish as model capacity increases. In contrast, our method yields consistent improvements even with high-performing backbones. In addition, by transferring knowledge through pretrained embeddings rather than model weights, the approach remains model- and task-agnostic, enabling broad applicability across downstream architectures and property prediction tasks. III. M ETHODOLOGY This work aims to enhance crystal property prediction by learning transferable representations from unlabeled crystal graphs that can be used to improve downstream supervised tasks. We propose a dual-branch pretraining framework combining an unsupervised denoising autoencoder and a selfsupervised contrastive learning objective. The two branches are trained in parallel, encouraging the encoder to capture both local structural information and global semantic consistency, as illustrated for the unsupervised branch in Fig. 1. Graph representation and augmentations
dropped edges, and outputs node embeddings hi , which are then passed to the two decoders.
Fig. 1. Unsupervised learning branch. CartNet encodes a masked, edgedropped graph. The decoders reconstruct atomic numbers and the periodicaware adjacency matrix.
1) Node Feature Reconstruction (decoder 1): The node embeddings hi ∈ Rd are first L2-normalized and then passed through a linear layer that projects each embedding vector into a 118-dimensional space, corresponding to the number of atomic element classes. The resulting logits are then passed through a softmax function to obtain the predicted probability distribution ŷi : ŷi = softmax Wnode ·
hi + bnode , ∥hi ∥2
ŷi ∈ RNZ , (1)
Each crystal is represented as a periodic graph G = (V, E), where nodes i ∈ V correspond to atoms and edges (i, j) ∈ E denote atomic interactions under periodic boundary conditions. Node features are atomic numbers Zi , and edge features include interatomic distances and direction vectors. These features are then mapped to high-dimensional embeddings, hi ∈ Rd for nodes and hij ∈ Rde for edges, encoding chemical and structural information for GNN layers to propagate and aggregate information across the graph. Lastly, to improve embedding robustness, we apply two augmentations to the graphs: node masking, which randomly hides a subset of node features to encourage inference from context, and edge dropping, which randomly removes edges to promote learning from remaining interactions.
for each node i ∈ V , where Wnode ∈ RNZ ×d and bnode ∈ R are learnable parameters, and NZ denotes the number of possible atomic numbers. The model is trained by minimizing the negative log-likelihood loss Lnode , aligning the predicted atomic number distribution with the ground-truth labels. 2) Adjacency Matrix Reconstruction (decoder 2): This decoder reconstructs the graph connectivity by predicting the number of edges originally present between each node pair after random edge removal. Due to Periodic Boundary Conditions (PBC), multiple edges can exist between the same pair of atoms. Edge reconstruction is therefore formulated as a six-class classification problem, corresponding to edge multiplicities: 0, 1, 2, 3, 4, or 5+. For each node pair (i, j), the decoder first computes a bilinear interaction vector from the normalized node embeddings:
A. Unsupervised Learning Branch: Denoising Autoencoder
sij = h⊤ i Wb hj + bb ,
The unsupervised branch follows a denoising autoencoder architecture. The encoder receives an augmented graph (as described in the section III) and reconstructs the original graph, learning robust embeddings hi that capture local chemical environments and structural context. This denoising objective encourages inference from context rather than memorization, improving downstream performance when combined with supervised fine tuning (Section IV-B). For the encoder, we adopt CartNet [15] due to its strong and lightweight performance relative to other crystal graph architectures. The CartNet encoder consists of an atom encoder, an edge encoder, and several CartLayers. It processes an augmented graph G = (V, E), with masked nodes and
where Wb ∈ Rd×6×d and bb ∈ R6 are the weights and bias of the bilinear layer. This interaction vector sij is then projected through a linear layer to produce logits, which are normalized with a softmax to obtain edge probabilities (ŷij ):
NZ
ŷij = softmax(Wa sij + ba ),
(2)
(3)
where Wa ∈ R6×6 and ba ∈ R6 are the trainable parameters of the linear layer. The model is trained using a weighted cross-entropy loss, denoted as Ladj , assigning lower weight to the zero-edge class to mitigate class imbalance caused by the large number of non-connected node pairs.
B. Self-supervised Learning Branch: Contrastive Objective In parallel to the denoising autoencoder, we use a contrastive self-supervised branch to encourage globally consistent and semantically structured embeddings. Instead of reconstructing features, this branch drives the encoder to organize the embedding space so that similar crystal structures are embedded closer together, improving generalization and robustness. Graph-level embeddings are obtained by mean-pooling node embeddings produced by the encoder and passing them through a two-layer MLP with SiLU activation:
z = W2 · SiLU W1
1 X hi |V |
!
! + b1
+ b2 .
(4)
i∈V
For each graph, we generate two random augmented views via node masking and edge dropping. Let (zi , zj ) be a positive pair (embeddings originating from the same graph), and all other embeddings in the batch act as negatives. For each anchor embedding zi , the contrastive InfoNCE loss [11] encourages similarity to its positive counterpart zj while discouraging similarity with all negatives: (i) LInfoNCE = − log
Softmaxk̸=i
! 1 , · sim(zi , zk ) τ j
(5)
with temperature τ . The final loss, denoted as LInfoNCE , is the average over all 2N augmented embeddings. This objective enforces invariance to graph perturbations while preserving discriminative structure, complementing the reconstruction branch and enhancing learned representations. C. Total Pretraining Objective Fig. 2 summarizes the overall training setup and how the loss terms in Eq. (6) are obtained. The final objective combines the reconstruction and contrastive losses: Ltotal = α Lnode + β Ladj + γ LInfoNCE
(6)
where α, β, γ control the contribution of each loss term.
Fig. 2. For each crystal graph, we generate multiple augmented views by randomly masking nodes and edges. A shared encoder produces node embeddings that a decoder uses to reconstruct the masked atomic features and adjacency relations (denoising objective). In parallel, graph-level embeddings from each augmented view are optimized with a contrastive loss that pulls together embeddings of the same crystal and pushes apart those from different crystals.
adapts the pretrained embeddings before they are processed by the message-passing layers. Apart from this modification, the downstream architecture is trained from scratch, without reusing pretrained weights. IV. R ESULTS We first describe the datasets and computational resources used for evaluation, followed by a quantitative analysis of the proposed pretraining framework. All results are reported as the average over four independent runs to account for the variability inherent to graph neural network training. Datasets
D. Knowledge Transfer via Pretrained Atomic Embeddings Knowledge learned during pretraining is transferred to downstream crystal property prediction tasks through pretrained atomic embeddings. Specifically, raw atomic numbers are replaced by fixed, element-wise embeddings learned from unlabeled crystal structures, which then serve as node features for the downstream model. To construct these embeddings, the pretrained encoder is frozen and applied to the training dataset to extract node representations. The embeddings corresponding to each atomic number are aggregated by computing their mean across all occurrences, resulting in a single, stable representation per element. The final embeddings are L2-normalized to ensure consistent scaling. In downstream models, the Atom Encoder (component of CartNet [15]) is replaced by a lightweight linear projection that
We use two datasets from the Materials Project [1]. The unlabeled dataset contains ∼139,000 crystal structures represented as crystal graphs, used for the generation of the embeddings. The labeled dataset (2018.6.1 release) comprises 69,239 inorganic crystals with various property labels. Band gap was selected as the target property because it is available for a much larger number of materials than bulk or shear modulus, which are limited to ∼5,500 crystals, and has been commonly used to evaluate representation learning methods in prior work [9]. Training details All experiments were executed on a high-performance computing cluster using a single GPU per run, with 11 to 48 GB of GPU memory and up to 6 CPU cores allocated per experiment, depending on the computational requirements.
CartNet was employed as the backbone. During training for embedding generation, the Adam optimizer [6] was used with a learning rate of 3·10−2 and batch size of 128 for 100 epochs. The parameters for the final loss (Eq. 6) were set to α = 225, β = 4, and γ = 3. Hyperparameter tuning was performed using grid search. A. Results and Impact of Labeled Dataset Size From Table I, we compare supervised band gap prediction performance using 128-dim embeddings across different pretraining strategies using CartNet as the backbone. The results show that UNATE embeddings improve the MAE by 2.7% with respect to the baseline, whereas CrysAtom embeddings do not yield gains over the baseline. TABLE I S UPERVISED BAND GAP PREDICTION (128- DIM EMBEDDINGS ). MAE (± STD . DEV.; ×10−3 (I MPROV. %) Method
Bandgap
CartNet w/ CrysAtom [9] w/ UNATE
194.02 ± 3.05 195.89 ± 2.61 (-0.6%) 188.78 ± 2.19 (+2.7%)
As this approach is aimed at reducing the dependence on labeled data, we progressively reduced the size of the dataset and measured the corresponding improvement provided by pretraining. Starting from the full set of ∼69K crystals, we also considered 50% and 25% subsets, while keeping the pretraining fixed on the full unlabeled dataset. TABLE II I MPACT OF LABELED DATASET SIZE ON MAE (± STD . DEV.; ×10−3 ) FOR MULTIPLE EMBEDDING DIMENSIONS USING C ART N ET AS BACKBONE . Dim
% Labeled Data
Baseline
Pretrained (Improv. %)
64
100% 50% 25%
202.47 ± 2.23 241.60 ± 2.02 291.67 ± 4.65
197.77 ± 1.51 (+2.3%) 235.22 ± 3.23 (+2.6%) 274.58 ± 5.97 (+5.9%)
128
100% 50% 25%
194.02 ± 2.62 240.55 ± 8.19 289.61 ± 10.24
188.78 ± 2.19 (+2.7%) 220.91 ± 4.27 (+8.2%) 260.63 ± 7.00 (+10.0%)
As shown in Table II, the benefit of pretraining grows as labeled data becomes scarce. With 64-dimensional embeddings, the relative improvement increases from 2.3% (full data) to 5.9% with only 25% of the labeled samples. For 128dimensional embeddings, the gains rise from 2.7% (full data) to 10% at 25% of labeled data. These results highlight the effectiveness of pretraining, particularly in low-data regimes, where pretrained node embeddings provide chemical and structural information beyond atomic numbers. B. Ablation Study Table II shows that 128-dimensional embeddings outperform 64-dimensional embeddings for both the baseline and pretrained models (larger embedding sizes could not be further explored due to GPU memory limitations). Pretraining consistently improves performance across both embedding sizes.
Although the main experiments use 128-dimensional embeddings, the following ablation studies are conducted with 64dimensional embeddings to reduce computational cost. Despite the lower dimensionality, the 64-dimensional setting preserves the same overall performance trends and is therefore sufficient for analyzing the contribution of individual components of the framework. Using the 64-dimensional embeddings, we analyzed the contribution of individual components of the proposed framework. Three aspects were considered: the effect of noise injection in the denoising autoencoder, the choice of selfsupervised objective, and the strategy used to transfer pretrained knowledge. Results are summarized in Table III. TABLE III A BLATION STUDY COVERING NOISE INJECTION , SELF - SUPERVISED METHOD , AND KNOWLEDGE - TRANSFER STRATEGY USING C ART N ET AS BACKBONE . MAE (± STD . DEV.; ×10−3 ) . Experiment / Setting CartNet (baseline - no pretraining)
MAE 202.47 ± 2.30
1. Autoencoder Noise No Noise Noise Added (UNATE)
201.77 ± 2.65 (+0.35%) 197.77 ± 1.51 (+2.32%)
2. Self-Supervised Method Barlow Twins [21] Contrastive Learning (UNATE) Deep Graph Infomax [17]
205.25 ± 2.61 (-1.37%) 197.77 ± 1.51 (+2.32%) 199.21 ± 2.18 (+1.61%)
3. Knowledge Transfer Strategy Weight Initialization (transfer learning) Node Embeddings (UNATE) Combined (weight init. + node embd.)
417.89 ± 4.01 (-106.4%) 197.77 ± 1.51 (+2.32%) 511.89 ± 5.67 (-152.9%)
From these ablation studies, three main observations emerge. Adding noise via node masking and edge dropping during pretraining consistently improves downstream performance, confirming the regularizing effect of the denoising autoencoder. Among the self-supervised objectives, contrastive learning yields the best results, suggesting that distinguishing between different structural contexts while maintaining invariance across augmented views leads to more informative representations. Finally, transferring pretrained atomic embeddings enhances performance compared to the baseline, whereas direct weight initialization from the pretrained encoder severely degrades performance, likely due to a mismatch with the downstream prediction task. C. Visualizing the Embeddings Finally, to assess the quality and interpretability of the learned representations, we visualize the atomic embeddings using t-SNE [16]. Figure 3 shows a t-SNE projection of the pretrained node embeddings (averaged by atomic number), later used as input for the downstream prediction task. Clear clusters emerge according to chemical families: transition metals (green) group in the upper left, lanthanides (gray) in the lower center, and reactive non-metals toward the upper right, etc. This demonstrates that the embeddings capture
R EFERENCES
Fig. 3. t-SNE visualization of the 128-dimensional atomic embeddings (first and second components). Each point corresponds to an element (mean embedding by atomic number), colored by its chemical category.
chemically meaningful information, reflecting both structural and elemental properties, which leads to improved downstream property prediction. V. C ONCLUSIONS We have presented a dual-branch pretraining framework for crystal property prediction, combining a denoising autoencoder and a contrastive self-supervised branch. Using pretrained node embeddings as input features consistently improved bandgap prediction performance. Specifically, improvements of 2.3% and 2.7% were observed for 64- and 128-dimensional embeddings, respectively, with larger gains (up to 10%) in low-data regimes, highlighting the benefits of pretraining when labeled data is limited. Our analysis shows that contrastive learning outperforms other self-supervised approaches such as Barlow Twins and Deep Graph Infomax. Furthermore, t-SNE projections of the learned embeddings reveal meaningful chemical structure, with atoms from the same group in the periodic table clustering together, despite training being fully unsupervised. These results demonstrate that pretraining on large unlabeled crystal datasets is an effective strategy to reduce the computational cost of creating labeled data and enhance downstream property prediction. Future work could focus on evaluating the proposed framework on additional material properties, such as formation energy and bulk modulus, to further assess its generalizability beyond bandgap prediction. In addition, exploring alternative self-supervised learning strategies, including triplet loss and contrastive learning with hard negative sampling, as well as different graph augmentation techniques, such as edge rewiring and node feature shuffling, may lead to richer and more informative node representations. ACKNOWLEDGMENT This work has been supported by the research project PID2024-161868OB-I00 [C3DRUM] funded by MCIN/AEI/10.13039/501100011033 and FEDER, EU.
[1] Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong. Graph networks as a universal machine learning framework for molecules and crystals. Chemistry of Materials, 31(9):3564–3572, 2019. [2] Kishalay Das, Bidisha Samanta, Pawan Goyal, Seung-Cheol Lee, Satadeep Bhattacharjee, and Niloy Ganguly. Crysgnn: Distilling pretrained knowledge to enhance property prediction for crystalline materials. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 7323–7331, 2023. [3] Nihang Fu, Lai Wei, and Jianjun Hu. Physics-guided dual selfsupervised learning for structure-based material property prediction. The Journal of Physical Chemistry Letters, 15(10):2841–2850, 2024. [4] Martin A Green, Anita Ho-Baillie, and Henry J Snaith. The emergence of perovskite solar cells. Nature photonics, 8(7):506–514, 2014. [5] A. Hautier, G.and Jain and S. P. Ong. From the computer to the laboratory: materials discovery and design using first-principles calculations. Journal of Materials Science, 47:7317–7340, 2012. [6] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [7] Yuchao Lin, Keqiang Yan, Youzhi Luo, Yi Liu, Xiaoning Qian, and Shuiwang Ji. Efficient approximations of complete interatomic potentials for crystal property prediction. In International conference on machine learning, pages 21260–21287. PMLR, 2023. [8] Steph-Yves Louis, Yong Zhao, Alireza Nasiri, Xiran Wang, Yuqi Song, Fei Liu, and Jianjun Hu. Graph convolutional neural networks with global attention for improved materials property prediction. Physical Chemistry Chemical Physics, 22(32):18141–18148, 2020. [9] Shrimon Mukherjee, Madhusudan Ghosh, and Partha Basuchowdhuri. Crysatom: Distributed representation of atoms for crystal property prediction. In Proceedings of the Third Learning on Graphs Conference (LoG 2024), volume 269 of Proceedings of Machine Learning Research, Virtual Event, November 26–29 2024. PMLR. [10] Jens Kehlet Nørskov, Thomas Bligaard, Jan Rossmeisl, and Claus Hviid Christensen. Towards the computational design of solid catalysts. Nature chemistry, 1(1):37–46, 2009. [11] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv:1807.03748, 2018. [12] John P Perdew and Karla Schmidt. Jacob’s ladder of density functional approximations for the exchange-correlation energy. In AIP Conference Proceedings, volume 577, pages 1–20, 2001. [13] Rampi Ramprasad, Rohit Batra, Ghanshyam Pilania, Arun MannodiKanakkithodi, and Chiho Kim. Machine learning in materials informatics: recent applications and prospects. npj Computational Materials, 3(1):54, 2017. [14] Àlex Solé, Albert Mosella-Montoro, Joan Cardona, Daniel Aravena, Silvia Gómez-Coca, Eliseo Ruiz, and Javier Ruiz-Hidalgo. PRISM: Periodic representation with multiscale and similarity graph modelling for enhanced crystal structure property prediction. npj Computational Materials, 2026. [15] Àlex Solé, Albert Mosella-Montoro, Joan Cardona, Silvia Gómez-Coca, Daniel Aravena, Eliseo Ruiz, and Javier Ruiz-Hidalgo. A cartesian encoding graph neural network for crystal structure property prediction: application to thermal ellipsoid estimation. Digital Discovery, 4:694– 710, 2025. [16] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9(86):2579–2605, 2008. [17] P. Velickovic, W. Fedus, W. L. Hamilton, P. Liò, Y. Bengio, and R. D. Hjelm. Deep graph infomax. ICLR (poster), 2(3):4, 2019. [18] Tian Xie and Jeffrey C Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical review letters, 120(14):145301, 2018. [19] K. Yan, C. Fu, X. Qian, X. Qian, and S. Ji. Complete and efficient graph transformers for crystal material property prediction. In International Conference on Learning Representations (ICLR), 2024. [20] Keqiang Yan, Yi Liu, Yuchao Lin, and Shuiwang Ji. Periodic graph transformers for crystal material property prediction. Advances in Neural Information Processing Systems, 35:15066–15080, 2022. [21] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 12310–12320. PMLR, 2021.