RUM: Reinforcement Unlearning for 3D Molecular Flow Matching Models

Generative models, particularly flow matching models, have shown strong potential for large-scale 3D molecular design. However, models trained on existing molecular datasets may also learn to generate molecules with toxic or other undesired properties, limiting their reliability and practical utility. Retraining a model from scratch using only retained molecules can address this issue but is computationally expensive for large 3D molecular datasets. Model unlearning provides a more efficient alternative, but faces two key challenges: removing undesired generation while preserving the original retained distribution, and preventing the forgotten capability from being easily recovered through subsequent fine-tuning. In this study, we propose a reinforcement unlearning framework for pretrained 3D molecular flow matching models. Reinforcement forgetting suppresses undesired generation, while Flow Matching replay on retained molecules sampled from the pretrained model maintains the retained distribution. We further simulate relearning through fine-tuning on undesired molecules and train the model to resist this recovery. Experiments show that RUM achieves effective forgetting while maintaining retained generation and improving resistance to subsequent fine-tuning.