Paper: Crystalformer: Infinitely Connected Attention for Periodic Structure Encoding (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
Thank you for pointing out a relevant work, PotNet, by Lin et al. (2023). After a careful assessment of this work, we consider that our contributions remain sufficient even with the presence of PotNet for the following reasons.

First, **their PotNet and our Crystalformer are built upon different architectural paradigms, GNN and Transformer.** While there are many GNN variants for crystal structure encoding, Transformer variants that have been successfully applied to crystal structure encoding are, to the best of our knowledge, limited to Matformer (Yan et al. 2022) and ours. Furthermore, as explained in Section 4, the Matformer architecture significantly differs from the standard Transformer architecture (Vaswani et al., 2017), for example, in the change of the softmax function to sigmoid and in the definitions of query/key/value features (see Figure 5 of the Matformer paper for more details). Although Matformer makes an important first step in introducing Transformers to the field of material modeling, we believe that **this field still lacks a viable Transformer baseline with minimal modifications to the standard Transformer architecture, except for our work.** In the field of ML, improvements and extensions of Transformer models have been and will be accumulated, mostly by assuming the standard architecture. Thus, **in order to continuously import these advances in ML into the field of material modeling, it is vital to have such a baseline Transformer model in the field.** This point is independent of the existence of PotNet.

Second, **establishing a viable Transformer model for crystal structures is not trivial even when PotNet is known, because Transformers differ significantly from GNNs in their architectures and are generally more difficult to train.** For example, Transformer models usually require an initial warm-up phase with a low learning rate during training. Without this phase, the training typically fails to reach valid solutions. For our model, the presence of infinite summations makes the training even harder, and this widely used warm-up strategy could not lead to training success. As described in Section 3.2, **we found that the normalization-free Transformer architecture with an improved weight initialization strategy by Huang et al. (2020) is essential for a Transformer to successfully work for crystal structures**. Furthermore, **we have identified an architectural factor that enables proper modeling of the crystal structures of single unit-cell atoms** (see “Value position encoding for periodicity-aware modeling” in Section 3.1). As this modeling problem stems from softmax’s weight normalization, it is usually not recognized by GNNs. **These non-trivial findings will help future studies in this area** and are beyond the focus of the PotNet work.

[*Remark: The third discussion below is better explained in a subsequent post:* https://openreview.net/forum?id=fxQiecl9HB&noteId=JqMYCMjipP *where we discuss that our method deeply fuses potential summantion calculations with neural networks by treating the whole attention operation as a neural interatomic potential summation.*] 

Third, **although PotNet uses the infinite summation of potentials similar to our $\alpha_{ij}$, its usage and definition are clearly different from ours**, as discussed below.

### PotNet's usage vs ours
Specifically, PotNet computes infinite summations for each unit-cell atom pair, as $S_{ij} = \sum_{n} \Phi(r_{ij(n)})$, where $\Phi(r_{ij(n)})$ is a known physical interatomic potential function. Then, transform scalar $S_{ij}$ to a feature vector via radial basis functions (RBF) and multi-layer perceptrons (MLP), as $e_{ij} = \text{MLP}(\text{RBF}(S_{ij}))$. Then, feed it to a message passing function $x_i \gets g(x_i, \sum_j \text{f}(x_i, x_j, e_{ij}))$.

In our case, we compute infinite summation $\exp(\alpha_{ij}) = \sum_n \exp(\phi_{ij(n)})$, which is similar to $S_{ij}$ in PotNet. However, it appears as part of softmax-based weighted averaging of features (i.e., attention): $\sum_n \exp(\phi_{ij(n)}) (v_j + \psi_{ij(n)})$ in Eq (3). Here, each potential value $\exp(\phi_{ij(n)})$ acts as a weight for feature vector $(v_j + \psi_{ij(n)})$.

Reviewer comment (Theory):
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.

Author response:
Sorry for causing the confusion by our misleading writing. Equation (S14) is derived concretely without ambiguity. Thus, “may be” is now rephrased by “is” in the paper.

Reviewer comment (Theory):
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}}$

Author response:
Thank you for pointing this out. **We have provided the error bound in the Supplementary Material (Section of “Bound for approximation error”)**.

In summary, given the theoretical expression of the summation as

$Z = \sum_{n} \exp \left(-\frac{\\| p_j + Ln - p_i \\|^2}{2\sigma_i^2} \right),$

we approximate it as

$\tilde{Z} =\sum_{n_1=-R_1}^{R_1}\sum_{n_2=-R_2}^{R_2}\sum_{n_3=-R_3}^{R_3} \exp \left(-\frac{\\| p_j+Ln-p_i \\|^2}{2\sigma_i^2} \right).$

Here, as explained in the implementation details (Sec 3.5), ranges $R_1,R_2,R_3 \in \mathbb{Z_+}$ are adaptively determined to cover a spherical radius of $c\sigma_i$ in 3D space  (we emprically set $c=3.5$), while ensuring $R_1,R_2,R_3 \ge 2$.

Then, residual error $\epsilon$ between $Z$ and $\tilde{Z}$ has the following upper bound:

$\epsilon := Z - \tilde{Z} \le \frac{3}{2} \left( \frac{2}{\rho} \right)^3 \Gamma \left( \frac{3}{2}, \left( R-\frac{\rho}{2} \right)^2 \right)$

where $\Gamma(z,x) = \int_x^\infty t^{z-1}e^{-t} dt$, $\rho = \min \\{ \frac{1}{\sqrt{2}\sigma_i} \\|Ln \\| \\, \vert n \in \mathbb{Z}^d, n \neq 0\\}$, $R = \max(c/\sqrt{2}, R')$, and $R' = \min \\{ \frac{1}{\sqrt{2}\sigma_i} \\| p_j + Ln - p_i \\| \\, \vert n \in \mathbb{Z}^3 , \max(|n_1|,|n_2|,|n_3|) > 2\\}$.

Reviewer comment (Experiments):
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.

Author response:
Thank you for sharing your concern. As mentioned in our first response, **we aim to develop a new Transformer framework for crystal structure encoding, with equal emphasis on its effectiveness as well as conceptual and architectural simplicity.** This is because we consider the field of material modeling ($\neq$ molecular modeling) still lacks an effective baseline Transformer model with minimal modifications to the standard Transformer architecture. In this objective, we have presented our model as simpmle as possible in our framework, leaving room for further improvements as discussed in Sec 6. We believe that the results without SWA successfully demonstrate the effectivenss of this simple method.

Reviewer comment (Experiments):
The paper currently provides a somewhat limited scope in its experiments, especially for the datasets, and could be improved by including additional relevant datasets such as OQMD [1] and 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.

Author response:
We appreciate your advice on how we can further strengthen our work.

We agree that the inclusion of experimental comparisons with large-scale datasets can enhance the value of our work. However, training the proposed method and other competitors on such large-scale datasets requires rich computational resources or a not-cheap financial budget for renting them through cloud GPU services. For now, we would like to leave such large-scale experiments as our future work.

Meanwhile, training on large-scale datasets itself is often done as independent topics, especially when Transformer-based large models are considered. For our case, it naturally attracts interest whether our Crystalformer can benefit from the increased dataset size by increasing its model size. Therefore, **we would like to consider Reviewer u628’s suggestion as one of the interesting future directions, and we have mentioned it from Sec 5.5 with references to OQMD and NOMAD**.

See also our response to the question “What is the reason for choosing your datasets?” below.

Reviewer comment (Experiments):
The paper could be improved by including additional relevant models such as FAENet [4] and Equiformer [5]. While new experiments would be ideal to strengthen the paper claims, a discussion of these models in related work would also be beneficial.

Author response:
Thank you for pointing out relevant references.
We see that both FAENet (GNN-based) and Equiformer (Transformer-based) focus on developing techniques for SE(3)/E(3)-*equivariant* encoders, while our work targets SE(3)/E(3)-*invariant* encoders. Also, although Equiformer is a Transformer variant like ours, it is presented for finite molecule structures only. Therefore, we consider that they are not direct competitors to ours. 

However, we think that **their techniques for SE(3)/E(3) equivariance are conceptually orthogonal to our framework, and can therefore be utilized when extending our model** to be SE(3)/E(3)-equivariant depending on the target task, or to incorporate angular/directional information into our model. **We have mentioned such a possibility in Sec 6 (Section of “Angular and directional information”)** with new references to FAENet and Equiformer. Also, **we have added a reference to Equiformer in Sec 1 (3rd paragraph)** as a Transformer-based molecule encoder, cited along with Graphormer (Ying et al., 2021).

Reviewer comment (Presentation):
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.

Author response:
Thank you for your valuable suggestion. Following your advice, **we have updated Figure 2** as follows.
- We have added a sub-figure showing an example of a crystal structure.
- We have clarified that the spatial information (ie, point coordinates $\mathcal{P}$ and lattice vectors $\ell_1,\ell_2,\ell_3$) is used by “periodic spatial encoding” and “periodic edge encoding”.

Reviewer comment (Writing):
The paper should include a description of stochastic weight averaging, which is first mentioned in Table 1.

Author response:
Thanks for the question. In both the initial and latest manuscripts, **stochastic weight averaging (SWA) is described in the main text of Sec 5.2**: “For test and validation model selection, we use stochastic weight averaging (SWA) (Izmailov et al., 2018) by averaging the model weights for the last 50 epochs with a fixed learning rate.”

We believe that SWA can suppress performance fluctuations due to model selection and can therefore help to more reliably analyze the performance in the ablation studies.

Reviewer comment (Experiments):
The paper should address whether other decay functions were tried, and explain why or why not.

Author response:
For now we use only the Gaussian decay function.

In Sec 6, we discuss the possibility of using other functions that approximate real interatomic potential functions. However, we think that such extensions are more domain-specific, and therefore **we plan to try them in a follow-up work aimed at the audience with more interest in materials science** (e.g., for a journal on materials science and informatics).

Besides, the Gaussian decay function has the following benefits:
1. It exponentially decays in real space. 
2. It has the analytic solution in Fourier (reciprocal) space, as shown in Eq (S14) in the Supplementary Material.
3. This Fourier-transformed form is also a Gaussian decay function in reciprocal space with an easy-to-understand complementary relation with the real-space Gaussian function. (Specifically, when the real-space Gaussian function has the tail-length of $\sigma$ in 3D space, the reciprocal-space Gaussian function has the tail-length of $1/\sigma$ in spatial frequency.)

For these reasons, we think that the Gaussian decay function is the best choice to demonstrate the core idea of our method.

Reviewer comment (Theory):
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)$.

Author response:
Thank you for sharing your concern with us. We think that our current discussion on the necessity of $\psi_{ij(n)}$ is basically fine for the following reasons.

First, as a common practice in mathematics, **it is sufficient to give a single counterexample when proving that a mathematical statement/equation is problematic**. In “Value position encoding for periodicity-aware modeling” in Sec 3.1, we point out that when $\mathcal{X} = \\{ x_1 \\}$, Eq (3) without $\psi_{ij(n)}$ is reduced to $y_1 = v_1$, which is insensible to input lattice vectors. This modeling flaw can be fixed by explicitly inserting $\psi_{ij(n)}$ into Eq (3). Mathematically speaking, this observation sufficiently shows the necessity of $\psi_{ij(n)}$. (We have added detailed mathematical explanations in the Supplementary Material, as explained later.)

Second, **from a materials science perspective, the shown counterexample ($\mathcal{X} = \\{ x_1 \\}$) is called *monatomic structures* and represents an important class of existing materials rather than imaginary ones.** Well-known monatomic structures of the same atom include *graphite* and *diamond*. These materials 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. Thus, $\psi_{ij(n)}$ is required from a materials science view too.

Finally, **Table 4 shows that the inclusion of $\psi_{ij(n)}$ consistently leads to performance improvements in all the prediction tasks**, emprically showing the the necessity of $\psi_{ij(n)}$.

We have added the Section of **“Necessity of value position encoding” in the Supplementary Material** that summarizes the above discussion, and added a reference to it from the aforementioned part of Sec 3.1.

Reviewer comment (Theory):
Lack of discussion on reflection and periodic transformations [1].

Author response:
As requested, **we have added a thorough discussion on invariance properties of the proposed model for various transformations**, including the permutation of unit-cell atoms, rotation/reflection and translation of 3D coordinates (so-called SE(3) and E(3) invariance), and periodic transformation. The section of **“Invariance properties” in the Supplementary Material** provides the definitions and proofs of these invariance properties. We have mentioned it from Sec 4 ("Invariance encoders").

Reviewer comment (Presentation):
In Figure 2, the representation of the infinitely connected attention is not clear.

Author response:
Thank you for your comments. 

Figure 2 illustrates our attention operation of Eq (4) in a matrix/tensor-based computational diagram instead of the vector-based representation of Eq (4). We adopt this style because **papers on Transformer architectures often illustrate a QKV attention operation by such a diagram, and following a similar manner should promote the understanding of readers who are familiar with Transformer papers.** Specifically, input vectors $x_1,...,x_N$ are concatenated as matrix $X = [x_1,...,x_N]$ in the diagram, and query, key, value, and output vectors are likewise expressed as matrices $Q, K, V, Y$. In this manner, scaled-dot product $q_i^Tk_j/\sqrt{d_K}$ for all ij pairs is expressed as $Q^TK/\sqrt{d_K}$. Also, $\alpha_{ij}$ for all ij pairs is expressed as matrix $A$ to be added to $Q^TK/\sqrt{d_K}$. Similarly, $\beta_{ij}$ for all ij pairs is expressed as rank-3 tensor $B$ to be added to $V$.  These notations are shown in the diagram. After all, Figure 2 is an illustration of Eq (4) in a common computational diagram of QKV attention and presents no new information than Eq (4).

However, we think Figures 2 and 3 had some clarity problems, which may have caused confusions. Specifically, how $\mathcal{P}$ and $\ell_1,\ell_2,\ell_3$ are used in Figure 2 was unclear, and how Figure 2 and 3 are related was unclear. 

Thus, **we have updated Figures 2 and 3** as follows.
- For Figure 2, we have clarified that $\mathcal{P}$ and $\ell_1,\ell_2,\ell_3$ are used as input to “Periodic spatial encoding” ($\alpha_{ij}$) and “Periodic edge encoding” ($\beta_{ij}$), and further clarified that Figure is a “matrix-tensor diagram” in the caption.
- For Figure 3, we have clarified that its “Multi-head attention” block corresponds to Figure 2.

Reviewer comment (Theory):
"With simple algebra, we can rewrite Eq. (3) as..." This rewriting should be included in the appendix somewhere, and please point to it.

Author response:
Thank you for your suggestion. **In the Supplementary Material (Section of “Derivation of pseudo-finite periodic attention”), we have added a step-by-step derivation from Eq. (3) to Eq. (4), and pointed to it from the main text.**
Note that in the initial manuscript, Eq. (4) had a mis-formulation: ${q_i^T k_{j}}$ had to be ${q_i^T k_{j}}/\sqrt{d_K}$. We apologize if it caused any confusion. This mis-formulation has been fixed too.

Reviewer comment (Evaluation):
The results in Table 4 demonstrate that the inclusion of the value position encoding systematically improves the performance of the model. However, there are 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?

Author response:
Thank you for your question. 
As pointed out, Table 4 shows that the inclusion of 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 are absolute energy metrics, whereas the bandgap and energy above hull are relative energy metrics that evaluate gaps between two energy levels**, as $\Delta E = E_1 - E_0$.

The accurate prediction of such energy gaps is generally difficult, because it needs accurate predictions of both $E_1$ and $E_0$. Furthermore, even when an algorithmic improvement in the prediction method leads to improvements in absolute energy prediction (e.g., formation energy prediction), it does not necessarily lead to improvements in energy gap prediction. Specifically, assume that an improved algorithm predicts energies $E_1' = E_1 + \Delta_1$ and $E_0' = E_0 + \Delta_0$ as better predictions than $E_1$ and $E_0$, respectively. These improvements can directly lead to lower prediction errors of absolute energy metrics. However, when improvements $\Delta_1$ and $\Delta_0$ are made by the same mechanism, it may be possible that $\Delta_1 \simeq \Delta_0$ so that they are canceled out in energy gaps $\Delta E' = E_1' - E_0'$.

We conjecture that the limited improvements in the bandgap and energy above hull prediction are due to these difficulties of energy gap prediction.

Reviewer comment (Evaluation):
I am 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.

Author response:
Thank you again for your question. Although we are not aware of the exact mechanism for this behavior, we conjecture that it is related to the fact that **the energy above hull is the only metric in Table 5 that involves the energy evaluation of different crystal structures than the targeted one**.

A more in-depth discussion needs to know the definition of the energy above hull. Let’s again express it as $\Delta E = E_1 - E_0$, where $E_1$ corresponds to the formation energy of the target structure and $E_0$ 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 the solid state. For example, a structure with a chemical formula of $AB$ has phases that have different structures with chemical formulas of $A_{x}B_{1-x}$, $x \in [0,1]$ (i.e., structures with the same element composition but different ratios of elements).  
When we plot y = FormationEnergy($A_{x}B_{1-x}$) for all possible phase structures of AB, we can draw a convex hull of the plots, which will have some corner points that represent stable phases of AB. (One can obtain an intuitive illustration of this by searching for “energy convex hull”, such FIG. 5 in https://www.researchgate.net/figure/A-schematic-convex-hull-in-the-A-B-chemical-space-Phases-Si-lie-on-the-convex-hull-and_fig2_312191506, where S1 to S4 represent stable phases of AB.) 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.

When $E_1$ and $E_0$ involve the energy evaluation of different structures, it may be possible that the inclusion of long-range interactions has a positive impact on some of the stable phases of a given structure in $E_0$ while not so for the given structure itself in $E_1$. This asymmetric relation between $E_1$ and $E_0$ may mitigate the aforementioned cancellation effect and may lead to improved prediction of $E_0$ and $\Delta E$.

We will add these insights in an upcoming revision.