Reinforcement learning with formation energy feedback for material diffusion models

Generative models are emerging as foundation tools for the discovery of new materials with remarkable efficiency. Existing works introduce physical constraints during the generation process of diffusion models to improve the quality of the generated crystals. However, it is difficult to accurately capture the distribution of stable crystal material structures, given the complex periodic crystal structure and the limited available crystal material data, even with the incorporation of symmetries and other domain-specific knowledge. Thus, these models still struggle to achieve a high success rate in producing stable crystal materials. To further improve the stability of generative crystal materials, we propose a novel fine-tuning framework RLFEF. We formulate the material diffusion process as a Markov Decision Process with formation energy serving as rewards. Moreover, we prove that optimizing the expected return in reinforcement learning is equivalent to applying policy gradient updates to a diffusion model. Additionally, we prove that the fine-tuned model adheres to the unique symmetry of crystal materials. Extensive experiments are conducted on three real-world datasets. The results show that our model achieves state-of-the-art performance on most tasks related to property optimization, ab initio generation, crystal structure prediction, and material generation.