Paper: Periodic Set Transformer: Material Property Prediction from Continuous Isometry Invariants (ICLR-2024)

Reviewer comment (Experiments):
Although the authors' presentation is quite clear in general, the details of the experiments provided in the paper are not enough, especially why the experiments are designed in this way, what are the datasets and targets, and what the difference across experiments is and how they collaborate to support the statements in the manuscript.

Author response:
The intention of the two experiments is to illustrate the flexibility of the model and the effectiveness of using invariants outside of their usual context of crystal comparison and instead in representations for ML. 
    
In the first case, the datasets used in section 4 and section 5 are significantly different. Section 4 uses molecular crystals with each set (e.g. T2) containing the same composition but changes in structure while the MP crystals are mostly inorganic crystals with varied compositions. 

In the second case, the PDD and PST are able to perform comparatively to graph-based models that contain additional structural information. The graph-based models also use $k$-NN distances but through the use of edges, they can specify between which atoms a distance corresponds to. The PDD does not have this information, it only uses the distances themselves. Despite this, the PST still outperforms CGCNN (the most comparable model in terms of information used) in 5/6 properties. The use of edges is discontinuous under atomic perturbations. If a small change in atomic positioning occurs, the set of $k$-NN can change. In a graph-based model, this means edges can shift, changing the graph's topology in a way that is disproportionate to a small change in atomic position. In the case of the PDD, when this occurs the distances will be changed continuously because there is no reference to the neighboring atoms.

Reviewer comment (Theory):
In terms of the explanation of isometry on page 2, I'm wondering why the isometry has the form $f(S)=Q$ and $g(Q)=S$. From my understanding, isometry means $d_S(a,b) = d_Q(f(a),f(b)), a,b \in Q$, and I can't tell that this is equivalent to the explanation in the manuscript.

Author response:
In the context of the explanation $f: Q \rightarrow S$ and $g : S \rightarrow Q$ are bijective isometries and should the condition $f(Q) = S$ and $g(S) = Q$ be met, then the two periodic point sets $S$ and $Q$ are isometric. This latter part could be generalized to $S$ and $Q$ being metric spaces in $\mathbb{R}^n$. In our explanation, we don't directly define an isometry and instead mention it is a distance-preserving mapping between metric spaces; this will be adjusted in the revision.

Reviewer comment (Experiments):
Earth Mover's Distance is mentioned in the Introduction part, but I do not see detailed descriptions about how to use it for crystal representation in the manuscript. Besides, have the authors considered comparing with ElMD [1], which has also introduced Earth Mover’s Distance as metrics for chemical similarity and inorganic compound embeddings?

Author response:
The Earth Mover's Distance establishes a continuous metric on periodic sets and is used to measure differences in PDDs. It is not used to create features or utilized in the model. It is used in plots 4a and 4b in the appendix to show how distances between structures can be a useful indicator for differences in their property values.

Reviewer comment (Experiments):
In the second experiment of prediction of lattice energy, if I'm understanding it correctly, the datasets consist of crystals with different compositions while the compositional information is not included. Then how does the model make predictions for two similar lattices with different compositions? And even if the model can outperform the baseline on this task, I'm afraid it cannot demonstrate that the model is applicable to practical usage.

Author response:
It is true the compositional information is not included (in section 4). The experiment shows that the structural differences produced from having different compositions are reflected in the distances in the PDD and this is enough to make accurate predictions. This is already shown to be the case when distances are measured via the EMD and then projected onto 2D space (Figure 4a of the appendix). Whether this is true when the PDD is input to an ML algorithm is what this experiment aims to address.

Reviewer comment (Evaluation):
Why is the PST model evaluated on the training set for the first two tasks of lattice energy prediction? And what is the reason for supplementing P2M data to training data to reduce error? From the results in Table 1 & 2, I cannot be persuaded of the effectiveness of PST.

Author response:
The table was trying to indicate what crystals make up the training and test sets. The reported error is for the test set. The first task follows the same setup as Ropers et al. The second extends this to multiple crystal sets. 

We tried to show that the model can be improved with a small amount of generated data for a new crystal type (P2M) which has not been seen by the model before. In a practical setting, DFT calculations (computationally expensive) can be done on a small number of crystals and be added to the training set. These additional crystals can improve the prediction accuracy on this new type of crystal. The PST can then be applied during the rest of the CSP process in which (tens of) thousands of additional structures could be generated.

Reviewer comment (Experiments):
Could you clarify how the contribution is calculated in the ablation study? And I think the errors here are sufficient to demonstrate the impact of compositions and PDD.

Author response:
If $e_{PDD}$, $e_{Comp}$ and $e_{PST}$ are the mean-absolute-errors for the contribution for the PDD component would be $c_{PDD} = r_{PDD} / (r_{PDD} + r_{Comp})$ and the composition contribution would be $c_{Comp} = r_{Comp} / (r_{PDD} + r_{Comp})$
    where $r_{PDD} =  e_{PST} /  e_{PDD}$ and $r_{Comp} = e_{PST} / e_{Comp}$. 

Please let us know if further explanation is needed.

Reviewer comment (Writing):
The description of the PDD representation and the architecture is often unclear and confusing (clarity).

Author response:
We have tried to supplement the definition of the PDD with an example in the appendix. If you could specify what about the explanation of architecture and PDD representation was unclear that would be greatly appreciated.

Reviewer comment (Novelty):
The contribution appears limited to the inclusion of the composition on top of the PDD representation, which appears to be prior work (significance, originality).

Author response:
The PDDs original purpose was to measure differences in the structure of crystals in order to compare them. We believe we have provided evidence that it (and potentially other invariants) has a place in machine learning as the input representation. Additionally, the model we propose can be applied to data with the structure of a weighted set (or discrete distribution). The method of turning multisets into weighted sets by grouping inputs based on their features and using weights to describe multiplicity can be applied to similar Transformer input to decrease the number of tokens in the input, improving computational performance. For the experiments in Section 5 the size of the set of the input is reduced to $\sim 45$ percent of the original number of atoms in the unit cell.

Reviewer comment (Experiments):
The experiments performed are relatively small in scale with the results often not well presented (clarity, quality).

Author response:
Could you please expand on this? Several other models report between 1-3 datasets including CGCNN  (only the Materials Project), MEGNet (the Materials Project and QM9, which is a molecular not crystal dataset), and SchNet (3 molecular datasets). We will be adding an additional ablation study and the results on the effect of the collapse tolerance on input size compared to other graph approaches (based on another reviewer's comments).

Reviewer comment (Experiments):
The experiments in Section 4 are not well described making it difficult to assess their significance (clarity, quality). Given that only PDD representations were used, it is unclear what contributions of the paper are being highlighted here. Also the model architectures used are unclear. My best guess is that it involves Gaussian processes similar to the AMD case.

Author response:
A similar critique has been mentioned by other reviewers, we will be more explicit about the setup and purpose of the experiment in section 4 of the revised version of the manuscript. The comparison done in this experiment is between the PST using only the PDD and Gaussian regression with AMD as input.

Reviewer comment (Reproducibility):
What are the model architectures used in Section 4?

Author response:
This model uses the Transformer described in section 3.1, but instead of using PDD encoding to incorporate structural information with compositional information, it uses just the PDD alone. The model we compare to uses the AMD as the representation for a crystal and the algorithm applied is Gaussian regression.

Reviewer comment (Reproducibility):
Can you describe in more detail how the rows of the PDD representation are collapsed into each, specifically how identical rows are identified?

Author response:
Two rows are identified as identical if their distance is less than or equal to the collapse tolerance using a valid distance metric (definition 3.2). In our application, we use the $L_\infty$ metric. These two rows are collapsed into one by taking the average of the rows.

Reviewer comment (Reproducibility):
How are the rows of the PDD representation ordered? Does this ordering matter?

Author response:
They are ordered lexicographically (definition 3.2). The order does not matter as the attention mechanism is permutation equivariant and the pooling layer is permutation invariant.

Reviewer comment (Reproducibility):
How do atoms get counted in the PDD construction described in Section 3.1? Since composition is not present yet, are the atoms indexed without atom types?

Author response:
The number of atoms in the unit cell defines their count. Through collapsing, however, which unit cell is used does not matter. The PDD contains no knowledge of the composition, only the pairwise distances.

Reviewer comment (Experiments):
Is there a predetermined way to choose k for the PDD? Based on the information in the appendix it appears to be a hyperparameter that seems significant. It would be good to have more details on this.

Author response:
$k$ is selected to be large enough such that the generic completeness property of the PDD is satisfied. That is, the distances in the final column of the PDD are greater than twice the covering radius of the lattice. This varies from crystal to crystal so $k$ is selected such that it is adequate for 99\% of the crystals in the largest dataset.

Reviewer comment (Experiments):
What types of crystals are studied in Section 4? You mention both molecules and crystals here, so are these molecular crystals? Is there a reason you claim that only the lattice matters for these structures? Clarity could be substantially improved by providing more detailed information about the task.

Author response:
Yes, they are molecular crystals. Specifically, simulated crystals created during the crystal structure prediction process. For each crystal set (e.g. T2) the crystals share the same composition but have different structure. 

The PDD, is used for measuring differences in structure via EMD. Figure 4a in the appendix shows how the pairwise distances manifest on a 2D plane, creating clusters by composition. The purpose of the experiment is to determine whether such an inference can be made in a ML application wherein the PDD is the input.

Reviewer comment (Experiments):
**Inadequate Experimental Results**: The paper's experimental section reveals suboptimal performance in predicting key electronic properties of crystals, such as formation energy. Notably, the proposed method performs poorly in comparison to other methods listed in the results table. Furthermore, the paper lacks a comparison with the state-of-the-art method coGN [1] at Matbench, which is a significant oversight.

Author response:
While this is true, it still performs well on other properties in comparison to other similar graph-based models. Additionally, our model outperforms the other Transformer model we compare to (and the only one on Matbench) in 5/6 properties.

While it is true that line graph-based methods such as ALIGNN and coGN perform with higher accuracy, they also incur a significantly higher computational cost. In the crystal graph construction proposed by CGCNN and used in ALIGNN, if a crystal contains $n$ atoms in the unit cell, the resulting graph will also have $n$ vertices. Each vertex will have $k$ directed edges assuming a large enough cutoff radius (which is suggested to prevent isolated portions of the graph). The line graph for this graph would then contain $k$ vertices and $\frac{n}{2} (k^2 - k)$ edges. coGN adds an additional line graph to account for dihedrals. For a crystal with just 4 atoms in the unit cell, the line graph used in ALIGNN for the angles with $k=12$ would have 48 nodes and 264 edges. The dihedral line graph would contain 264 vertices and 17,424 edges. The crystals in the materials project dataset on average contain $\sim 24$ atoms in the unit cell. In contrast, the approach we propose decreases the size of the representation by grouping similar atoms together. We have mentioned how the collapse tolerance affects the cardinality of the input set in response to another reviewer, we include it here as well. The size of our representation compared to the number of atoms in the unit cell at a collapse tolerance of 1e-4 is as follows:

Phonon Peak: 0.464

Formation energy: 0.444

Band Gap: 0.437

Bulk Modulus: 0.414

Shear Modulus: 0.414

Refractive Index: 0.434

 where on average if a dataset has 24 atoms in the unit cell, our representation's average size would be roughly  $24 \times 0.464 \approx 11$ (in the case of phonon peak).

Reviewer comment (Novelty):
**Lack of Novelty in k-Nearest-Neighbor Construction**: The paper does not sufficiently differentiate its approach from k-nearest-neighbor graph construction of message-passing methods. Common methods for material prediction, such as CGCNN [3] and ALIGNN [4], also consider both atomic properties and distances, raising questions about what exactly is the authors’ method beyond those message-passing methods with k-nearest-neighbor graph construction.

Author response:
Graph construction in other work often needs the use of a cutoff radius to prevent ambiguous graph construction between neighbors with the same distance. This is not needed in our case as we only use the distance and have no reference to the neighbor's atomic type. This radius and the value of $k$ are often not justified in their selection and are instead based on trial and error. We propose a heuristic for finding an adequate value of $k$ based on the theoretical results of generic completeness for the PDD. This is mentioned in the experiments in the appendix, but it seems pertinent so it will be moved to the main body.

Graph-based methods also use edges to indicate between which atoms the distances lie. This provides additional structure to the data but at the cost of being discontinuous under atomic perturbations. Small changes in the position of atoms can change the set of atoms in the set of $k$-NN. If this occurs edges will shift, changing the graph's structure. In our representation we use the same information but in a different way. Despite not using any references to neighboring atoms to indicate which distances correspond to which neighboring atoms, our model still performs on par or better than similarly informed graph-based models on several properties. Additionally, because only distances to the neighbors are used, perturbations will change the distances in a continuous manner.

Reviewer comment (Experiments):
**Omission of Citations**: The authors don't include important baselines coGN [1] and PotNet [2].

Author response:
These two will be added to the review of related work and the previous rationale in the first point will be included in the section which discusses the selection of models to compare our model to.

Reviewer comment (Theory):
A key property of the PDD is its Lipschitz continuity, making it robust to perturbations in the positions. The way the authors use the $k$-NN distances with the hard collapse, then treating the rows as separate tokens, loses this property. Currently, however, the authors are suggesting that the continuity of the PDD is a benefit to their method. The authors should clarify that.

Author response:
During the collapse of any two rows, the newly produced row is created taking the average of the two rows. This way even if the collapse tolerance is greater than zero rows are not selected in an either/or manner. This detail was left out, thank you for pointing it out. It will be added to the updated manuscript.

Reviewer comment (Evaluation):
Could the authors comment on how often the rows of the PDD are collapsed in practice, so how much it matters that the used tokens are aggregates, rather than individual atoms?

Author response:
This is dependent on the dataset and the collapse tolerance used. For a collapse tolerance of 1e-4 (the same as used in the main Materials Project results) the number of grouped tokens in the representation over the number of atoms in the unit cell, averaged for each dataset is:

Phonon Peak: 0.464

Formation energy: 0.444

Band Gap: 0.437

Bulk Modulus: 0.414

Shear Modulus: 0.414

Refractive Index: 0.434

We will add information to the article to show how this changes depending on other collapse tolerances.