Alignment is Key for Applying Diffusion Models to Retrosynthesis

Najwa Laabid,Severi Rissanen,Markus Heinonen,Arno Solin,Vikas Garg

2024-05-28

Abstract:Retrosynthesis, the task of identifying precursors for a given molecule, can be naturally framed as a conditional graph generation task. Diffusion models are a particularly promising modelling approach, enabling post-hoc conditioning and trading off quality for speed during generation. We show mathematically that permutation equivariant denoisers severely limit the expressiveness of graph diffusion models and thus their adaptation to retrosynthesis. To address this limitation, we relax the equivariance requirement such that it only applies to aligned permutations of the conditioning and the generated graphs obtained through atom mapping. Our new denoiser achieves the highest top-$1$ accuracy ($54.7$\%) across template-free and template-based methods on USPTO-50k. We also demonstrate the ability for flexible post-training conditioning and good sample quality with small diffusion step counts, highlighting the potential for interactive applications and additional controls for multi-step planning.

Machine Learning,Quantitative Methods

What problem does this paper attempt to address?

This paper mainly discusses how to effectively apply diffusion models in retrosynthesis tasks. Retrosynthesis is the process of determining the precursors of a given molecule, which is crucial for chemical synthesis planning. The study found that standard graph diffusion models have limitations in retrosynthesis due to the use of permutation equivariant denoisers, which restrict their expressive power in retrosynthesis.

To address this issue, the paper proposes the concept of "aligned permutation equivariance". By leveraging atom mapping information, the permutations between reactants and products are aligned, maintaining equivariance only when atom mapping remains unchanged. This approach improves the accuracy of the model, and the newly proposed denoiser achieves the highest accuracy (54.7%) in both template-free and template-based methods on the USPTO-50k dataset.

Furthermore, the paper demonstrates the model's ability to incorporate conditional information after training, such as completing and optimizing reaction properties, making the model suitable for interactive applications and multi-step planning. The paper emphasizes the importance of correctly applying diffusion models in graph conditional problems and proposes various alignment implementation methods that significantly improve the generation quality of directly applying graph diffusion models.

In summary, this paper addresses the limitations of using standard graph diffusion models in retrosynthesis tasks and proposes a new aligned permutation equivariant denoiser, enhancing the flexibility, accuracy, and adaptability of the model.