Hierarchical GFlownet for Crystal Structure Generation

W1: The presentation of the proposed method lacks important details. For instance, the paper does not mention the initialization method. To improve understanding, it would be beneficial to provide a detailed algorithm outlining the generation process.

A1: We provided the initialization as well as the sampling process in Algorithm 1.

W2: "The proposed method does not consider translation invariance. In Section 4.1, fractional coordinates are directly used as inputs for the GCN model, which violates the translation invariance of 3D crystals. This issue has been well-studied in previous predictive [1] and generative [2] methods, and its omission is a notable weakness.

A2: The atomss list only mantain a reference atom of crystallographic orbit to make sampling process light-weight. Because it is not a complete crystal structure, translation invariance is not required. At the end of the trajectory where the complete crystal structure is generated to get the formation energy, the graph representation as well as the formation energy model is SE(3) invariant.

W1: The paper lacks clarity in explaining the motivation behind the proposed method and its components, making it challenging to discern the specific advantages or necessity of employing GFlowNet in crystal material generation. The introduction of physics-informed reward functions also lacks motivation. There are totally four new reward functions introduced. While their individual effectiveness is showcased through an ablation study, without a robust theoretical underpinning or discussion compared with existing approaches, it is not convincing enough. Does any existing work try similar stuff already? I could possibly raise some random ideas of some physic-related reward functions but how to prove the authors' proposed set of functions are optimal?

A1: On the physical motivation of the minimum bonds information: it was obtained by querying the entire Materials Project database, and computing a bond dictionary for each structure using the function $get\_bonds()$ defined in the OgStructure (https://github.com/sheriftawfikabbas/oganesson/blob/main/oganesson/ogstructure.py). The bond distance between two atoms can fall in a range, but will not go below a specific value for any given structure, otherwise the atoms will experience a repulsive force. The large variability of materials in the Materials Project database is a basis for using a statistical minimum of available bond data for two atoms as a meaningful distance between those two atoms. Such minimum was also used in Ref. [https://www.nature.com/articles/s41524-023-00987-9] for a similar purpose. On the physical motivation of the density parameter, we have already explained in the manuscript the reason for considering the density of the generated structure. Quoting our text from the manuscript: “Because the bond distance preference term applies a strict penalty for structures violating the minimum distance constraint, the generative model tends to generate a structure with long distances between atoms pairs with only a few neighbor atoms. This leads the model to generate gas with low density rather than solid-state density. On the other hand, structures with very high density (i.e. larger than 10) are unlikely to be realistic as the units are crowded with atoms causing a very high formation energy.”

W2: The writing of the paper feels hurried and overloaded with technical details, at the expense of coherence and contextual depth. The introduction is full of many technical details but lacks many discussions related to current literature, what the exact problem is which existing methods either do not aim to tackle or cannot deal with. There is no illustrative figure until page 4., hindering a clear comprehension of the proposed components. A high-level overview and an initial presentation of the experimental evaluations and aims could mitigate these issues and enhance the paper's overall accessibility.

A2: As discussed in the introduction, in this work, we aim to tackle the problem of the complexity associated with the large search space (in this case is the large material space) by proposing a new generative model termed Hierarchical Generative Flow Networks.

The paper’s experimental section does not convincingly demonstrate the superiority of the proposed method. Observing Tables 1 and 2, it’s evident that the proposed methodology struggles to outperform baseline models consistently across various evaluation metrics. This modest performance raises questions about the fundamental contributions and practical viability of the proposed approach. A clearer demonstration of the proposed method's unique advantages, or a clearer justification of its design choices, is crucial. Without such improvements, the paper’s contribution remains uncertain and inadequately supported.

W3: The method's unique advantage is the ability to incorporate constraints into the sampling process and soft constraint to rewards function.

Thank you reviewer for your very detailed comments. It helps us to improve our manuscript alot. We really appreciate your comments regarding our manuscript. We hope we can answer your concerns.

W1: In Figure 1, a flowchart.

A1: Thank you for pointing this out. The wording in the Fig. 1 is indeed misleading. It should be worded "The current state". The figure is to explain how an action of choosing an atom reference position can be translated into multiple positions with symmetry operation.

W2: Sampling process, Gaussian distribution.

A2: Each step in the trajectory samples one single atom coordinates $x_0$. Then the sampled atom is added to the list of sampled atoms $A_{sampled}$. $A_{sampled}$ only stores the one reference atom $x_0$ of the rystallographic orbit (Eq. 1). At the end of the trajectory, we apply the symmetry operation given by the sampled space group to have the full list of atoms $A_{crystal}$.

During the trajectory sampling, we only use the graph representation from the sampled atoms list $A_{sampled}$ as the policy network input. There are two reasons we chose this implementation design. First, with the graph representation is permutation invariant. Second, if we apply symmetry operation on the sampled atoms list $A_{sampled}$ and use the full crystal structure atom list $A_{crystal}$, the number of atoms of the crystal structure can easily reach to hundred or thousand of atoms due to high multiplicity. The computation of a crystal graph of a thousand atoms is very computationally expensive. Creating a new crystal graph at every steps of the trajectory is not preferred.

W3: Continuous GFlownet assumptions, negative fraction coordinate

A3: We also change the claim to "As the original GFlowNets only work on the discrete space, we follow the recent work on continuous GFlownet (Lahlou et al., 2023) to work on the continuous space of the atoms’ coordinates and lattice parameters."

It is true that sampling from Gaussian distribution may result in the negative fraction coordinate. Therefore, we cap the sampled action to $[0,1]$ to ensure the validity of the action.

W4: Other aspects are the selection of lattice parameters, the criteria used in the evaluation ("A structure and its optimized structure are matched if their atoms translation and angle are within certain thresholds" (emphasis mine), the criteria for diversity in the evaluation, the definition of a mode in Section 5.2, etc.

A4: We have provided more information regarding sampling lattice parameters, criteria for stability experiments, and definition of a mode as following and added to Appendix section.

W5: Prediction model

A5: We add the reference of the prediction model in Sec. 4.3 to A.2.2 . We set the cut-off to $-10$ eV/atom as any structure having lower formation energy than that is more likely to be invalid. It is very difficult to tell exactly how often it occurs as the prediction model is a black box function. We can not tell exactly which type of structure causes this. In our own experience, it rarely happens, less than 10 times out of 100000 sampled crystal structures.

W6: The lattice parameters used to illustrate the sampling process in Figure 1 (a = 4, b = 6, c = 4) do not seem to match the drawing.

A6: We have adjusted the Figure to match the number.

W7: "HGFlowNets generalize GFlowNets" claim

A7: We removed this claim as this requires further experiments on different tasks to fully support this claim.

W8: "more meaningful actions" term

A8: The term "more meaningful actions" means that the actions are designed to have a more direct connection the target properties. The connection in the design is provided by prior knowledge, for example, perovskite structure class with conductivity. As discussed in the third paragraph in the Introduction section, the crystal structure class can associate with the properties that we are optimizing. However, to generate the crystal structure, we need to sample the lattice parameter and the atoms' coordinates. Instead of only relying on very complex interactions between atoms using only coordinates to optimize properties, we can choose simpler actions but can directly associate with target properties. In this paper, we limit our experiment to only space group. In future applications, we can extend to more sophisticated higher-level tasks such as choosing composition to associate with crystal structure neutral charge.

W9: "discrete concepts" term

A9: The term "discrete concepts" means that the actions or states are directly involved in the sampling process and can not be omitted. For example in the crystal structure generation task, we cannot have a structure without lattice parameters or atoms' type and coordinates. Meanwhile, the action of choosing a spacegroup is optional but can greatly assist the sampling process. Another example is in the control task. The action of moving the robot's joint is necessary and directly linked to the robot's movement. On the other hand, the action "move left", "move right", or "move to the door" is more at an abstract level.

W10: The paper contains a number of typos or errors in the description of the methods.

A10: Thank you for pointing out these errors. We have corrected those accordingly.

W11: Soft constraint is not optimal, hard constraint during the sampling process.

A11: Our proposed method in crystal structure generation incorporate both hard constaints and soft constraint. The hard constraint is the task we performed in the experiments. The generated crystal must have one of three alkali metals Li, Na, and K with coupled with other light elements Be, B, C, N, O, Si, P, S and Cl. To generates the required material, the hard constraint is choosing the one of three alkali and must coupled with other light elements. which is imposed during the atom's element sampling. The soft constraints are imposed in the reward functions.

W12: PGCGM results

A12: PGCGM performs the post-generation step to refine the crystal structure. To make it fair, in all methods, we only evaluate the generated structure without any further refinement or post-processing. As we pointed out in the 1st paragraph of Sec 5.1, PGCGM faces the common problem of sampling from data-induced distribution without any refinement, which is low structure validity and atoms are very close to each other thus leading to high formation energy.

W13: concurrent work

A13: Thank you for introducing us to the concurrent work. Our work is different as we sample the complete crystal structure while other work only samples crystal structure parameters without atom coordinates. The formation energy is only provided by a newly trained prediction function without any coordinates.

Thank you reviewer for your valuable time and consideration. We really appreciate your comments regarding our manuscript. We hope we can answer your questions.

W1: The experimental section seems quite weak. The method is tested on one task and one dataset. There is only one baseline model and the original GFlownet is used as a baseline.

A1: In the experiment, our goal is to compare the proposed method with the representative methods of different deep learning approaches introduced in the related works section. So in this case, we choose PGCGM for the sampling approach, GFlownet for iterative generation approach.

W2: There are many tasks such as the generation of MOF or Zeolites that are very well suited for the evaluation of this extension of GFlownets. There is a large number of methods that form the state of the art on these tasks that would be great candidates for baselines.

A2: Thank you very much for the task suggestion. Both MOF and Zeolites generation task generates crystal structure based on specific requirements. MOF generates structures with crystalline porous materials with high surface area and tunable pore sizes. Zeolites mainly consist of silicon, aluminium, oxygen with formula as $M^{n+}_{1/n}(AlO_2)^{-}(SiO_2)_x\cdot y H_2O$. In our experiment section, we demonstrate the flexibility based on our battery material discovery task, generating structure with defined formula $M_x N_y A_z$ where A is one of three akali metals Na, K, and Li, M and N are light elements Be, B, C, N, O, Si, P, S and Cl. All these three tasks can be used interchange to demonstrate the iterative crystal structure generation.

Q1: Given that the relaxation results in such large changes in the lattice structure and the atom positions, how much of a role does the lower level (in the hierarchy) policy actually play a role in the generation process?

A1: Structure relaxation is a common step in crystal structure discovery. The relaxation process may result in small changes or large changes in the lattice parameters and atoms' position. If the lower level (in the hierarchy) policy can learn to optimize the reward function, then the change to the lattice parameters and atoms' position should be small. Then we can say that the generated structure from the framework is close to optimal. As a result, it will reduce the time run for DFT optimization.

Q2: Can the crystal be defined only by the atom positions without the lattice parameters? Even though these two data structures are definitely coupled, does the model generally treat them independently?

A2: For the lattice and atoms position, we follow the previous works using the lattice parameter and atoms' fraction coordinate to define a crystal structure. When the policy network takes actions of choosing lattice parameters and atoms' positions, it also has to consider the distance between atoms due to the bond distance preferences term in the reward function. Therefore, the policy needs to treat them dependently.