SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Namkyeong Lee $^{1\dagger}$ , Heewoong Noh $^{1\dagger}$ , Sungwon Kim $^{1}$ , Dongmin Hyun $^{2}$ , Gyoung S. Na $^{3}$ , Chanyoung Park $^{1*}$
$^{1}$  KAIST  $^{2}$  Yahoo Research  $^{3}$  KRIXT
{namkyeong96,heewoongnoh,swkim,cy.park}@kaist.ac.kr
dhyun@yahooinc.com, ngs0@krict.re.kr
The density of states (DOS) is a spectral property of crystalline materials, which provides fundamental insights into various characteristics of the materials. While previous works mainly focus on obtaining high-quality representations of crystalline materials for DOS prediction, we focus on predicting the DOS from the obtained representations by reflecting the nature of DOS: DOS determines the general distribution of states as a function of energy. That is, DOS is not solely determined by the crystalline material but also by the energy levels, which has been neglected in previous works. In this paper, we propose to integrate heterogeneous information obtained from the crystalline materials and the energies via a multi-modal transformer, thereby modeling the complex relationships between the atoms in the crystalline materials and various energy levels for DOS prediction. Moreover, we propose to utilize prompts to guide the model to learn the crystal structural system-specific interactions between crystalline materials and energies. Extensive experiments on two types of DOS, i.e., Phonon DOS and Electron DOS, with various real-world scenarios demonstrate the superiority of DOSTransformer. The source code for DOSTransformer is available at https://github.com/HeewoongNoh/DOSTransformer.
Along with recent advances in machine learning (ML) for scientific discovery, ML models have rapidly been adopted in computational materials science thanks to the abundant experimental and computational data in the field. Most ML models developed in computational materials science to date have been focused on crystalline materials property consisting of single-valued quantities [29], e.g., band gap energy [31, 69], formation energy [3, 56], and Fermi energy [61].
On the other hand, spectral properties are ubiquitous in materials science, characterizing various properties of crystalline materials, e.g., X-ray absorption, dielectric function, and electronic density of states [29] (Figure 1).
The density of states (DOS), which is the main focus of this paper, is a spectral property that provides fundamental insights on various characteristics of crystalline materials, even enabling direct computation of single-valued properties [17]. For example, DOS is utilized as a feature of materials for analyzing the underlying reasons for changes
Crystalline Material
SbPO4
Band Gap 3.34 eV Formation Energy -2.299 e Energy Above Hull 0.057 eV
Single-valued Property
Spectral Property
Figure 1: A crystal structure and various types of its properties.
*Corresponding author. † These authors contributed equally.
in electrical conductivity [15]. Moreover, band gaps and edge positions, which can be directly derived from DOS, are utilized to discover new photoanodes for solar fuel generation [47, 64]. As a result, the exploration of ML capabilities for predicting DOS takes us one step further toward the core principles of materials science, thereby expediting the process of materials discovery. However, credible computation of DOS with traditional density functional theory (DFT) requires expensive time/financial costs of exhaustively conducting experiments with expertise knowledge [7, 14]. Therefore, alternative approaches for DOS calculation are necessary, whereas demonstrating ML capabilities for learning such spectral properties of the crystalline materials is relatively under-explored.
In this paper, we focus on predicting DOS by following the nature of DOS calculation: DOS determines the general distribution of states as a function of energy. That is, while single-valued properties can be described solely by the crystalline material, DOS is determined by not only the material itself but also the energy levels at which DOS is calculated. Therefore, integrating the heterogeneous signals from both the crystalline material and the various energy levels is crucial for DOS prediction. However, previous works for DOS prediction focus on obtaining a qualified representation of material [7, 14, 17, 9], thereby overlooking the unique characteristics of DOS.
On the other hand, learning from multiple input types, i.e., crystalline material and various energy levels, is not trivial due to their heterogeneity. To this end, we formulate the DOS prediction problem into a multi-modal learning problem, which has recently gained significant attention from ML researchers in various domains thanks to its capability of extracting and relating information from heterogeneous data types [6, 38, 54, 5, 2]. For example, VisualBERT proposes to align elements of input modalities, i.e., images and text, thereby achieving SOTA performance in various vision-and-language tasks [36]. Moreover, Dall-E shows its power in generating images based on text data by encoding images and texts together [43]. Among various multi-modal learners, the multi-modal transformer demonstrates its intrinsic advantages and scalability in modeling multiple modalities [62], by associating heterogeneous modalities with the cross-attention mechanism [65, 50].
Inspired by the recent success of the multi-modal transformer, we propose a multi-modal transformer model for DOS prediction, named DOSTransformer, which incorporates crystalline material and energies as heterogeneous input modalities. Unlike previous works that overlook the energy levels for predicting DOS [7, 14, 17, 9], DOSTransformer learns embeddings of energy levels which are used for modeling complex relationships between the energy levels and crystalline material through a cross-attention mechanism. By doing so, DOSTransformer obtains multiple representations for a single crystalline material according to various energy levels, enabling the prediction of a single DOS value on each energy level. However, simply combining a crystalline material with energy levels may fail to consider the significant impact of the crystal's structure, which sometimes results in distinct material properties even when it is composed of identical elements (e.g., carbon). For example, although graphite and diamond are made entirely out of carbon, their arrangement (or structure) of carbon atoms differs significantly, resulting in distinct material properties. Therefore, we utilize learnable prompts to inform DOSTransformer of the input crystal system (among the seven crystal systems $^2$ ), which navigates the model to capture structure-specific interactions between a crystalline material and energy levels. By doing so, DOSTransformer learns to extract not only the relational information that is shared with a crystal system, but also across the crystal systems. In this work, we make the following contributions:
- Inspired by the fact that DOS is determined by not only the materials themselves but also the energy levels at which DOS is calculated, we propose a multi-modal transformer model for DOS prediction, called DOSTransformer, which incorporates crystalline materials and energy as heterogeneous input modalities.
- To capture structure-specific interactions between a crystalline material and energy levels, thereby being able to predict DOS that significantly differs according to the crystal's structure, we utilize learnable prompts to inform DOSTransformer of the input crystal system.
- Our extensive experiments on two types of DOS, i.e., Phonon DOS and Electron DOS, demonstrate that DOSTransformer outperforms a wide range of state-of-the-art methods in various real-world scenarios, i.e., in-distribution scenarios and out-of-distribution scenarios.
To the best of our knowledge, this is the first work that considers various energy levels during DOS prediction and introduces prompts for crystal structural systems.
$^{2}$ https://en.wikipedia.org/wiki/Crystal_system
2.1 Machine Learning for Materials and Density of States
Traditional materials science research heavily relies on theory, experimentation, and computer simulations, but these approaches are time-consuming, costly, and inefficient, hindering their ability to keep up with the field's development [57]. As an alternative research tool, ML, which requires neither expensive trial and error experiments nor domain expert knowledge, has been attracting a surge of interest from materials scientists [67]. Consequently, various ML methods have been established to learn the relationship between materials and their properties which are usually obtained from ab-initio calculations [11]. Inspired by the recent success of GNNs in biochemistry [19, 48, 25, 33, 34, 32], CGCNN proposes a message-passing framework for crystal structure providing high-accuracy prediction for 8 different materials properties [61]. However, most studies have focused on single-valued properties [29], whereas various spectral properties, such as DOS, are also crucial.
The DOS describes the number of different states at a particular energy level that particles are allowed to occupy, determining various properties of crystalline materials. For example, phonon DOS has a crucial influence in determining the specific heat and vibrational entropy of crystalline materials as well as the interfacial thermal resistance [60], while electron DOS is used to derive the electronic contribution to heat capacity in metals [21] and the effective mass of electrons in charge carriers [18].
Despite the importance of DOS, how to incorporate ML capabilities for decoding such spectral properties is under-explored since the vast majority of methods concentrate only on encoding crystalline materials. Specifically, Chandrasekaran et al. [7] proposes to predict electron DOS based on the hand-crafted fingerprint of each grid point in the crystalline materials, while Del Rio et al. [14] leverages that of each atom in the crystalline materials. Chen et al. [9] predicts phonon DOS by learning the representations of the crystalline materials that are equivariant to 3D rotations, translations, and inversion with Euclidean neural networks [49, 28, 58], achieving high-quality prediction with a small number of training materials with over 64 atom types. In this work, we concentrate on decoding spectral properties regarding the complex relationship between the various energy levels and crystalline materials.
2.2 Multi-modal Transformer, Prompt Tuning, and Positional Encoding
Multi-modal Transformer. In recent years, much progress in multi-modal learning has grown rapidly by extracting and aligning the rich information from heterogeneous modalities [2, 6]. Among the various multi-modal learning methods, the multi-modal transformer demonstrates its superior capability in learning multiple modalities [62], by associating the heterogeneous modalities with a cross-attention mechanism. Specifically, MulT [50] learns the representations directly from unaligned multi-modal streams with multi-modal transformer, while Yao & Wan [65] improves machine translation quality by learning the image-aware text representations. In this paper, we investigate how to decode DOS of a crystalline material by learning the energy-aware crystal representations via a multi-modal transformer.
Prompt Tuning. On the other hand, prompt designing has gained significant attention from NLP researchers as an efficient method for fine-tuning large language models (LLMs) for various tasks [39]. There are two main approaches to prompt designing: discrete prompt designing [44] and continuous prompt tuning [37]. While discrete prompt designing involves manually creating task description tokens for LLMs through trial and error, continuous prompt tuning trains learnable prompts in a continuous latent space without requiring expert knowledge. Inspired by the success of continuous prompt tuning in NLP domain, it has been widely adopted in various domains such as computer vision [24] and vision-language [68]. While previous works utilize prompts for efficient fine-tuning of the models, we employ prompts to capture the heterogeneity of various crystal structures.
Positional Encoding. In the original proposal of the Transformer architecture, a sinusoidal function was suggested as a method for positional encoding. However, the sinusoidal function may have limitations in terms of learnability and flexibility, which can affect its effectiveness [53]. To address this issue, most pre-trained language models [16, 40] employ learnable vector embeddings as positional representations. By conceptualizing each energy value as the position of a word in a sentence, various energy embeddings in Section 4.2 can be analogous to the positional encoding used in traditional transformers.
Notations. Let  $\mathcal{G} = (\mathcal{V},\mathcal{A})$  denote a crystalline material, where  $\mathcal{V} = \{v_{1},\ldots ,v_{n}\}$  represents the set of atoms, and  $\mathcal{A}\subseteq \mathcal{V}\times \mathcal{V}$  represents the set of edges connecting the atoms in the crystalline material  $\mathcal{G}$ . Moreover,  $\mathcal{G}$  is associated with a feature matrix  $\mathbf{X}\in \mathbb{R}^{n\times F}$  and an adjacency matrix  $\mathbf{A}\in \mathbb{R}^{n\times n}$  where  $\mathbf{A}_{ij} = 1$  if and only if  $(v_{i},v_{j})\in \mathcal{A}$  and  $\mathbf{A}_{ij} = 0$  otherwise.
Task: Density of States Prediction. Given a set of crystalline materials  $\mathcal{D}_{\mathcal{G}} = \{\mathcal{G}_1,\mathcal{G}_2,\dots ,\mathcal{G}_N\}$  and a set of energies  $\mathcal{D}_{\mathcal{E}} = \{\mathcal{E}_1,\mathcal{E}_2,\ldots ,\mathcal{E}_M\}$ , our goal is to train a model  $\mathcal{M}$  that predicts the DOS of a crystalline material given a set of energies, i.e.,  $\mathbf{Y}^i = \mathcal{M}(\mathcal{G}_i,\mathcal{D}_{\mathcal{E}})$ , where  $\mathbf{Y}^i\in \mathbb{R}^M$  is an  $M$  dimensional vector containing the DOS values of a crystalline material  $\mathcal{G}_i$  at each energy  $\mathcal{E}_1,\dots ,\mathcal{E}_M$ , and  $\mathbf{Y}_j^i\in \mathbb{R}$  is the DOS value of  $\mathcal{G}_i$  at energy level  $\mathcal{E}_j$ .
4 Methodology: DOSTransformer
In this section, we introduce our proposed method, named DOSTransformer, a novel DOS prediction framework that learns the complex relationship between the atoms in the crystalline material and various energy levels by utilizing a cross-attention mechanism of the multi-modal transformer.
4.1 Crystalline Material Encoder
Before modeling the pairwise interaction between the crystalline material and the energies, we first encode the crystalline material with GNNs to learn the representation of each atom, which contains not only the feature information but also the structural information. Formally, given a crystalline material  $\mathcal{G} = (\mathbf{X},\mathbf{A})$ , we generate an atom embedding matrix for the crystalline material as follows:
$$
\mathbf {H} = \operatorname {G N N} (\mathbf {X}, \mathbf {A}), \tag {1}
$$
where  $\mathbf{H} \in \mathbb{R}^{n \times d}$  is an atom embedding matrix for  $\mathcal{G}$ , whose  $i$ -th row indicates the representation of atom  $v_{i}$ , and we stack  $L'$  layers of GNNs. Among various GNNs, we adopt graph networks [4] as our crystalline material encoder, which is a generalized and extended version of various GNNs. We provide further details on the GNNs in Appendix C.
4.2 Prompt-based Multi-modal Transformer
After obtaining the atom embedding matrix  $\mathbf{H}$ , we aim to capture the relationship between a crystalline material and various energy levels by utilizing the cross-attention layers of a multi-modal transformer [62] with self-attention layers [52]. In a nutshell, we expect the cross-attention layers to integrate the heterogeneous signals from a crystalline material and various energy levels, while self-attention layers aim to integrate the information of material-specific energy representations.
Cross-Attention Layers. Specifically, we expect the cross-attention layers to generate the material-specific representation of the energies by repeatedly reinforcing the energy representation with the crystalline materials. To do so, we first introduce a learnable energy embedding matrix  $\mathbf{E}^0\in \mathbb{R}^{M\times d}$ , whose  $j$ -th row, i.e.,  $\mathbf{E}_j^0$ , indicates the embedding of energy  $\mathcal{E}_j\in \mathcal{D}_{\mathcal{E}}$ . Then, we present cross-modal attention for fusing the information from the crystal structure into various energy levels as follows:
$$
\begin{array}{l} \mathbf {E} ^ {l} = \operatorname {C r o s s - A t t e n t i o n} \left(\mathbf {Q} _ {\mathbf {E} ^ {l - 1}}, \mathbf {K} _ {\mathbf {H}}, \mathbf {V} _ {\mathbf {H}}\right) \in \mathbb {R} ^ {M \times d} \\ = \operatorname {S o f t m a x} \left(\frac {\mathbf {E} ^ {l - 1} \mathbf {H} ^ {\top}}{\sqrt {d}}\right) \mathbf {H}, \tag {2} \\ \end{array}
$$
where  $l = 1,\dots ,L_{1}$  indicates the index of the cross-attention layers. In contrast to the conventional Transformer, which introduces learnable weight matrices for query  $\mathbf{Q}$ , key  $\mathbf{K}$ , and value  $\mathbf{V}$ , we directly employ the previously obtained energy embedding matrix  $\mathbf{E}^{l - 1}$  as the query matrix, and the atom embedding matrix  $\mathbf{H}$  as the key and value matrices. Based on the above cross-attention mechanism, we obtain the material-specific energy embedding  $\mathbf{E}^l\in \mathbb{R}^{M\times d}$  by aggregating the information regarding the atoms in the crystalline material that were important at each energy level. The final material-specific energy embedding matrix  $\mathbf{E}^{L_1}\in \mathbb{R}^{M\times d}$  generated by the cross-attention layer reflects the relationship between the atoms in the crystalline material and various energy levels.
Figure 2: Overall model architecture and attention layers in prompt-based multi-modal Transformer.
Global Self-Attention. In addition to cross-attention layers, we propose to enhance the material-specific energy embedding matrix  $\mathbf{E}^{L_1}$  by aggregating the information from other energy levels. To do so, we employ variants of self-attention layers in conventional Transformer [52] as follows:
$$
\begin{array}{l} \tilde {\mathbf {E}} ^ {p} = \text {S e l f - A t t e n t i o n} \left(\mathbf {Q} _ {\tilde {\mathbf {E}} ^ {p - 1}}, \mathbf {K} _ {\tilde {\mathbf {E}} ^ {p - 1}}, \mathbf {V} _ {\tilde {\mathbf {E}} ^ {p - 1}}\right) \in \mathbb {R} ^ {M \times d} \\ = \operatorname {S o f t m a x} \left(\frac {\tilde {\mathbf {E}} ^ {p - 1} \tilde {\mathbf {E}} ^ {p - 1 \top}}{\sqrt {d}}\right) \tilde {\mathbf {E}} ^ {p - 1}, \tag {3} \\ \end{array}
$$
where  $p = 1,\dots ,L_{2}$  indicates the index of the self-attention layer. We use an enhanced material-specific energy embedding  $\tilde{\mathbf{E}}^0$  as an input of the self-attention, which is obtained by concatenating the material-specific energy embedding  $\mathbf{E}^{L_1}$  with the material representation  $\mathbf{g}_i\in \mathbb{R}^d$ , i.e.,  $\tilde{\mathbf{E}}_j^0 = \phi_1(\mathbf{E}_j^{glob})$ , where  $\mathbf{E}_j^{glob} = (\mathbf{E}_j^{L_1}||\mathbf{g}_i)$ ,  $\phi_1:\mathbb{R}^{2d}\to \mathbb{R}^d$ , and  $||$  indicates the concatenation operation. Note that  $\mathbf{g}_i$  is a sum pooled representation of the material  $\mathcal{G}_i$ , and  $\mathbf{E}_j^{L_1}$  indicates the  $j$ -th row of the energy embedding matrix  $\mathbf{E}^{L_1}$  computed by Equation 2 where  $j = 0,\ldots ,M$ . Although the output of cross-attention layers, i.e.,  $\mathbf{E}^{L_1}\in \mathbb{R}^{M\times d}$ , captures the local atom-level information regarding the material, we also recognize the importance of incorporating the global material-level information  $\mathbf{g}_i$  that may not have been fully encoded during the cross-attention phase. Note that the importance of considering the global contextual information for enhancing the learning of dependencies among neural representations has been widely studied in various domains such as computer vision and natural language processing [51, 55, 22]. By additionally considering the global material-level information, we aim to enhance the material-specific representation of the energies more comprehensively.
System Self-Attention with Crystal System Prompts. However, solely aggregating the information of material-specific energy embeddings may neglect the crucial influence of the structural properties in crystalline materials, whereas the structural properties play a significant role in determining the unique characteristics and properties of these materials [45]. To illustrate this point, let's examine the case of graphite and diamond, both composed entirely of carbon atoms. In graphite, carbon atoms are arranged in stacked layers, where each carbon atom is bonded to three neighboring carbon atoms in a hexagonal lattice. In contrast, the carbon atom in diamond forms strong covalent bonds with four neighboring carbon atoms, resulting in a tetrahedral arrangement. Even though graphite and diamond are made entirely out of carbon, due to the structural difference, graphite exhibits properties such as electrical conductivity and lubricity due to its layered structure [59, 26], while diamond is renowned for its hardness and thermal conductivity owing to its tightly bonded tetrahedral network [8, 27].
Despite the importance of structural information, effectively incorporating it into the model is not trivial. Naively concatenating the structural information as an input feature of materials is straightforward but may lead to interference, hindering the model's ability to learn and generalize knowledge across different crystal structures. To this end, we propose to provide additional information to the self-attention layers with learnable prompts, which indicate the structural information of the crystalline materials. Specifically, we adopt learnable prompts  $\mathbf{P} \in \mathbb{R}^{7 \times d_p}$ , whose  $k$ -th row  $\mathbf{P}_k$  represents one of the seven widely known crystal systems: Cubic, Hexagonal, Tetragonal, Trigonal, Orthorhombic, Monoclinic, and Triclinic. Then, given a crystalline material  $\mathcal{G}_i$ , whose crystal system is given as  $k$ , we incorporate the learnable prompts into the material-specific energy embedding  $\tilde{\mathbf{E}}^0$  as an input
of the self-attention, which is obtained as follows:  $\tilde{\mathbf{E}}_j^0 = \phi_2(\mathbf{E}_j^{sys})$ , where  $\mathbf{E}_j^{sys} = (\mathbf{E}_j^{L_1}||\mathbf{g}_i||\mathbf{P}_k)$  and  $\phi_{2}:\mathbb{R}^{2d + d_{p}}\to \mathbb{R}^{d}$ . By doing so, we obtain high-level energy features that contain the crystal structural system information (i.e.,  $\tilde{\mathbf{E}}^{L_2}$ ), while keeping the low-level features to share the knowledge across all the materials (i.e.,  $\mathbf{E}^{L_1}$ ). In Section 5.4, we demonstrate the effectiveness of our prompt-based approach compared to other variants.
Besides the self-attention layers, we use additional cross-attention layers on top of the self-attention layers, whose input energy embedding  $\mathbf{E}^0$  in Equation 2 is given as  $\tilde{\mathbf{E}}^{L_2}$ , and output the final material-specific energy embeddings  $\mathbf{E}^{L_3}$ , where  $L_{3}$  is the number of additional cross-attention layers. By doing so, the model extracts the enhanced relationships between the crystalline materials and energy levels regarding the crystal structural system information.
4.3 Energy Decoder