DiffAlign

Abstract

Graph diffusion models, dominant in graph generative modeling, remain underexplored for graph-to-graph translation tasks like chemical reaction prediction. We demonstrate that standard permutation equivariant denoisers face fundamental limitations in these tasks due to their inability to break symmetries in noisy inputs. To address this, we propose aligning input and target graphs to break input symmetries while preserving permutation equivariance in non-matching graph portions. Using retrosynthesis (i.e., the task of predicting precursors for synthesis of a given target molecule) as our application domain, we show how alignment dramatically improves discrete diffusion model performance from $5\%$ to a SOTA-matching $54.7\%$ top-1 accuracy.

Background

TODO: define graphs + diffusion model.

Equivariance and Self-Symmetry

It has been noted in the literature [CITE] that permutation equivariant functions struggle to map between a self-symmetrical input space to a less self-symmetrical output space. To confirm this observation, we train two models to translate a ring molecular structure into a less-self symmetrical structure. The first model is an equivariant GNN (GraphTransformer) and the second is a non-equivariant GNN (MPNN, achieved by adding unique positional encodings to the input features of each node). We share the samples generated by the two models below and you can run the models yourself in this notebook.

What this means for Diffusion Models

Diffusion models are expected to transform a noisy input into a clean output by iteratively removing noise from the input. However, if the input is self-symmetrical, the model will struggle to produce a non-symmetrical output. To showcase this effect, we design a simple experiment where the goal is to copy a graph structure. In a single diffusion step, the effect we get is the same as noticed in the previous section. However, if we run the diffusion process for more steps, the model will start to produce more and more different outputs. It can even recover the source graph eventually.

What happens, mathematically?

It turns out that the equivariant denoiser, in an attempt to balance the requirement to break symmetries while maintaining equivariance, will eventually learn a distribution equal to the marginal distribution of the input labels. Formally, we write... (or refer to the paper for the theorem + proof?)

Solution: Aligned Equivariance

So how can we help the model break self-symmetries while maintaining equivariance? We can use node identifiers to match the source and target graphs! Assume you have a source graph $G$ and a target graph $G'$. If you know some connection between the nodes of $G$ and $G'$, you can use that to align the two graphs. For instance, if the graphs represent two configurations of the same graph (e.g. evolution of a graph over time), you can assume that a portion of the nodes in the graph remain intact while the rest are rearranged/added/removed. In retrosynthesis, we can identify where nodes in the source graph (products) are in the target graph (reactants) through atom-mappping. We can then tell the denoiser somehow which nodes are paired through atom-mapping. We will see in the following section how we can do this.

How to ensure alignment?

How can we tell a denoiser that specific nodes are paired? We explore three different methods: Note that the methods can be combined to strengthen the alignment signal. We show that aligned equivariant denoiser remain equivariant to the non-paired nodes in the generated graph.

Results on retrosynthesis

We show that our method achieves a SOTA-matching $54.7\%$ top-1 accuracy, compared to a $5\%$ accuracy without alignment.

What's next? (downstream applications)

Inpainting + inference guidance. Other features of diffusion: distillation, etc.

BibTeX

@proceedings{DiffAlign2025Laabid,
              title={Equivariant Denoisers Cannot Copy Graphs: Align your Graph Diffusion Models},
              author={Laabid, Najwa and Rissanen, Severi and Heinonen, Markus and Solin, Arno and Garg, Vikas},
              booktitle={International Conference on Learning Representations (ICLR)},
              year={2025},
              url={https://openreview.net/forum?id=onIro14tHv}}

Equivariant Denoisers Cannot Copy Graphs: Align your Graph Diffusion Models

We use node identifiers to align source and target graphs, thus breaking the self-symmetries in the input space while preserving permutation equivariance in non-matching graph portions.