Crystalformer: Infinitely Connected Attention for Periodic Structure Encoding
Authors: Tatsunori Taniai
Ryo Igarashi
Yuta Suzuki
Naoya Chiba
Kotaro Saito
Yoshitaka Ushiku
Kanta Ono
Venue: ICLR 2024 poster (2024)
Keywords: Materials Science
Attention Networks
Transformer
Physics-Informed ML
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Propose a transformer for crystal property prediction by mimicking interatomic potential summations via self-attention.

Abstract: Predicting physical properties of materials from their crystal structures is a fundamental problem in materials science. In peripheral areas such as the prediction of molecular properties, fully connected attention networks have been shown to be successful. However, unlike these finite atom arrangements, crystal structures are infinitely repeating, periodic arrangements of atoms, whose fully connected attention results in *infinitely connected attention*. In this work, we show that this infinitely connected attention can lead to a computationally tractable formulation, interpreted as *neural potential summation*, that performs infinite interatomic potential summations in a deeply learned feature space. We then propose a simple yet effective Transformer-based encoder architecture for crystal structures called *Crystalformer*. Compared to an existing Transformer-based model, the proposed model requires only 29.4% of the number of parameters, with minimal modifications to the original Transformer architecture. Despite the architectural simplicity, the proposed method outperforms state-of-the-art methods for various property regression tasks on the Materials Project and JARVIS-DFT datasets.

Published as a conference paper at ICLR 2024
CRYSTALFORMER: INFINITELY CONNECTED ATTEN- TION FOR PERIODIC STRUCTURE ENCODING
Tatsunori Taniai', Ryo Igarashi', Yuta Suzuki,, Naoya Chiba?, Kotaro Saito Yoshitaka Ushiku', Kanta Ono
OOMONN SINIC X Corp., Toyota Motor Corp.. Tohoku University, *Randeft Inc., $Osaka University
github wtwp://...aortttaaaat.a.aoo former
ABSTRACT
Predicting physical properties of materials from their crystal structures is fun- damental problem in materials science. In peripheral areas such as the prediction of molecular properties, fully connected attention networks have been shown to be successful. However, unlike these finite atom arrangements, crystal structures are infinitely repeating, periodic arrangements of atoms, whose fully connected attention results in infinitely connected attention. In this work, we show that this infinitely connected attention can lead to a computationally tractable formulation, interpreted as neural potential summation, that performs infinite interatomic po- tential summations in a deeply learned feature space. We then propose simple yet effective Transformer based encoder architecture for crystal structures called Crystalformer. Compared to an existing Transformer-based model, the proposed model requires only 29,45 9.4% of the number of parameters, with minimal modifica- tions to the original Transformer architecture. Despite the architectural simplicity, the proposed method outperforms state-of-the-art methods for various property regression tasks on the Materials Project and JARVIS-DE -DFT datasets.
INTRODUCTION
Predicting physical properties of materials from their crystal structures without actually synthesizing materials is important for accelerating the discovery of new materials with desired properties (Pollice al. 102T). While physical simulation methods such density functional theory (DFT) calculations can accurately simulate such properties, their high computational load largely limits their applicability, for example, in large-scale screening of potentially valuable materials. Thus, high-throughput machine learning (ML)-based approaches are actively studied (Pollice et al. [2021; Choudhary et al. 2002)).
Since crystal structures and molecules are both 3D arrangements of atoms, they share similar challenges in their encoding for property prediction, such as permutation invariance and SE(3) invariance (i.e., rotation and translation invariance). Hence, similar approaches using graph neural networks (GNNs) are popular for invariantly encoding these 3D structures (Reiser et al. (022).
Against this mainstream of GNN variants, approaches based on Transformer encoders (Vaswani et al., (017) are emerging recently and showing superior performance in property prediction of molecules (Ying et al., 2021: Liao & Smidt 2023) and crystal structures (Yan et al., (022) Particu- larly, Graphormer by Ying et al. (2021) adopts fully connected attention between atoms in a molecule, by following the standard Transformer architecture, and showed excellent prediction performance.
Crystal structures, however, have unique structural feature-periin periodiiiiiii------- that produces infinitely repeating periodic arrangements of atoms in 3D space. Because of the periodicity, fully connected attention between atoms in crystal structure leads to non-trivial formulation, namely infinitely connected attention (see Fig. [). which involves infinite series over repeated atoms. The previous method, called Matformer (Yan al. 2022). avoids such a formulation and presents itself rather as a hybrid of Transformer and message passing GNN. Thus, whether the standard Transformer architecture is effectively applicable to crystal structure encoding is still open question.
In this work, we interpret this infinitely connected attention as a of interatomic potentials performed deeply in abstract feature space, which we call neural potential
infinite summation

---

Published as a conference paper at ICLR 2024
1111
2(0.1)
Unit eel)
10
Aloue in wais cell
Trauslated aloue sin) @ jim)
Self-interaction i
Other in-cell interaction:
Translated seliinteraction + ((m)
Note = i00 and
Other out-oell internction +- j(m)
Figure 1: 2D diagrams of crystal structure and infinitely connected attention.
summation, In this view, attention weights are formulated as interatomic distance-decay potentials, which make the infinitely connected attention approximately tractable. By using this formulation, we propose a simple yet effective Transformer encoder for crystal structures called Crystalformer, establishing a novel Transformer framework for periodic structure encoding. Compared to the previous work by Yan et al. (2022), we aim to develop the formalism for more faithful Transformer- based crystal structure encoding. The resulting architecture is shown require only 29.4% of the total number of parameters of their Matformer to achieve better performance, while sharing useful invariance properties. We also point out that direct extension of Graphormer (Ying et al. 2021 for periodic structures leads to inadequate modeling of periodicity, which the proposed framework overcomes. We further show that the proposed framework using the infinitely connected attention formulation is beneficial for efficiently incorporating long-range interatomic interactions.
Quantitative comparisons using Materials Project and JARVIS-DFT datasets show that the proposed method outperforms several netuallnetwork-ooo Schütt et al., 20188 Chen et al. 2019; Louis et al. 2020 Chen & Ongl [2022] Choudhary & DeCost
tfate-of-the-ae art methods (Xie & Grossman, 2018
2021
Yan al. 2022) for various crystal property prediction tasks. We release code online.
2
PRELIMINARIES
We begin with introducing the unit cell representation of crystal structures and recapping standard self-attention in Transformer encoders with relative position representations.
2.1 CRYSTAL STRUCTURES
Assume periodic structure system lattice in 3D space, which consists of finite sets of points and their attributes, P (Pi, PP.... PN) and X = { an) }. in a unit cell as well as lattice vectors, 41,42,45 E R defining the unit cell translations. As crystal structure, each p; E R and x! E Rd represent the Cartesian coordinate and species of an atom in the unit cell, respectively. Any point in the system is then located a translated position, or an image, of unit-cell point
Pi(m)) = P: + n1li + T22/2 + n3la.
where three integers in n (n1, n2, n33 E 3 define 3D translation with 4 L2, L3. See Fig. [ (lett) for an illustration. For brevity, let i denote the index of the i-th unit-cell point without translation, and let i(n) denote the index of translated image of i (including i itself as i(0)). Σn, denotes the infinite series over Z3, Indices j and j(n) are used similarly.
2.2 SELF-ATTENTION WITH RELATIVE POSITIONS
Self-attention Vaswani et al. 2017) with relative position representations (Shaw et al. 2018) trans- forms an ordered finite sequence of feature vectors to another, as (æ1, &N) + y
yN):
1 Σ exp (g. / (1))(1;))))))))))))))) +b))
y: Z,
2

---

Published as a conference paper at ICLR 2024
where query q. key k, value are linear projections of input æ, Z, = EE= exp(q k; /vdx ois) is the normalizer of softmax attention weights, and dk is the dimensionality of k and q. Since each input z; is position-agnosta and so are pairwise similarity q, k,/ /d and value Vs they are augmented with scalar Øij and vector biy biases that encode/embed the relative position, j
3 CRYSTALFORMER
We consider problem of estimating physical properties of a given crystal structure. Following Xic Grossman 2018) and Schütt et al. (2018), we represent physical state of the whole structure as finite set of abstract state variables for the atoms in the unit cell, X {I1, T2, .... Iy }. assuming that any atom in the structure shares the state with its corresponding unit-cell atom, as ()()) = zq Given input state 2 that only symbolically represents the species of the unit-cell atoms, we evolve state 2 through repeated interactions between atom- wise states, another X' reflecting the target properties for prediction. To this end, we propose an attention operation for state X' which is induced by the input structure specified with the unit cell points (P1, P2. .... .... PN } and lattice vectors {11, b2 13).
3.1 INFINITELY CONNECTED ATTENTION AS NEURAL POTENTIAL SUMMATION
Inspired by Graphormer (Ying et al. 2021)ss 's fully connected attention for atoms in molecule, we formulate similar dense attention for a crystal structure for its state evolution. Compared to finite graph representations in GNNs, this dense formulation more faithfully represents the physical phenomena occurring inside crystal structures and should provide good starting point for discussion. Because of the crystal periodicity, such attention amounts to pairwise interactions between the unit- cell atoms as queries and all the infinitely repeated atoms j(n) in the structure as keys/values as
y: Z
exp (q) k/tkn + @is(m)) (D)(m) + 4))(()))
(3)
where Z, = Σソ= En expla /5tm\//dr + Pij(n)). Notice kj(()= = k, and (j(n) = v, since Tz(n) = Vy. Fig. I (right) illustrates these dense connections. We regard this infinitely connected attention in Eq.
neural potential summation. In physical simulations, energy calculations typically involve infimite summation Ex(m) ((Il)(() (5))) for potential function ((r) and physical quantity v (e.g., electric charge). Analogously, Eq. (3) computes the state of each unit-cell atom, yı, by summing abstract influences, (B)(m) ))s)(m)), from all the atoms in the structure, j(n). with their weights provided by abstract interatomic scalar potentials, exp(q k5(n)/ Vdx Oij(m)) //z Since of kz(m)/ V and Uz(m) are position-agnostic,, they are augmented with relative position encodings, (is(n) and ()f(())) reflect the interatomic spatial relation (i.e., Pý(n) pt). Thus, exp(dig(m)) is interpreted as a spatial dependency factor of the interatomic potential between i and (n.), and ()))()) as an abstract position-dependent influence on from jnn..
Compared to the standard finite-element self-attention in Eq. (2). the main challenge in computing Eq. (3) is the presence of infinite series Σ,, owing to the crystal periodicity.
Pseudo-finite periodic attention. With simple algebra (see Appendix A), we can rewrite Eq. 3 as
exp (q. k2/ Vdx ais) (v, + BBy,
y:
(4)
i=1
where
aij = loog Σ exp (dij(m))
(5)
n2
Bi Zij Σ
exp (ois())) pg(())))
(6)
and Zij exp(asy). Eq. (4) now resembles the standard finite-element attention in Eq. 2 if asj
and B.j are viewed as quantities encoding the relative position between i and j in finite element set
(1, 1,2,... N). In reality, og the log-sum of spatial dependencies exp(dxg(n) between and the j's
3

---

Published as a conference paper at ICLR 2024
Kay Valun
Y 31
yxl
3
Σ P.
Query
Softmax
BIy
A
5
Periodic spatial encoding
Q' Q'K/ "K/vall /dk
Periodic edge encoding
Scaled dot product
Broadcast
Spatialint
Spotial kg,4)
K
V
Crystal structure
Linear
Linear
Linear
1g
Matzix
Rank-3 tensor
X = [
Øx]
Spacisa info
(Fe Po Fe Fe N)
Figure 2: Pseudo-finite periodic attention with periodic spatial and edge encoding in a matrix- tensor diagram. Scalar (kj and vector 3:g integrate the spatial relations between unit-cell atom and the j's all repeated atoms, allowing the infinitely connected attention to be performed as standard fully connected attention for finite unit-cell atoms. (Unlike usual notation, X, Q, K, V, Y here denote celumn-vector-baseee feature matrices for better consistency with the notation in the main text.)
all images. Likewise, Bey is the softmaxx-wiihtee weighted average of position-dependenn influences ())j(n) on i from the j's all images, weighted by their spatial dependencies. We call oij and Bij the periodic spatial encoding and periodic edge encoding, respectively, and call the infinitely connected attention with these encodings the pseudo-finite periodic attention. If aij and B. are tractably computed, this attention can be performed similarly to the standard finite-clement attention, as shown in Fig. 2
Distance decay attention. In the physical world, spatial dependencies between two atoms tend to decrease as their distance increases. To approximate such dependencies with exp @ij(n)), we adopt the following Gaussian distance decay function simple choice among other possible forms. cxp(dij(m)) = exp -lp))(m))) - /
(7)
Here, d; >> is scalar variable controlling the Gaussian tail length. When o, is not too large, series Σ,, in Eqs. (5) and (6) converges quickly as ||>a|| grows with provable error bound (see Appendix B). thus making Okg and Bi computationally tractable. To ensure the tractability, we model a; as function of current state qi with fixed upper bound, as a, < Jub (see Appendix for the complete definition of o.). We found that, although allowing relatively long tails (e.g., Ouh 7Å) still tractable, the use of a shorter tail bound (e.g., .. 2Å) empirically leads to better results. The methodology eliminate the bound for a; is further discussed in Sec.
Value position encoding for periodicity aware modeling. Value position encoding Pistt)) in Eq. (3) influences on from j(n), and is another key to properly encoding periodic structures. In fact, attention without ())()) cannot distinguish between crystal structures consisting of the same single unit-cell atom with different lattice vectors. This is obvious since Eq. (3) with (n) = O and N 1 degenerates lattice vectors (see Appendix D for more discussions). As a simple choice for dhaj(n) we borrow edge features used by existing GNNs (Schütt et al. 2017 Xie & Grossman 2018). These edge features
represents position-dependee
y/1 V1. which is completely insensible the input
are Gaussian radial basis functions b(r) = (b,, b2, )x defined as
(r( = cxp (-(r- - mb) / K)))
(8)
where μi = kT K, and K and Tmax are hyperparameters. Intuitively, vector b(r) quantizes scalar distance via soft one-hot encoding using K bins equally spaced between 0 and Tmax. We provide Diit(n) as a linear projection of b(r) with trainable weight matrix E as
U vij(m) = b (Ilp)()) - P:ll)
Implementation details. When computing any and B. in Eqs. (5) and (6). the Gaussian functions in series Σn mostly decays rapidly within a relatively small range of || S 2 (i.e,, supercell of 5
4

---

Published as a conference paper at ICLR 2024
Prediction
x
Feed forward
Feed forward
Pooling
self-attention block
Self-attention block
Linear
Concat
aelf-attention block
P: 2
Multi-head attention iFigure 2)
Self allf--ttention block
Xo; atom species
P: atom positions
(atom embeddings)
L 22 ly lattice vectors
Xt
Figure 3: Network architecture of Crystalformer.
unit cells), but structures with small unit cells often require larger ranges. We thus adaptively change
the range of for each i to sufficiently cover a radius of 3.5Ø; in Å. This is done by setting the range
of nı as max (R, 2) I nı S max max(R1, 2) where R₁ = [3.50, || x tall/ (1,, 12, da)l (ī-1 is the ceiling function) and doing similarly for m2 and ng. For Eq. we use 64 and Tmax 14Å.
3.2 NETWORK ARCHITECTURE
As illustrated in Fig. the Crystalformer architecture basically follows the original Transformer encoder architecture (Vaswani et al. 2017) with stacked self-attention blocks, each consisting of two residual blocks connecting a head attention (MHA) layer and a shallow feed-forward network (FFN). As an important difference from the original architecture, our self-attentio, block entirely removes Layer Normalization. We found that a normalization-fres architecture with an improved weight initialization strategy proposed by Huang et al. (2020) beneficial to stabilize the training.
Given a set of trainable embedding vectors (atom embeddings) representing the species of the unit- cell atoms as initial state X² Crystalformer transforms it to abstract state 2" through four stacked self-attention blocks using the attention formulation provided in Sec. 3.]] The atom- wime states in X' are then aggregated into a single vector via the global average pooling. This material-wiss feature is further converted through FFN of linear, ReLU, and final linear layers to predict the target properties. More architectural details are provided in Appendix
RELATED WORK
Invariant encoders. Crystal structure encoders for property prediction must satisfy various invariance properties against artificial differences in data representations. The most elementary one is the permutation invariance, which initially studied by Zaheer al. (2017) and Qi al. (2017) and is widely adopted into the GNN framework. Our method is permutation-invarian thanks to the Transformer architecture. The translation and rotation invariance, or the SE(3) invariance, are also essential for ML on 3D point clouds, including molecules and materials. Our method simply employs the fully distance-based modeling (Xie Grossman 2018) ensure the SE(3) invariance. Recent point-cloud encoders, such as convolution-baut based (Zhang et al. (019 KIM et al.) 2020 Xu et al. 2021> and Transformer-basee (Qin et al., 2022; Yu et al. 2023) methods, exploit richer information while maintaining the SE(3) invariance. These techniques can be possibly incorporated to our framework with proper extensions for crystal structures. Lastly, the periodic invariance (i.e. supercell invariance and periodic- cisooddc-bbuudrryy shift invariance) has recently been pointed out by Yan et al. 2022) as a property particularly important for crystal structure encoders. Our formalism provides periodic-invarian encoder, if and ( are computed until convergence (see Appendix
5

---

Published as a conference paper at ICLR 2024
GNNs for erystal structure encoding, The initial successful work on encoding crystal structures with neural networks can be traced back to CGCNN by Xie & Grossman (2018). Their multi- edge distance graph represents crystal structures as finite graphs, and has provided foundation for the GNN framework for crystal structures. Almost concurrently, [Schütt their molecular encoder, SchNet, for erystal structures using similar approach. Subsequently, several GNNs adopting similar approaches have been proposed for universal property prediction not limited to energy prediction. MEGNet (Chen al. 2019) proposed to sequentially update atom, bond, and global state attributes through GNN layers. GATGNN (Louis et al. 2020) incorporated attention mechanisms into convolution and pooling layers of the GNN framework. iCGCNN (Park & Wolverton 2020) and ALIGNN (Choudhary& & DeCost 2021) proposed 3-body interaction GNNs to exploit interatomic angular information, while GeoCGNN (Cheng et al.] (022)) used rotation- invariant plane-wave-basee edge features to encode directional information. Kosmala et al. (2023) recently proposed Ewald message passing for exploiting long-range interatomic interactions with GNNs. PotNet (Lin et al. 2023) proposed new type of thysics-inffrmee edge feature that embeds the infinite summation value of several known interatomic potentials and allows a standard fully- connected GNN to be informed of crystal periodicity. (We more discuss PotNet in Appendix G Apart from material-level property prediction, M3GNet (Chen & Ong. 2022) extended MEGNet for interatomic potential prediction. Our focus is apart from the line of these studies in the GNN framework, and we aim to establish an effective Transformer framework for crystal structure encoding in standard fashion of the Transformer using fully connected attention.
(eahrr et al. 2018) extended
Transformers. The Transformer was originally proposed by Vaswani et al. (201)) for machine translation in natural language processing (NLP), as a sequence-t--ssequeeee model in an encoder- decoder architecture. Since then, it has been widely applied to various tasks in many fields, such as NLP, computer vision, and speech processing (Lin et al. 2022). , using its encoder-decoder or encoder-only architecture. Compared to standard convolutional neural network encoders, Transformer encoders have an outstanding ability to model complex interdependencies among input elements (e.g., resolving ambiguous meanings of words in sentence context) and have great flexibility in handling irregularly structured or unordered data such as point clouds (Zeng et al., (022). These capabilities should benefit the encoding of crystal structures and molecules because atoms in these structures interact with each other in complex ways to determine their states, and also they are structured in 3D space rather than regularly ordered. While there have been attempts to partly replace key modules of GNNs with attention mechanisms, Ying et al. (2021) first presented a complete Tormmr architecture (Graphormer) for graph encoding and showed state-of-the-ard performance for molecular property prediction. Later. Yan et al. (2022) proposed Matformer as a Transformer-inspire for crystal structures. Matformer fixes the periodic invariance break in existing GNNs by using radius nearest neighbors (radius-NN) instead of k-NN for graph construction. Note that works on language models for materials by Wei et al. (2022) and et al. (2023) are clearly distinguished from ours.
encoder
Graphormer YYing et al. 2021) is specifically designed for finite graphs such as molecules and does not ensure distance-decal attention, which leads to its inapplicability to periodic structures. Moreover, even an extended model mending this problem suffers from another modeling problem of periodicity, such as discussed in See B.T The problem is that Graphormer encodes all the spatial information, including distance edge features, only softmax biases (i.e., Φij in Eq. (2) and does not use value position encoding (i.e., vh))). Such a formulation is fine for molecules but fails to distinguish between crystal structures of the same single atom in differently sized unit cells. We will show that this modeling issue leads to performance degradation in Sec. 5.4
Matformer (Yan al. 2022 employs finite graphs with radius-NN edges similarly GNNs. Their abstract bead attention does not ensure distance decay and is thus inapplicable to the infinitely connected edges. Moreover, they modify the original Transformer in various ways, by extending query/key/value by concatenation as qū = (q. q (1), k:j (k.lkyleg) Vay = (e:/v, /ei)) (ei an edge feature), changing the scaled dot product Hadamard product, softmax sigmoid function, etc. We consider that the lack of explicit distance decay attention required such modifications, as they report worse results with standard softmax attention. They also propose self-connecting edges (i.e., a part of the orange dashed edges in Fig. (right)) although unsuccessful. By contrast, we employ much denser edges made possible by explicit distance-decay attention. This formulation leads to an architectural framework closely following the original Transformer (see Figs. [ [nd 3), specifically its relative- siriiive-posttion-baae (Shaw et al. 2018) and normalization- free (Huang et al. 2020) variants.

---

Published as a conference paper at ICLR 2024
5 EXPERIMENTS
We perform regression tasks of several important material properties, comparing with several neural-
-base state-of of-the-ar art methods (Xie Grossman 2018 Schütt et al.] 2018 Chen al.
2019
Louis et al. 2020; Chen & Ong. 2022] Choudhary & DeCost 2021 Yan et al. 2022 Lin et al.
2023). Following our most relevant work by Yan et al. (2022). we use the following datasets with
DF -Fllcaluultted properties.
Materials Project (MEGNett is collection of 69,239 materials from the Materials Project database retrieved by Chen et al. (2019). Following Yan et al. (2022), perform regression tasks of formation energy, bandgap, bulk modulus, and shear modulus.
SARVIS-DFT (3D 2021) is a collection of 55,723 materials by Choudhaary et al. 2020 Following Yan et al. (022), we perform regression tasks of formation energy, total energy. bandgap. and energy above hull (E hull). For bandgap, the dataset provides property values obtained by DFT calculation methods using the OptB88vdW functional (OPT) or the Tran-Blaha modified Becke-Johnson potential
(MBJ). While MBJ is considered more accurate (Choudhary et al. 2020),
both for evaluation.
Thanks to the great effort by Choudhary & DeCost (2021) and Yan et al. (2022) many relevant methods are evaluated on these datasets with consistent and reproducible train/validation/tess splits. We partly borrow their settings and results for fair comparison. Doing so also helps to reduce the computational load needed for comparisons.
5.1 TRAINING SETTINGS
For each regression task in the Materials Project dataset, we train our model by optimizing the mean absolute error loss function via stochastic gradient descent (SGD) with batch size of 128 materials for 500 epochs. We initialize the attention layers by following Huang et al. (2020). We use the Adam optimizer (Kingma & Ba 2015) with weight decay (Loshchilov & Hutter 2019 of 10-5 and clip the gradient norm at I. For the hyperparameters of Adam we follow Huang et al. (2020); we use the initial learning rate a of 5 x 10 and decay it according to a. /4000/ ( (4000 + t) by the number of total training steps t, and use (B1, B2) (0.9, 0.98). For test and validation model selection, stochastic weight averaging (SWA) (Izmailov et al. 2018) by averaging the model weights for the last 50 epochs with fixed learning rate. For the JARVIS-DFT dataset, we increased batch size and total epochs, as specified in Appendix for its relatively smaller dataset size.
5.2 CRYSTAL PROPERTY PREDICTION
Tables I and summarize the mean absolute errors (MAEs) for totally nine regression tasks of the Materials Project and JARVIS-DFT datasets, comparing our method with eight existing methods Our method consistently outperforms all but PotNet in all the tasks, even without SWA (i.e., evaluating a model checkpoint with the best validation score). Meanwhile, our method competitive to PotNet. This is rather remarkable, considering that 1) PotNet exploits several known forms of interatomic potentials to take advantage of strong inductive bias while we approximate them with simple Gaussian potentials, and that 2) our model is more efficient as shown in the next section. It is also worth noting that our method performs well for the bulk and shear modulus prediction tasks, even with limited training data sizes without any pre- pre-training, which supports the validity of our model.
5.3 MODEL EFFICIENCY COMPARISON
Table 3 summarizes the running times and model sizes, comparing with PotNet (Lin et al. and Matformer (Yan et al., [022) as the current best GNN-based and Transformer-basse models. Notably, the total number of parameters in our model is only 48.6% and 29.4% of the number in
2023
"Except for M3GNet, their MAES are cited from previous works (Choudhary & DeCost 2021 Yan et al.
Lin et al. (023). We M3GNet by using the code provided in Materials Graph Library (attps 7/github. wwt.co/////aamm commmatteri a svi rtua 1 / /b/matg)).......
2022
We evaluated the running times of all the methods by using single NVIDIA RIX A6000 and a single CPU
core the same machine. We used the codes of PotNet and Matformer in the AIRS OpenMat library (Zhang)
2023). The inference time of PotNet dominated by the precomputation of potential summations (309 ms).
7

---

Published as a conference paper at ICLR 2024
Table MAE comparison on the Materials Project (MEGNet) dataset. The three numbers below each property name indicate the sizes of training, validation, and testing subsets. The results of the proposed method without stochastic weight averaging (SWA) are also shown. The numbers in bold indicate the best results and the numbers with underline indicate the second best results.
Formation energy
Bandgap 6000/ 5000 4239
Bulk modulus
Shear modulus 4664
60000) 5000/4239
2664/ /33/ 393
Method
logtGPa)
log(GPa)
CGCNN Xie & Grossman
2018
0.031
0.292
0.047
0.077
SchNet Scholl
2018
0.033
0.345
0.066
0.099
MEGNe Iin er 20
2019
0.030
0.307
0.060
0.099
GATGNNI
0.033
0.280
0.045
0.075
M3GNee Chen & Ong
1122
0.024
0.247
0.050
0.087
ALIGNN [Cooudlary DeCost
2021
0.022
0.218
0.051
0.078
Matformer
0
0.211
0,043
0.073
PotNet (Lin et 2124
0.0188
0.204
0.040
0.065
Crystalformer (proposed)
0.0186
0.198
0.0377
0.0689
Crystalformer w/o SWA
0.0198
0.201
0.0399
0.0692
Table 2: MAE comparison on the JARVIS-DFT 3D 2021 dataset. The bandgap prediction is evaluated with the ground truth values provided by two DFT calculation methods, OPT and MBJ.
Form. energy
Total energy
Bundgap (OPT)
Bandgap (MBJ)
E hull
44578/5572/500 5573
4477/5577/9572
A478/5572/5577
14537/ 1837//817/1817
44296/5587/557
Method
eVlatom
eViatom
eV
CGCNN
0.063
0.078
0.20
0.41
0.17
SchNet
0.43
MEGNet
0.047
0.058
0.145
0.34
0.084
GATGNN
0.047
0.056
0.17
0.51
0.12
M3GNet
0,039
0.041
0.145
0,362
0.095
ALIGNN
0.0331
0.037
0.142
0.31
0.076
Maiffonnrr
0.0325
0.30
PotNet
0.0294
0.032
0.127
0.27
0.055
Crystallormer (proposed)
0.0306
0.0320
0.128
0.274
0,0463
Crystalformer w/o SWA
0.0319
0.0342
0.131
0.275
0.0482
Table 3: Efficiency comparison. We show the average training time per epoch, total training time, and average testing time per material, all evaluated on the KARVI-DFTT formation energy dataset using a single NVIDIA A6000 GPU. We also show the total and number of parameters and number of parameters per self-attention or graph convolution block.
Model
Type
Time/Epoch
Total
Test/Mater.
Params.
# Params./Black
PotNet
GNN
s
h
ms
1.8 M
Matformer Crystalformer
Transformer
60 s
8.3
20.4 ms
2.9 M
544
Transformer
32
7.2 h
6.6 ms
853 K
206 K
w/o 1/ 9
Transformer
12
2
5.9 ms
820 K
198 K
Table 4: Ablation studies on the validation sets the JARVIS-DE DFT 3D 2021 dataset.
Settings
+b
# Blocks
Form. E.
Total E.
Bandgap (OPT)
Bandgap (MBJ)
E hull
Proposed
0.0301
0.0314
0.133
0.287
0.0487
Simplified
0.0541
0.0546
0.140
0.308
0.0517
PotNet and Matformer, respectively. We believe that our neural potential summation induces a strong inductive bias and leads to more compact model, Although our total training time is relatively long, our inference speed is the fastest. Note that PotNet reduces its training time by precomputing infinite potential summations for training data. For reference, we also show the efficiency of our model without 2b/3, indicating that the computations for & largely dominate the entire training process.
5.4 ABLATION STUDIES
We compare the proposed model with its simplified version removing value position encoding tb. as a variant close Graphormer (Ying et al. (02T). Table 4 shows consistent performance improvements
8

---

Published as a conference paper at ICLR 2024
by the inclusion of 1b, most notably on the formation and total energy predictions but less so on the others. In Appendix IJ we discuss this behavior in terms of the definitions of these metrics, and also show further performance improvements by changing the number of self-attention blocks.
6 LIMITATIONS AND DISCUSSION
Angular and directional information. We currently adopt fully distance-based formulations for position econdings & and &a to ensure SE(3) invariance. Although straightforward. such formulations limit the expressive power and the addition of angular/directional information is preferred (Pozd-
nyakov & Ceriotti 2022). Some existing works on SE(3)-invariant GNNs explore this direction by
using 3-body interactions (Park& & Wolverton 2020; Choudhary & DeCost 2021 Chen & Ong 2022)
20en
or plane-wave-bedd based edge features (Cheng et al. [202T). Others propose (qq)equuaaria encoding
techniques for similar purpose (Liao & Smidt 2023] Duval al., (023). Extending Crystalformer
to a 3-body, higher-order Transformer (Kim et 2021I, incorporating plane wave features into 1b,
or introducing (Equ)--quv transformations are possible future directions.
Forms of interatomic potentials. Explicitly exposing distance-decay functions in attention helps the physical interpretation of our model. Our current setting limits them to the Gaussian decay function, intending to approximate the sum of various potentials in the real world, such as the Coulomb potential (1/r) and the van der Waals potential (I/r"). Despite the simplicity, the experiments show empirically good results, and we believe that the overall Transformer architecture (i.e., MHA, FFN, and repeated attention blocks) helps to learn more complex potential forms than the Gaussian function. Still, we can explore other choices of functions by explicitly incorporating known potential forms into our model. This can be done by assigning different forms of exp(o) for different heads MHA, which possibly leads to further performance gains model efficiency. We leave it as our future work.
Attention in Fourier space for long-range interactions. Gaussian tail length o; in Eq. 7 is upper-bounded by certain constant, ubb 2Å, to avoid excessively long-tailed functions, thus making a and 33 computationally tractable. Despite the empirical success, such a bound possibly overlooks the presence of long-range interactions, such as the Coulomb potential. Here, our infinite- series formulation becomes advantageous for computing such long-range interactions, with the help of reciprocal space. Reciprocal space, or Fourier space, is analogous to the frequency domain of dime-dependent functions, and appears in the 3D Fourier transform of spatial functions. When Eq. (5) is written as aij = log fpp P;) with spatial function fr). f is a periodic function that can be expressed in reciprocal space via Fourier series. In reciprocal space, f's infinite series of Gaussian functions of distances in real space becomes an infinite series of Gaussian functions of spatial frequencies. (See Appendix frr detailed formulations.) These two expressions are complementary in that short-tail and long-tail potentials decay rapidly in real and reciprocal space, respectively. We use these two expressions for parallel heads of each MHA, by computing their a and , differently in real or reciprocal space. To ensure the tractability in reciprocal space, 6; should be lower-bounde a; >> Tbp By setting the bounds for the two spaces as oh < aubs this dual-space MHA can cover the entire range of interactions in theory. As tentative experiment, we evaluate this dual-space variant of Crystalformer on the JARVIS-DF dataset. Results in Table|A6]in Appendix J.4 show that adding the reciprocal-space attention results in worse performance for formation and total energy, comparable for bandgap, and much better for E hull predictions. Adaptively switching each MHA head between the two attention forms depending on the task will lead to stable improvements.
as
7 CONCLUSIONS
We have presented Crystalformer as Transformer encoder for crystal structures. It stands on fully connected attention between periodic points, namely infinitely connected attention, with physically- motivated distance-decay attention to ensure the tractability. The method has shown successful results in various property prediction tasks with high model efficiency. We have further discussed its extension to a reciprocal space representation for efficiently computing long-range interatomic interactions. We hope that this simple and physically-motie motivated Transformer framework provides a perspective in terms of both ML and materials science promote further interdisciplinary research. As Transformer based large language models are revolutionizing AI. how Crystalformer with large models can absorb knowledge from large-scale material datasets is also ambitious open question.

---

Published as a conference paper at ICLR 2024
AUTHOR CONTRIBUTIONS
T.T. conceived the core idea of the method, implemented it, conducted the experiments, and wrote the paper and rebuttal with other authors. R.I. suggested the Fourier space attention, provided knowledge of physics simulation and materials science, and helped with the paper writing and rebuttal. Y.S. provided codebase for crystal encoders, provided materials science knowledge, and helped with the experiments, survey of erystal encoders, and paper writing. N.C. provided MI knowledge and helped with the experiments, survey of invariant encoders, and paper writing. K.S. provided materials science knowledge and helped with the paper writing. Y.U. co-led our matcrials-relater collaborative projects, provided ML knowledge, and helped with the experiments and paper writing. K.O. co-led our materials- collaborative projects, provided materials science knowledge, and helped with the paper writing.
ACKNOWLEDGMENTS
This work is partly supported by JST-Mirai Program, Grant Number JPMJMI19GI. Y.S. was supported by JST ACT-I grant number CMMIPRIUUE during the early phase of this project. The computational resource of AI Bridging Cloud Infrastructure (ABCI) provided by the National Institute of Advanced Industrial Science and Technology (AIST) was partly used for numerical experiments.
REFERENCES
Chi Chen and Shyue Ping Ong. A universal graph deep learning interatomic potential for the periodic table. Nature Computational Science, 2(11):118-728 Nov 2022. ISSN 2662-8457 doi: 10310
143588-022-00344
URLbttts:://ol//
1038/
Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong. Graph networks universal machine learning framework for molecules and crystals. Chemistry of Materials, (1(9):3564 3572-522,
May 2019. ISSN 0897-4756. doi: 10.1021/acs. http:/a...ccmmmaaee/0000999
Jiucheng Cheng, Chunkai Zhang, and Lifeng Dong. A geometiic-ifffmaaionnnnaaan crystal graph network for predicting properties of materials. Communications Materials, 2(1):92, Sep 2021. ISSN 2662-4443, doi: 10. 011888/442246002-00909999
s43246 - 021 - 00094-1110000
Kamal Choudhary and Brian DeCost. Atomistic line graph neural network for improved materials property predictions. npj Computational Materials, 7(1):185, Nov 2021. ISSN 2057-3960. doi: 10. htRp::t::: /woi..i..////// /wwg/10...0 100.:...000//111131302
1038//415244001-006550-
021 00650-1
Kamal Choudhary, Kevin F. Garrity, Andrew C. F. Reid, Brian DeCost, Adam Biacchi, Angela Hight Walker, Zachary Trautt, Jason Hattrick-Simpers, A. Gilad Kusne, Andrea Centrone, Albert Davydov, Jie Jiang, Ruth Pachter, Gowoon Cheon, Evan Reed, Ankit Agrawal, Xiaofeng Qian, Vinit Sharma, Houlong Zhuang, Sergei V. Kalinin, Bobby G. Sumpter, Ghanshyam Pilania, Pinar Acar, Subhasish Mandal, Kristjan Haule, David Vanderbilt, Karin Rabe, and Francesca Tavazza. The joint automated repository for various integrated simulations (jarvis) for data-driver materials design. npj Computational Materials, 6(1):173, Nov 2020. ISSN 2057-3960. doi: 10. 1038/ s41 541152--020000440- URL.https://don http:////ooooog////////o 10 1 1038/ 00044414--220 020-00400-0000-0-0
Kamal Choudhary, Brian DeCost, Chi Chen, Anubhav Jain, Francesca Tavazza, Ryan Cohn, Cheol Woo Park, Alok Choudhary, Ankit Agrawal, Simon J. L. Billinge, Elizabeth Holm, Shyue Ping Ong. and Chris Wolverton. Recent advances and applications of deep learning methods in materials science. npj Computational Materials, 8(1):59 Apr 2022. ISSN 2057-3960. doi: 10.
1033/s411524-022-007344444
URL/atts:://oo// /ttp::///oiiiie www.:/o..ggg11.. LBL 4445241244 022 00734-
Bernard Deconinck, Matthias Heil, Alexander Bobenko, Mark van Hoeij, and Mareus Schmies. Computing riemann theta functions. Mathematics of Computatioon 730247)1441-1444 2004. ISSN 10255718, 10886842. URI. httt:://wwt jstor www./orgggseee. h 4099903
Claudia Draxl and Matthias Scheffler. The nomad laboratory: from data sharing artificial intelli- gence. Journal of Physics: Materials, 233000 may 2019. doi: 10.1088/2515-7633 7639//bb 13bb.
URL.tttp /x doi /tt.:////....1///111555////////
10

---

Published as a conference paper at ICLR 2024
Alexandre Agm Duval, Victor Schmidt, Alex Homandez--arrcfn Santiago Miret, Fragkiskos D. Malliaros, Yoshua Bengio, and David Rolnick. FAENet: Frame averaging equivariant GNN for materials modeling. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 9013-90333 press
PMLR, 23-29 Jul 2023. URL https: https////ooceeeeng.ooo ott.e/ee.in.s......pees/////////
html
Nihang Fu, Lai Wei, Yuqi Song, Qinyang Li, Rui Xin, Sadman Sadeed Omee, Rongzhi Dong, Edirisuriya M. Dilanga Siriwardane, and Jianjun Hu. Material transformers: deep learning language models for generative materials design. Mach. Learn. Sei. Technol. 4(1):15001, 2023, doi: 10. (0.1088//632-2153/acaal
wttpooo....///222222222 21111118/2222211111222222 2153 2153/acdddd acaded
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 13-18 Jul 2020.
PMLR,
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson Averaging weights leads to wider optima and better generalization, 2018. URL hellhttpp/aaan org/abs /8033 05407
Jinwoo Kim, Sacyoon Oh, and Seunghoon Hong. Transformers generalize deepsets and can be extended to graphs &amp; hypergraphs. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems, volume 34, pp. 28016-28028. Curran Associates, Inc., neur i wwppppcppppppppppppppppppp fi es/paper/2021/ 2021. URL https::/prroceddiiiiii 3eb8 http::/eeeeeeeeaaaaeeeeeeeeeeeee
pdf
SEOHYUN KIM, Jae Yoo Park, and Bohyung Han.
to-lobaa representation learning for 3d point cloud. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 8174-8185. Cur- ran Associates, Inc., 2020. URL UBL/nntpe::/prrceedinns...
wttppppp.cpppppppppppppppppppp fileerrf_lees_eeeeeeeeeereeeeese / pdf
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015. URLhttp
//axxiv. //wxxii.rggaaas////// awtpg//ab//// 1112. 6980
Scott Kirklin, James E. Saal, Bryce Meredig, Alex Thompson, Jeff W. Doak, Muratahan Aykol, Stephan Rühl, and Chris Wolverton. The open quantum materials database (oqmd): assessing the accuracy of dft formation energies. npj Computational Materials, 1(1):15010,, Dec 2015. ISSN 2057-3960. doi: 10. 0008/mpicompumaa.2055.5
URL http::::://ooooog//10100000
npjeompumats wttp:/ppuuuuts.....11.1 2015.10
Arthur Kosmala, Johann Gasteiger, Nicholas Gao, and Stephan Gunnemann. Ewald-based long-range message passing for molecular graphs. In International Conference on Machine Learning, 2023.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023. URL
https: /
wet.eiiwwweeeeeee..e.eee wwwpee/ewwweeeeeeuuuuruuu forum? FARGOT)
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. Al Open, 111- 2022. ISSN 2666-6510. doi: http:://w....oggl.0..oo......pp. 10.000 URL https:
/ww.ssiieneeireet.ttt comm http:///ww............./////////
Yuchao Lin, Keqiang Yan, Youzhi Luo, Yi Liu, Xiaoning Qian, and Shuiwang Ji. Efficient approx- imations of complete interatomic potentials for crystal property prediction. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 21260-21287 21287. PMLR, 23-29 Jul 2023. URL https:://rroceee igg. www.rrprrsseeee / n2.3m. html

---

Published as a conference paper at ICLR 2024
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL httppss//oeeeerreeeeentttee/eeee wwwp:/wwwwwwwwt.twweeeee.eeeaeee
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. Phys. Chem. Chem. Phys., 22:18141-18148, 2020. doi: 10. 10.1039/D0CP014777 URL http:
dx www.doi..gggoooo.
1039 D000000414447
Cheol Woo Park and Chris Wolverton. Developing an improved crystal graph convolutional neural network framework for accelerated materials discovery. Phys. Rev. Mater, 4:063801, Jun 2020. doi: 000110//PhysseeeMeeria.....0 4.08801. htR.hhtts::/Iinn aps . wttp///////ooooo/oo/o//////////e 1103/
elysseewarteiaass.4.0663
Robert Pollice, Gabriel dos Passos Gomes, Matteo Aldeghi, Riley J. Hickman, Mario Krenn, Cyrille Lavigne, Michael Lindner-D' Addario, AkshatKumar Nigam, Cher Tian Ser, Zhempeng Yao, and Alán Aspuru-Guzik. Data-driven strategies for accelerated materials design. Accounts of Chemical Reseurch, 54(4): Feb 2021. ISSN 0001-4842. doi: 10.1021/aa0 hacl///a.cooonsssooo Oc0o078 URL https://doi.orpg:
awcoutaccccccccssoooaocccc....t www.aaottttttooooooooooo 0c00785
Sergey Pozdnyakov and Michele Ceriotti. Incompleteness of graph neural networks for points clouds in three dimensions. Machine Learning: Science and Technology, 3(4):045020, nov 2022. doi: 10. 0011088/263222155aaaa11 URL 0ttp:::/dd..ool..rr/1.....0 http:///...100/222222220000
aca 1 F8
Charles Ruizhongtai Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning on point for 3d classification and segmentation. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pp. 77-85. IEEE Computer Society, 2017. doi: N0.100//VVP..201...6 URL https: /oi wwt.:ogggg.oo.....
11PPR.011..0
Zheng Qin, Hao Yu, Changjian Wang, Yulan Guo, Yuxing Peng, and Kai Xu. Geometric transformer for fast and robust point cloud registration. In /EEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2022.
Patrick Reiser, Marlen Neubert, André Eberhard, Luca Torresi, Chen Zhou, Chen Shao, Houssam Metni, Clint van Hoesel, Henrik Schopmans, Timo Sommer, and Pascal Friederich. Graph neural networks for materials science and chemistry. Communications Materials, 3(1):93, Nov 2022. ISSN 2662 4443. doi: 10. 10.1038//43246-0222 s43246- 0222444-22-003155-
URL https: httpp:://a...org/.....
K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24): 241722, 03 2018. ISSN 0021-9606. doi: 10.1063/1.501977, URL URL wtt.ttttss... pattps://doi.co org/10 1063/1.5019777
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Klaus-Robert Müller, Schnet: continuous-filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems, volume 30, 2017.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In North American Chapter of the Association for Computational Linguisties, 2018. URL https 7appi rapi.semanticachon
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL ltttpa://proceeeee/e F1le/3f5ee243542 9 c4a00 hoceeeeeeaeeooaoacocccaaaaaaaeee
www.iierrucc.mo laallappeeeeeee__aaaaaaalala www.fiieeeppppe p11es//paper/2017
12

---

Published as a conference paper at ICLR 2024
Lai Wei, Qinyang Li, Yuqi Song, Stanislav Stefanov, Edirisuriya M. Dilanga Siriwardane, Fanglin Chen, and Jianjun Hu. Crystal transformer: Self-learnin neural language model for generative and tinkering design of materials. CoRR, /tlp://24..1153..... 2022. doi: 10.48550/rXXx.2244.1955
URI.tttps. //doi/doi.oo org. 48550/ att.5///rxiv.. 2204. 11953
Tian Xie and Jeffrey C. Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Phys. Rev. Lett., 120:145301, Apr 2018. doi: 10. 10.1103/PhysRevLett 120.1253000 145301. URL prttps:///inng aps wttpp/..p.........ooooooooo
PhysRevIett. 120. 145301
Jianyun Xu, Xin Tang, Yushi Zhu, Jie Sun, and Shiliang Pu. Sgmnet: Learning rotation-invarian point cloud representations via sorted gram matrix. In IEEE/CVF International Conference Computer Vision (ICCV), 2021.
Keqiang Yan, Yi Liu, Yuchao Lin, and Shuiwang Ji. Periodic graph transformers for crystal material property prediction. In S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 15066-15080 Curran Associates, Inc., 2022.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform badly for graph representation? In Dauphin. P.S. Liang, and J. Wortman Vaughan (eds.), Ad-
M. Ranzato, A. Beygelzimer, vances in Neural Information Processing Systems, volume 34, pp. 28877-28888. Curran Asso- cc/ppppr..eep.....pppp files/paper/
ciates, Inc., 2021. URL htt wnw.poooweddddddooooooooo neuri o 2021/F: e/ f c 92 0882258881111 3ac.a02222220000ccccfccc03333303
Hao Yu, Zheng Qin. Ji Hou, Mahdi Saleh, Dongsheng Li, Benjamin Busam, and Slobodan Ilic. Rotation-invarianl transformer for point cloud matching, In JEEE/CVF Conference on Computer Vision and Pattrrn Recognition (CVPR), 2023.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ Salakhut- dinov, and Alexander J Smola. Deep sets. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.. 2017. URL httpps//pprooceddngnn..rrrerrres http://nnnnnnnni.c.cccooopaaaeee E #lee2228711011222/55556666444444
Jiahao Zeng, Decheng Wang. and Peng Chen. A survey on transformers for point cloud processing:
An updated overview. IEEE Access, 0086510-86527,, 2022. doi: 10.1009/ACEBSS.222.3113333
Xuan Zhang, Limei Wang. Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Haiyang Yu. YuQing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du, Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Nzizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial intelligence for science quantum. atomistic, and continuum systems. arXiv preprint arXiv:2307.088423 2023, URL https:: : /ttptb...uu/ggg///g//// hww.hhu..oooooecmm d ve ab/a AIRS
Zhiyuan Zhang, Binh-Son Hua, and Sai-Kit Yeung. Shellnet: Efficient point cloud convolutional neural networks using concentric shells statistics. In JEEE/CVF International Conference on Computer Vision (ICCV), 2019.
13

---

Published as a conference paper at ICLR 2024
A
DERIVATION OF PSEUDO-FINITE PERIODIC ATTENTION
Starting from Eq.
we can show that
N
yi exp (a k/tmm// + des(mn) (Uy(m) + ve)(()))
(A1)
Z
j=1
exp k// /x + )rs(n) + De(mn))
(A2)
7
Σ exp ( k// V dk) Σ exp(d + +s((m))
(A3)
En exp exp()))(n)) (ey6n)
Σ exp kg V dx) exp( +
(A4)
Z j=1
exp(oa5)
1
exp (a k// ddx + (0g) + B,g).
(A5)
Z
which equals Eq.
B BOUND FOR APPROXIMATION ERROR
The computations of aij and Bi involve the following infinite summation:
|IP, Ln
1
Zis Σ exp
(A6)
20
where L - e2, Ez]. As explained in the implementation details (Sec. 3.1 finite summation:
we approximate Ziy by
R
R2
Ra
[p; + Ln - An-p.ll
Σ
exp
(A7)
20²
12= f2
R3
where ranges R1 R2, Rg E 2 are adaptively determined so that the summed points (i.e., Pj Ln - Pi) sufficiently cover a spherical radius of co, 3D space (c set 3.5) while ensuring R1, R2, Rs 2 2. This derives an upper bound for the residual error between Zig and Ziy as
Fij
A Zig
Z:j S
Σ
exp (-1
(A8)
rEAg
|r 2R
where Agj = { 1 (Ln + P,- p.) 2²), R max c/ V2) R min
/20, P, + Ln
V2o:
p:ll In E 23, maxmmaann (/11 |12), (ng)) > 2}.
Here, Deconinck et al. (2004) prove the following inequality:
|r|Pexxp (-11-113 S d + p
2 d
? 7)²)
0 2
Σ
(A9)
2
2
where = (Xn + x/n E zd4 X E Rdxd is a full-rank matrix, x E Rdd, = min Xn|| In e zd n + 0), and I'(2,z) = [ -- the incomplete Gamma function.
Substituting p = 0, 3, and X =
1 L into this inequality derives a bound for our residual error:
V20.
3
3
R - 2) 2)²)
Eig < 2 p
r
(A10)
2
where p min 1/2,
Lnll nn E zd # 0).
14

---

Published as a conference paper at ICLR 2024
C DEFINITION OF o;.
To increase the expressive power of the model, we treat tail length o, of the Gaussian potential in Eq. 7 as variable parameterized via current state qi Here, the exponent of the Gaussian function, pijim) = - (P)))) - p.ll" /20 is intended to represent the relative position encoding by Shaw et al. (2018), whose original form is dij = q agj with embedding aij of relative position i. As such a linear form of should ease the training of the network, we model o; instead of o; as function:
= ro² p((g w - mm)/s)
(A11)
where ro hyperparameter determining distance scale of spatial dependencies in A, w is trainable weight vector, m and are constant parameters, and p IR Rt is function to ensure positive value for o,
The adjustment of ro could be difficult, since it has to bridge between network's ' internal states (i.e., qf w) and physical distances (i.e., 0,) across different scales. To ease it, we design p so that the initial values of p and hence o; are distributed around and fu, respectively. To this end, initialize m and as the mean and standard deviation of the values of q. w computed for the mini-bater samples the initial training step. Then, for p introduce shifted exponential linear unit (shifted ELUU): p(2; a, b) = -11-b))LU b))LLU(az/(( b)) +1,
(A12)
where ELU(x) = exp(x) - 1 if x < 0 and ELU(z) = otherwise. Assuming that æ initially follows (0, 1), shifted ELU p(x; a, b) is designed be p(0)
so that o; initially distributed around ro. 2) lower-bounded as p{xx) > b so that a; < rn/ / avoid excessively long tails, and 3) linear response with derivative for x 2 0 to increase the training stability. We empirically set the hyperparameters to (ro.a, a, b) (1.4, 0....
D NECESSITY OF VALUE POSITION ENCODING
Shaw et al. 2018) discuss that value position encoding is required depending on the application, and we show that crystal structure encoding is indeed such a case.
As mentioned in the main texts, Eq. (3) without value position encoding vhis(m) cannot distinguish between crystal structures of the same single atom in differently sized unit cells. This can be confirmed by substituting X' { or N = into Eq. as below.
1
V1 exp (qf ki/Vdxx an) B11)
(A13)
Z1
= U1 + Bn
(A14)
= + En exp(omm))))
(A15)
En exp((pp(m))
Note that the lattice information (i.e., L1,22, Lz) used only in key position encoding @11(n) and value position encoding 1))())) for 011 and B11. and does not exist in k, in the initial state. The above equations show that when N = 1, the weight of (v) + +,11) is canceled out by Z1 and therefore lattice information only present in B1.- Without value position encoding (i.e., 1111()) = 0), the attention formula becomes in 3D space by l1.6g. l2 Lz.
v1, showing that result yı does not reflect how the atom is repeated
From materials science perspective, crystal structures of single unit-cell atoms are called monatomic structures and represent an important class existing materials. Well-known monatomic structures of the same atom include graphite and diamond. They have the single unit-cell atom of Carbon but have different lattice vectors and thus different crystal structures. These differences make graphite a soft substance used for pencil tops while making diamond the hardest substance in the world. The value position encoding is thus required to distinguish such important materials.
E DETAILED NETWORK ARCHITECTURE
Table Al shows the detailed specifications of the proposed Crystalformer architecture. Note that changing the ReLU activation function to GELU in the FFNs of the self-attention blocks did not lead to meaningful performance improvements.
15

---

Published as a conference paper at ICLR 2024
Table A1: Detailed network architecture.
Blocks
Oulput dims
Specifications
Atom embedding
(128, NN)
{e AtomicNumber(()), -,,,,,,, 1,, .V)
where e are initialized by N/(,, 128 12.5
Self-attention black 4
(128, N)
Model dim d: 128
Number of heads:
Query/Key/Value dims:
FFN: Linear-ReeU--Lnee ReLUU-Lnnar (dims: 128 - 512 - 128)
Activation fune: ReI.U
Normalization: No
Pooling
128
Global average pooling
Linear with output dim: 128
Feed forward
ReLU
Linear with output dim: (target property dim.)
F INVARIANCE PROPERTIES
F.1 DEFINITIONS
Following notation by Yan et al. 2022). we represent a crystal structure with tuple of three matrices, (X, P. L), where x= 31 EN E d is the states of N unit-cell atoms, R3xN is the 3D Cartesian coordinates of these atoms, and I, [e1, l2, L3] E R3x3 is the lattice vector matrix. The coordinates of the unit-cell points are defined to be within the unit cell region defined by L. that is, they have fractional coordinates of L- [0, 1)3xN, When the overall network architecture is seen as function f(X, P,L) + X, they satisfy the following invariance properties.
[p..... P.v E
Definition 1:
Unit cell permutation invariance. function f : P. L) - X is unit cell per- mutation invariant such that for any permutation function a {1,2, .., N) + (1,2,..NN) we have f(X, L) = f(o(X), (P)), L), where (X) [=0o11) [o(N)) and (P) [Po(1) .. Po(N))
Definition 2: Unit cell E(3) invariance. A function f (X, P,L) - is unit cell E(3) invariant 3 R. R-1, det(R) +1, and b R have f(X, P, L) I(X,T, RL), where T = R. R3x4 is rigid transformation, R is a rotation matrix, b is translation vector, and P P I 1xN the homogeneous coordinates P in 3D space. (Note that the definition gives unit cell SE(3) invariance when R is restricted to rotation matrix, as det(R) = 1.)
such that for all R
Due to the arbitrary choice of extracting a unit cell structure from the infinitely expanding crystal structure in 3D space, there is an infinite number of invariants (X, P. L) that represent the same crystal structure as (X, P. L). Similarly to the notation by Yan et al. (2022), we define such unit-cell slice of given crystal structure
(A16)
The slice is parameterized via periodic boundary scale, &s 2 that defines supercell of L as
i - [ki k222, kala):
(A17)
and periodic boundary shift, p E R33 that provides corner point of the periodic boundaries. The unit-cell atoms in the slice are then provided as collection of the atoms contained in new unit cell i with the scaled and shifted periodic boundaries as
p = {p: + In p/i. - (p. + In p) E 10,1), i (1,..., N). nnnee EZ))
(A18)
X (x, Ln p) [0,1)3, i (1..., N). , 23}.
(A19)
The periodic invariance described by Yan et al. (002) is then summarized as follow.
16

---

Published as a conference paper at ICLR 2024
Definition 3: Periodic invariance. A unit cell E(3) invariant function f (X.P.I) - X is =
periodic invariant - P,, p)) holds for all k E 2 and R
F.2 PROOFS OF INVARIANCE
Proof of unit cell permutation invariance.
Let us represent overall network transformation f as =(x,
(A20)
where X' = P, is the transformation by stacked self-attention blocks, = h(X') is global pooling operation, and o(2) is the final output network given global feature vector z. We can easily show that g is permutation equivariant such that (XX') g(o((). (P), L) holds for any
permutation σ. (This is obvious from that fact that substituting i <- ai and j(n) t (n)) into
the right-hand side of Eq. 3 yields Mo(()) and thus the attention operation by Eq. 33 is permutation
equivariant.) We can then prove that ((o(X), (P), L) o(h(o(XX)) o(h(
f(X,
holds if h is permutation invariant pooling operation such as the global mean or max pooling.
Proof of unit cell E(3) invariance. In f. the spatial information (ie., unit-cell atom positions P and lattice vectors L) is only used by the position encodings, Pij(n) and vij((n). in Eq. 33. Currently, they are defined to simply encode relative distances
(ej() <(pp; Pj(n)
(A21)
(6()))) <i(pp -P(())))) I])
(A22)
thus invariant to rigid transformation
(P',L') + (TP.RL)
l|p P)((n)
lpl - (p's L'n)]l
(A23)
[Rpp b) (Rp; b RIm)]
(A24)
|R(p: P)(m)))l
(A25)
Ip: - -Pi(m)l--
(A26)
Note that framework allows the of directional information by using relative position (p; P)(n)) instead of relative distance Ilp; - P)(m)]l Il. provided that ((r) and (rr are rotation invariant.
Proof of periodic invariance. is sufficient to show that Eq. is invariant for minimum repeatable unit-cell representation (X, P. L) and its invariant representation (X, P.P.Li i) = 4(X, P. L, k. p). Periodic invariance of Eq. 3 intuitively understandable, because Eq. 3 recovers the original, infinitely expanding crystal structure from its unit cell representation and attends to all the atoms in it with relative position encoding. A formal proof is provided below.
Comparing (X, P, L) and (X, P,), each unit-cell point (1,..., NN in (X,P, is populated
to k2 unit-cell points in (X, P, f.,, which we index by i[m] where m (m,, m2, ma) E z and 0 S m S ki 1, 0 < m2 < l2 1,0 0 S mg S k3 1 (or simply 0 m S 1). These
populated points i[m] in (X, P, L) have the state:
Tmm = 1i
(A27)
and periodically shifted positions:
Pi(m) = p: + L(m +;)
(A28)
where c; E 3 is a constant chosen appropriately for each so that Piim) for all m: 0 S m $ k
reside within the unit cell region of L (i.e., i --'p:(m) 1)3).
By simply applying Eq.
3 for (X, P, Ż). we obtain
1 N
y:
Σ
Σ
exp (qT //s/Vdx + dig(n)) + (.(m))
(A29)
2,
17

---

Published as a conference paper at ICLR 2024
where q, k and ú are key, query, and value of Ø in X. Scalar P1z(n) and vector 4b.j(m) encode the interatomic spatial relation as (())(m) - -D;) and (P)(())) - -P;). respectively, where Py(n) = Pj + Ln. We rewrite Eq. (A29) by rewriting populated point i|m] and j as j(m'), obtaining
N
exp (1%mm) (immm// / dx + (emml)smmm(mm)) (0gmml) + 1b mllmm)mmm
7i(m] (mmm) j=1 m
(A30)
where Em is short for 2mi=0 2mb=0 m with m = (mi, m2, my) Since 3.(m) in X equals 2; in X, we substitute qi/mml = q: kjim'l = kj. and 0g[mel] = Uy to obtain
Uilml Zi(mi J=1 j=1m TE
ΣEΣ exp (q ky/ /dx + (mm)) (vy + (m(p)mmmmm))
(A31)
Similarly to Eq. (4) this infinitely connected attention can be written in the following pseudo-finite periodic attention form:
N (q k; lt///d dx + (s(m))) + Bimmls)
[((ml
exp
(A32)
7:mmml j=1
where
d((m)] = log Σ
exp (di(m]](mm]())
(A33)
T2
m
Biimll =
exp exp(a (f))))))) (()) (mmm))me1mmm
(A34)
mmmll m
and Zopmii cxp(ás(mlj) Here, [immmmmm[e encodes the interatomic spatial relation:
Pi(me(((n)) Peimms] (P)(ml) i.) Pi|ma
(A35)
= (p, + Lm' + Le, - -p + in) (p, + I. + Le, - p)
(A36)
Pj + (in Lm') L(c, m-c:)= -pi-
(A37)
Notice here that where nt' = k2 + m4, ni2k2 + mly Nigks + m²/) T E 23, Thus, we obtain
in + Lm' = (niki + m4 + (n2k2 + mb))e + (nak' + mmy)es
(A38)
= Ln',
(A39)
(a(m)](m)) Pi(m) = Pj + L(n' + + Pi:
(A40)
where c = (c, m c.) 3. Because En Em for terms of ní' is equivalent infinite series En over all n' E Z3, we obtain the following equivalence:
d:(mi] = llog Σ
Σ
exp
exp (d(P)(m')(a)) - -Pi(m))
(A41)
mn'
log Σ
exp (d(p) L(n' + + - -p:))
(A42)
n
= log
exp (dij(n'+a)
(A43)
n²'
Therefore, ds[m]) converges to the same value of azj in Eq. 5 if they converge. Likewise, 3.[m] converges to the same value of Bij in Eq. (6). Provided that @j(m) Oay and [A(m]] = Bij. the right-hand side of Eq. A32)) equals the right-hand side of Eq. 4.. thus proving Ui(m) = y:-
At the overall architecture level, the above result indicates that applying the stacked self-attention blocks for an invariant representation, as X' g($(X,, P. k. p)). gives kykgk3 duplicates of X' g(,, P, L). For these X' and X', pooling operation h such as global mean or max pooling holds the invariance: h(X') = b([X',... .'.,,'])) X')) = h(X'). Thus, overall network transformation f is periodic invariant.
18

---

Published as a conference paper at ICLR 2024
G DIFFERENCE FROM POTNET
PotNet is a GNN-based crystal- crystal-structure encoder proposed by Lin et al. (023)) PotNet shares similarity with Crystalformer in that uses infinite summations of interatomic potentials to model the periodicity of crystal structures. However, in addition to the obvious difference of being a GNN or Transformer, there is more fundamental or conceptual difference in their representations of potential summations. To see this, let us discuss how these methods perform state evolution.
State evolution in PotNet. Using our notation, the PotNet's state evolution can be essentially expressed as
N Bi. j=1
x x 9
(A44)
where f and g are learnable GNN layers that perform message passing all pairs of hidden state variables 2))) and edge features Eij between them. Thus, PotNet presents itself in a standard GNN framework for fully connected graphs of nodes of unit-cell atoms. The most important part of PotNet is its edge feature Cjj, which is an embedding of an infinite summation of scalar potentials. This summation is defined for each unit-cell atom pair (i,j)
Sig- = Σ
O (1P,s0)) - pll)
(A45)
72
using scalar potential function ((r): R. + R. PotNet exploits three known forms of actual interatomic potentials, that is, Coulomb potential 1/r, London dispersion potential 1/r", and Pauli repulsion potential e e-ar, via their summation:
+ k1/r + N2/r" + Kze
(A46)
where fi,, lig Ri; a are hyperparameters. Edge feature ety is then provided as embedding expanded and encoded via radial basis functions (RBF) and multi-layer perceptron (MLP)
Eij = MLP(RBF(SS)))
(A47)
Here, we see that PotNet exploits the infinite summations of known interatomic potentials to derive physics- and geriodicityiiffommer edge features eij. to be processed in a standard GNN framework.
State evolution in Crystalformer. The core part of our state evolution is the attention operation provided in Eq. 3 and written again below:
y: Z
ko + Pis()) + 4p))())
(A48)
As described in Sec. 3.1 this attention operation, as whole, interpreted an infinite summation of interatomic potentials defined and performed in an abstract feature space.
One may find some similarities of our oiyy and Biy with S.y and e.y of PotNet, because they enable standard fully-connected Transformer or a fully-connected GNN to incorporate the crystal periodicity. However, our Ossj and B., are rather bi-products of the explicit potential summation feature space performed as the infinitely connected attention, whereas S.j is potential summation calculated in potential space before the GNN message passing (typically as preprocessing) and eij is an abstract embedding of S.j.
Summary.
PotNet proposes type of physics-inte -nnforeed edge feature by using the values of infinite summations of interatomic potentials. This edge feature implicitly embeds the information of crystal periodicity and can thus allow standard fully-connected GNN be informed of crystal periodicity.
Crystalformer deeply fuses the calculations of infinite potential summations with neural networks, by performing the infinitely connected attention infinite potential summation expressed abstract feature space. This attention operation can be performed in a standard fully-connected Transformer by using proposed position encoding oay and Big-
19

---

Published as a conference paper at ICLR 2024
H HYPERPARAMETERS FOR THE JARVIS-DFT DATASET
Compared to the Materials Project dataset, the JARVIS-DFT dataset is relatively small in both its training size and average system size (i.e., average number of atoms in unit cells, which is approximately 10 atoms in the JARVIS-DFT while 30 in the Materials Project). To account for these differences, increase the batch size and number of total epochs shown in Table A2
Table A2: Hyperparameters for the JARVIS-DFT dataset.
Form. energy
Total energy
Bandgap (OPT)
Bandgap (MBJ)
hull
Training set size
44578
44578
44578
14537
44296
Batch size
256
256
256
256
256
Total epochs
800
800
800
1600
800
I DETAILED ABLATION STUDIES
I.1 PHYSICAL INTERPRETATIONS OF RESULTS IN TABLE
Table 4 indicates the inclusion the value position encoding yields greater improvements for the formation energy and total energy while relatively limited improvements for the bandgap and energy above hull (E hull). This result can be understood from the fact that the formation energy and total energy absolute energy metrics, whereas the bandgap and energy above hull are relative energy metrics that evaluate gaps between two energy levels, as AE = Fj Eg-
The accurate prediction of such energy gaps is generally difficult, because it needs accurate predictions of both E₁ and Eg. Furthermore, even when an algorithmic improvement in the prediction method leads to improvements absolute energy prediction (e.g., formation energy prediction), it does necessarily lead to improvements in energy gap prediction. Specifically, assume that an improved algorithm predicts energies E' E1 Å and Eg Ep + Ag better predictions than E₁ and Eo. respectively. These improvements can directly lead to lower prediction errors of absolute energy metrics. However, when improvements A and An are made by the same mechanism, it may be possible that A Ao so that they are canceled out in energy gaps A F F--
We conjecture that the limited improvements in the bandgap and energy above hull prediction are due to these difficulties of energy gap prediction.
1.2 CHANGING THE NUMBER OF SELF-ATTENTION BLOCKS
We further study the effects of the number of self-attention blocks. We here utilize the simplified variant because of its relatively good performance and high efficiency (Table 3) The results in
Table A3 suggest that the performance moderately mounts on a plateau with four or more blocks.
Table A3: Ablation studies by changing the number of self-attention blocks.
Settings
Blocks
Form. F.
Total E.
Bandgap (OPT)
Bandgap (MBJ)
hull
Simplified
0.0664
0.0669
0.171
0.368
0.1313
2
0.0627
0.0609
0.153
0,329
0.0965
3
0.0558
0.0569
0.145
0,320
0.0676
0.0541
0.0546
0.140
0.308
0.0517
5 0.0525
0.0543
0.143
0,309
0.0484
0.0503
0.0533
0.140
0.323
0.0430
7
0.0498
0.0515
0.131
0.298
0.0407
Given the good performance of the seven-block setting, we test seven-block variant of the proposed full model on the Materials Project and JARVIS-DFT datasets. The results Tables A4 and A5 show that the larger model achieves slightly better accuracies, especially when the training set is large enough. Investigating how the performance changes with larger-scale datasets (Kirklin et al. 2015 Draxl & Scheffler 2019) is left future topic.
20

---

Published as a conference paper at ICLR 2024
Table A4: Performance of the larger model on the Materials Project (MEGNet) dataset.
Formation energy 60000/ 5000/4239
Bandgap 6000/ 5000/ 4239
Bulk modulus
Shear modulus
4664/ 393//993
4664/392/393
4 blocks (default)
0.0186 0.0185
0.198
0,0377 0.0372
0.0689 0.0717
blocks
0.184
Table A5: Performance of the larger model on the JARVIS-DFT 3D 2021 dataset.
Form. energy 44578/5572/
Total energy
Bandgap (OPT)
Bandgop (MBJ) 14577/1817/1817
E hull
(4578/55228557
(4578/5572/5577
44206/5537/555
4 blocks (default)
0.0306 0.0297
0.0320 0.0312
0,128 0.127
0.274 0.275
0,0463
7 blocks
0.0461
J RECIPROCAL SPACE ATTENTION FOR LONG-RANGE INTERACTIONS
J.1 BACKGROUND
Before presenting our attention formulation in reciprocal space, we briefly introduce its background from physical simulations. In physical simulations, the difficulty of computing long-range inter- actions in periodic systems is well known. Given potential function (r)) and physical quantities {ui...., UN (e.g., electric charges) associated with unit-cell atoms, these simulations typically involve the summation of interatomic interactions:
V2=
Σ ((IP)(n) -Pel) j(n)ys
(A49)
where Exm)) sums for all the atoms j(n) in the structure, as Ey= En. except for jnn = i. Its direct computation expanding n causes inefficiency for long-range ((r) due to slow convergence. Thus, DFT calculations often employ an efficient calculation method called Ewald summation.
In Ewald summation, potential (r)) is represented as the sum of short-range and long-range terms as
+ (r) = 4 shon (r) + Plong (r).
(A50)
In this manner, the summation of interactions is divided into two parts as
V,
Σ
Pstoon p:ll) + Σ (1P pxll)us- j(m)yf
(A51)
j(nzzt
The first part is the sum of fast-decay potentials in distance, which can be efficiently computed directly in real space. On the other hand, long-range potentials have concentrated lowfrrquueccy components and thus have fast decay in frequency. Therefore, the second part can be efficiently computed in reciprocal space via its Fourier series expansion.
J.2 RECIPROCAL SPACE ATTENTION
Kosmala et al. 2023) recently proposed Ewald message passing in the GNN framework, which performs graph convolutions in both real and reciprocal space. We also import the idea of Ewald summation into our Transformer framework.
To obtain reciprocal-space representations of ouj and B.g in Eqs. (5) and (6), it is sufficient derive for infinite series Σn, exp(dr)(n) (eit(n in the right-hand side of Eq. 6 We rewrite it as f'p; - p.) using the following 3D spatial function:
f(r)
Σ exp (r + nihi + ngl2 + nals)) (r + n14 + +12h2 + ngla)-
(A52)
With this spatial function, aij and Bij are expressed
Ogj = log f(p; pila/; + 1),
(A53)
1
Big = f(p; pt),
(A54)
Zip
21

---

Published as a conference paper at ICLR 2024
where f(11z,, < 1) denotes substitution s/ 1 into Eq. (A52)
Since f is a periodic function, can be alternatively expressed via Fourier series expansion with f's Fourier coefficient F(wm) as
f(r)= = Σ
F(Wm)oxp (iwm r),
(A55)
mEZ
where is the imaginary unit and
Wm miti m222 mgla
(A56)
is a 3D vector analogous to the angular frequency of the Fourier transform. Wm is defined with three integers m (m1, m2, my) E 2 with reciprocal lattice vectors:
2T
2T
(L,, 12, 13) = 2
L3
13
x L1,
V l x 1))
(A57)
where V = (1 x 12) 13 is the cell volume. Fourier coefficient F(wm) ) is obtained via 3D integration:
Flwm) JJ/ f(r) exp ( iwm dxdddd
(A58)
///e exp (T) exp ( iwm r) drddddz,
(A59)
where Re represents a low high frequency component of f(r).
R the 3D domain the unit cell. Intuitively, F(wm) with small or large ||m|||
Computing Ogj and Bij with Eqs. A53) and (544) in reciprocal space using Fourier series in
Eq. (A55) is effective when potential function exj()))(m) is long-tailed, because F(wm) in such
case decays rapidly as |mm|| increases.
J.3 GAUSSIAN DISTANCE DECAY ATTENTION IN RECIPROCAL SPACE.
The reciprocal-spaie space representation of azy in the case of the Gaussian distance decay function in Eq. (7) is expressed as follows.
(2mo?) 3/2
0. [?][[m]] 112 cos (wm (p, p.))
Oij = log
Σ exp
(A60)
2
Eq. (660) is essentially an infinite series of Gaussian functions of frequencies ||wm ll. which rapidly decay as frequencies increase. On the other hand, we omit B,, in the reciprocal space attention by assuming vaag(n) Bij O, because its analytic solution does not exist for our current (hay(n) provided in Eq. (9). Exploring an effective form of sb)j(n) for the reciprocal-space attention is left as future work.
When o² = o², the two expressions of an with Eqs. 5 and A60)) become theoretically equivalent. Similarly to Ewald summation, however, want to use the real-space attention for short-range interactions and the reciprocal-space attention for long-range interactions. To this end, we treat d; and o; independent variables by following parameterization method in Appendix C Specifically, we parameterize o² as
,: = Po w m)/s;a,
(A61)
In this way, o, is lower bounded as o, > Tth with оb = Vo. We use (Fo, a,b) = (2.2,0.1 0.5) for 0. Now d; and ō; have an upper and lower bound, respectively, as o; < Jub ~ 1.98Å and o. > db ~ 1.56Å. Therefore, the combined use of the real-space and reciprocal-space attention can cover the entire range of interactions in principle. We compute Eq. (A60) with a fixed range of |m]] < 2. Note that because of the symmetry, we can reduce the range of m, as 0 S m, < by doubling the summed values for my > 0.
22

---

Published as a conference paper at ICLR 2024
Table A6: Effects of attention in the Fourier space on the JARVIS-DFT dataset.
Method
Form. energy
Total energy
Bandgup (OPT)
Bandgup (MBJ)
E hull
Crystalformer (real space) Crystalformer (dual space) 0.0343
0.0306
0.0320 0.0353
0.128 0.126
0.274
0.0463
0.283
0.0284
J.4 RESULTS USING THE DUAL SPACE ATTENTION
As discussed in Sec. 6 we evaluate variant of our method by changing half the heads of each MHA to the reciprocal attention provided in Appendix J.3 Table A6 shows the results with this dual-space attention.
We observe that the dual-space variant yields a significant improvement in the energy above hull prediction. Although its exact mechanism is unknown, we provide hypothetical interpretations below, as we do for results in Table in Appendix [I] In short, we conjecture that it is related to the fact that the energy above hull is the only metric in Table (A6 Ao that involves the energy evaluation of different crystal structures than the targeted one.
more in-depth discussion needs to know the definition the energy above hull. Let's s again express it as AE = E1 - Eg. where Ej corresponds to the formation energy of the target structure and Eg is the so-called convex hull provided by stable phases of the target structure. As we know, substances around us change their phases between solid, liquid, and gas depending on the temperature and pressure. Likewise, a crystal structure has multiple stable phases with different structures within solid state. For example, structure with chemical formula of AB has phases that have different structures with chemical formulas of A, E [0, 1] (i.e., structures with the same element composition but different ratios of elements). When plot y FormationFnergy A,, Bi- for all possible phase structures of AB, we can draw a convex hull of the plots, as in Fig. AT which will have some corner points that represent stable phases of AB (i.e., S1 S4 in Fig. [A]p. The energy above hull is then defined as the vertical deviation between the formation energy of the target structure and the line of the convex hull, essentially representing the instability of the target structure.
SI(A)
$4(B)
Formstion energy
Target structure
AE
Convex hull
83
$2
0
Concentration A. A.B,,,
Figure A1: Energy above hull. Image courtesy of Ma, J., Hegde, V., Munira, K., Xie, Y., Keshavarz. S., Mildebrath. D., Wolverton, C., Ghosh, A., Butler, W. (2017). Computational investigation of half-Heusler compounds for spintronics applications. Phys. Rev. B, 95, 024411.
When Ej and Ep involve the energy evaluation of different structures, it may be possible that the inclusion of long-range interactions has positive impact on some of the stable phases of given structure in E, while not so for the given structure itself in Ej. This asymmetric relation between E, and E, may mitigate the aforementioned cancellation effect and may lead to improved prediction of Eo and AE.
23


Reviews:
[{"review_id": "DE400W7YXb", "reviewer": "Reviewer_S6vd", "summary": "This paper presents an attention mechanism and a transformer model for periodic structure encoding. The attention mechanism is formulated as an infinite sum over the periodic extensions of the unit cell in all three directions. In order to make the sum tractable, the paper proposes a distance decay attention by which the spatial dependencies between atoms decay exponentially with their distance. The authors present an evaluation of this attention mechanism, incorporated into a transformer architecture, on two data sets and several material property prediction tasks, comparing with a handful of methods from the crystal modelling literature.", "questions": "- The results in Table 4 demonstrate that the inclusion of the value position encoding systematically improve the performance of the model. However, we observe significant differences in the performance gap between the proposed and simplified models across properties. Do you think this gap in performance could be explained by the characteristics of each property?\n- I am also curious whether you have some intuitions about why the performance of the dual space version of the model is significantly better than the baseline in the case of predicting the energy above hull.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper is clearly written and easy to follow. I have enjoyed reading it and I believe it will likewise be relevant for the machine learning community working on materials modelling. Concretely, these are some of the strengths I see in this paper:\n\n- The introduction effectively describes the gap in the literature it addresses, namely the periodic nature of crystal structures, which is not considered in the core of the GNN literature devoted to molecular modelling. The paper provides a comprehensive review of graph and transformer models for related tasks and motivates the need for a method to account for the periodic nature of crystal structures.\n- Sections 2 and 3 clearly describe the proposed model and the figures effectively support the text descriptions.\n- The evaluation seems sound and the results achieved by the proposed method are remarkable compared to existing approaches, including in terms of training and inference time.", "weaknesses": "I do not have any major concerns. If I were to put on the hat of the average machine learning reviewer, I could argue that the technical contribution of this paper is limited because the method is a small adjustment of existing work. However, that is largely irrelevant to me provided the idea is sound, it is clearly described, it addresses a current challenge in materials modelling and the results are competitive.\n\nOne comment regarding the presentation of results that I can make is that it would be much easier to interpret the results if they were presented graphically instead of in tables, despite this being the common practice in the machine learning community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "moHUsBlQJi", "reviewer": "Reviewer_NG2X", "summary": "To encode infinitely repeating crystal structures, this paper introduces a distance-decay inductive bias into the attention mechanism to ensure tractability. The authors compare their method with five baseline methods on two datasets.", "questions": "None", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed decayed attention naturally incorporates a physical inductive bias, making the attention mechanism for infinite structures computationally tractable.\n2. Good experimental results, and each component of the proposed method is also validated through ablation experiments.", "weaknesses": "1. The necessity of $\\psi_{ij}(n)$ is unclear. Although the authors give an example in Page 4,  it remains questionable whether it is appropriate to assume $\\mathcal{X} = \\{x_1\\}$ when attempting to prove the necessity of $\\psi_{ij}(n)$. \n2. Lack of discussion on reflection and periodic transformations [1]. \n3. About writing\n\t- In Figure 2, the representation of the infinitely connected attention is not clear. \n\t- In Figure 3, the interaction of $X^t$ with $\\mathcal{P}$ and $l_1, l_2, l_3$ requires a more detailed and explicit illustration, as they are somewhat hidden within the formulas, which can be challenging for readers to locate.\n\t- \"With simple algebra, we can rewrite Eq. (3) as...\" This rewriting should be included in the appendix somewhere, and please point to it.\n\n\n[1] Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems\n\n=======================\n\nI sincerely appreciate the authors for taking the time and effort to address my concerns. I would like to raise my rating to 8.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y2Yw8hXzE3", "reviewer": "Reviewer_u628", "summary": "The paper proposes a new transformer architecture, called Crystalformer, for crystal structure property prediction that manages the periodic structure of crystals. To completely capture the periodic structure of crystals one would effectively require infinitely connected attention, which is intractable, prompting the papers to introduce a new attention method to effectively manage this requirement using a distance based decay.\n\nThe papers starts by introducing the problem of crystal structure property prediction and discusses recent approaches, including GNN-based neural networks and a couple of transformer models, specifically Matformer. The paper then introduces a set of preliminaries, including description of crystal structures and the self-attention mechanism. Subsequently, the authors describe the different components of Crystalformer with a focus on the attention formulation called pseudo-finite periodic attention that encodes relevant crystal information. Next, the authors describe distance decay attention, which applies a Gaussian decay function to place less importance on atoms that are further apart from each other and thereby simplifies the need for infinitely connected attention to a tractable problem. The authors describe relevant details of their attention formulation and distance decay function before outlining their network architecture consisting of multiple attention blocks.\n\nIn Section 4, the paper described related work in detail and emphasizes how Crystalformer differs from various competing approaches and showcase experiments on Materials Project and JARVIS in Section 5. The experiments in Section 5 show that Crystalformer generally outperforms relevant GNN and transformer baselines for both datasets. Next, the paper provided a set of ablation studies related to removing positional encodings and performing attention in Fourier space.", "questions": "The paper could be improved by addressing the following questions related mainly to clarity:\n* An example consisting of input features (atom types and atom positions) would make things clearer showing how the data is exactly processed. Figure 3 shows that positions and lattice parameters go into the attention blocks, but their processing is not explained in Figure 2.\n* Can you describe stochastic weight averaging? It's first mentioned in Table 1.\n* What is the reason for choosing your datasets? It seems like the Materials Project version is significantly smaller than the latest version. It would be good to have more details.\n* Did you try other decay functions? Why or why not?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper has the following strengths:\n* The paper presents a new attention formulation for crystal structure design that makes an infinite attention scenario computationally tractable (originality, significance).\n* The paper provides detailed problem descriptions and related work outlining how their work differentiates itself from prior methods (clarity, significance).\n* The paper provides relevant experimental results and ablations to supports its primary claims (quality).", "weaknesses": "The paper currently provides a somewhat limited scope in its experiments, especially for the datasets, and could be also improved by including additional relevant datasets and models:\n* Datasets: OQMD [1], NOMAD [2] which are crystal structure datasets available through the Open MatSci ML Toolkit [3]. Experiments on these datasets would significantly strengthen the claims of the paper since they include significantly more datapoints than the ones in the current experiments. The toolkit also includes a more updated version of Materials Project with ~150k crystal structures. \n* Models: FAENet [4], Equiformer [5]. While new experiments would be ideal to strengthen the paper claims, I would at least like to a discussion of these models in related work. \n\n\nAdditionally, the current draft would benefit from additional clarity related to the method, which are described in the section below.\n\n[1] Kirklin, S., Saal, J.E., Meredig, B., Thompson, A., Doak, J.W., Aykol, M., Rühl, S. and Wolverton, C., 2015. The Open Quantum Materials Database (OQMD): assessing the accuracy of DFT formation energies. npj Computational Materials, 1(1), pp.1-15.\n\n[2] Draxl, C. and Scheffler, M., 2019. The NOMAD laboratory: from data sharing to artificial intelligence. Journal of Physics: Materials, 2(3), p.036001.\n\n[3] Lee, Kin Long Kelvin, et al. \"MatSciML: A Broad, Multi-Task Benchmark for Solid-State Materials Modeling.\" arXiv preprint arXiv:2309.05934 (2023).\n\n[4] Duval, Alexandre Agm, et al. \"Faenet: Frame averaging equivariant gnn for materials modeling.\" International Conference on Machine Learning. PMLR, 2023.\n\n[5] Liao, Yi-Lun, and Tess Smidt. \"Equiformer: Equivariant graph attention transformer for 3d atomistic graphs.\" arXiv preprint arXiv:2206.11990 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7Q385aCeE5", "reviewer": "Reviewer_8AQw", "summary": "This paper proposes a method to construct fully connected graphs for crystal materials, with two attention biases between node pairs proposed. From my understanding, the first attention bias ($\\alpha$ ones) captures infinite interactions between two nodes, and the second one ($\\beta$ ones) captures the local interactions in the summation form. After this, the traditional fully connected transformer architectures can be applied. The experiments show a good improvements beyond Matformer.", "questions": "The major questions are provided in the above weakness section. Given the concerns for now, I vote for rejection.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The strengths of this method are as follows.\n\n- A reasonably good transformer network with high efficiency for crystal scalar property prediction\n- Good presentation, easy to follow writing\n- Good efficiency for a transformer architecture", "weaknesses": "However, there are some concerns of this paper.\n\n- One of the most important issue is that there is another work named PotNet[1], already published one year ago, proposed a method to calculate infinite summations for node pairs in crystal structures. The core idea in this paper of calculating the first attention bias ($\\alpha$ ones) capturing infinite interactions between two nodes of this paper are very similar to PotNet. But the similarity is not mentioned, and comparison is not provided. Besides the first attention bias ($\\alpha$ ones), the second attention bias ($\\beta$ ones) is not new. It can be seen as the summation of local radius features.\n\n- The derivation in the appendix might be problematic, the authors use word \"may be\", which is ambiguous. If authors want to use Equation (S14) in the Appendix to derive the first attention bias, which is one of their major contribution, the coefficient in the reciprocal space needs to be properly derived. \n\n- The paper uses “computationally tractable” in many places, but the error bound of the summation is still unknown to me. Particularly, the authors say if the $\\sigma$ of $\\sum_n e^{-\\frac{\\Vert r\\Vert^2}{2\\sigma}}$ or $\\sum_n \\Vert r \\Vert e^{-\\frac{\\Vert r\\Vert^2}{2\\sigma}}$ has an upper bound and lower bound then is tractable. If it is tractable, then what is the exact error between the authors’ evaluation and the value of every summation $\\sum_n e^{-\\frac{\\Vert r\\Vert^2}{2\\sigma}}$\n\n- The performances of the proposed method is not very significant without the additional stochastic weight averaging (SWA) (Izmailov et al., 2018) by averaging the model weights for the last 50 epochs with a fixed learning rate. If with this module, the comparison is not that fair, but without this module, the performance gains are not very significant. \n\n\n[1] Efficient Approximations of Complete Interatomic Potentials for Crystal Property Prediction. (Public in ICLR 22 and accepted in ICML 23) https://proceedings.mlr.press/v202/lin23m.html", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)