Efficient Symmetry-Aware Materials Generation via Hierarchical Generative Flow Networks

TL;DR

The Symmetry-aware Hierarchical Architecture for Flow-based Traversal (SHAFT) is proposed, a novel generative model employing a hierarchical exploration strategy to efficiently exploit the symmetry of the materials space to generate crystal structures given desired properties.

Abstract

Discovering new solid-state materials requires rapidly exploring the vast space of crystal structures and locating stable regions. Generating stable materials with desired properties and compositions is extremely difficult as we search for very small isolated pockets in the exponentially many possibilities, considering elements from the periodic table and their 3D arrangements in crystal lattices. Materials discovery necessitates both optimized solution structures and diversity in the generated material structures. Existing methods struggle to explore large material spaces and generate diverse samples with desired properties and requirements. We propose the Symmetry-aware Hierarchical Architecture for Flow-based Traversal (SHAFT), a novel generative model employing a hierarchical exploration strategy to efficiently exploit the symmetry of the materials space to generate crystal structures given desired properties. In particular, our model decomposes the exponentially large materials space into a hierarchy of subspaces consisting of symmetric space groups, lattice parameters, and atoms. We demonstrate that SHAFT significantly outperforms state-of-the-art iterative generative methods, such as Generative Flow Networks (GFlowNets) and Crystal Diffusion Variational AutoEncoders (CDVAE), in crystal structure generation tasks, achieving higher validity, diversity, and stability of generated structures optimized for target properties and requirements.

Figures (9)

• Figure 1: (a) Hierarchical crystal structure state. The space group level provides the set of symmetry operations for atoms' positions and lattice parameters constraints. (b) An example of applying the hierarchical state space. The current state has one Oxygen atom at position (0, 0, 0), lattice parameters $a=4,b=6,c=4,\alpha=60^\circ, \beta=80^\circ, \gamma=60^\circ$, and $P1$ spacegroup. The action of choosing space group P4(2)/mmc provides a symmetry operation to generate another Oxygen atom at position (0, 0, 0.5). The lattice parameter constraints reduce the unit cell's length search space from $\mathbb{R}^3$ to $\mathbb{R}$ and make the unit cell's angles constant at $90^\circ$
• Figure 2: Hierarchical policy for crystal structure state. First, the crystal structure graph state $s$ is decomposed into space group state $s_{sg}$ and atom-lattice state $s_{al}$. Then the transition probability $P(s'_{sg}|s_{sg},s_{al})$ in Eq. \ref{['eq:transist_prob']} is given by the space group policy network $\theta_{sg}$. The transition probability $P(s'_{al}|s_{sg},s_{al},s'_{sg})$ in Eq. \ref{['eq:transist_prob']} is given by the atom-lattice policy network $\theta_{al}$.
• Figure 3: Examples of generated crystal structures and the corresponding structure optimized by M3GNet framework Chen2022ATable.
• Figure 4: (a) Comparison of SHAFT and GFlowNet in exploring crystal modes using 3 steps. A mode is defined as a valid crystal structure with negative formation energy. A step is an action of choosing one atom in the spacegroup-lattice-atom hierarchical state space. (b) The average reward of crystal structure sampled by GFlownet and SHAFT in each epoch.
• Figure 5: The distribution of energy above the hull of DFT optimized structures. Most of the generated are stable ($E_{hull}=0 eV/atom$) or metastable $E_{hull}<0.05 eV/atom$. We can achieve low $E_{hull}$ without optimizing directly.