Efficient symmetry-aware materials generation via hierarchical generative flow networks

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, outside the training space, 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 with 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. To benchmark our approach, we first develop a novel, non-hierarchical GFlowNet for complete crystal structure generation. We then introduce SHAFT, a more expressive hierarchical model that leverages its architecture and increased model capacity to more efficiently explore the materials space. We demonstrate that SHAFT significantly outperforms the flat GFlowNet baseline and other state-of-the-art methods like CDVAE and DiffCSP in generating valid, stable, and diverse crystal structures.