Property-aware Reinforcement Learning with Retrieval Enhancement for Controllable 3D Molecule Generation

Zhoujun Ye, Liangjie Zhao, Weizhi Zhang, Philip S. Yu, Yiwei Fu

Keywords: controllable molecule generation, language model, property-guided retrieval enhancement, property-aware reinforcement learning

Abstract: This paper studies the problem of controllable 3D molecule generation, which aims to design 3D molecules that satisfy given conditions. Previous methods usually incorporate the condition tokens into language models, and reconstruct molecules from the generated tokens. Despite their progress, their performance remains unsatisfactory due to the neglect of the condition during the generation process. To address this limitation, we propose a novel approach named Property-aware Reinforcement Learning with Retrieval Enhancement (POETIC) for controllable 3D molecule generation. To be specific, POETIC first tokenizes 3D molecular structures and leverages a language model (LM) for molecular generation. More importantly, it retrieves relevant samples with similar properties from an external database, which are used as prefixes to enhance generation quality. Furthermore, we pre-train a prediction model to identify the molecular properties, which in turn provides property-aware rewards for evaluation. These rewards guide reinforcement learning to optimize the LM. Extensive experiments on benchmark datasets validate the effectiveness of the proposed POETIC in comparison with state-of-the-art approaches. The source code is available at https://anonymous.4open.science/r/POETIC-BEA3.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 7709