Alignment is Key for Applying Diffusion Models to Retrosynthesis

2024-05-27

arXiv.org

AI Summary

Standard permutation-equivariant denoisers suffer from limited expressivity in graph-to-graph translation tasks—such as retrosynthetic prediction—due to their inability to break the symmetry inherent in noisy inputs. This work reframes retrosynthesis as a conditional graph generation problem and introduces **alignment equivariance**, a novel symmetry constraint that replaces permutation equivariance with atom-mapping consistency, thereby overcoming the symmetry bottleneck. This design enables flexible post-training conditional control and unifies template-based and template-free paradigms. Built upon a graph diffusion framework, our approach integrates a customized denoising network with an atom-mapping alignment mechanism. On USPTO-50k, it achieves a state-of-the-art 54.7% top-1 accuracy, supports few-step sampling, and enables multi-step retrosynthetic planning.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Graph diffusion models lack exploration for graph-to-graph translation tasks
• Standard equivariant denoisers fail due to symmetry preservation in noisy inputs
• Alignment improves performance in retrosynthesis prediction tasks significantly

Innovation

• Align input and target graphs
• Break input symmetries effectively
• Preserve permutation equivariance non-matching parts