Equivariant Denoisers Cannot Copy Graphs: Align Your Graph Diffusion Models

Authors: Najwa Laabid, Severi Rissanen, Markus Heinonen, Arno Solin, Vikas Garg
Primary area: learning on graphs and other geometries & topologies
Keywords: Graph Diffusion, Discrete Diffusion, Equivariance, Symmetries

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 \emph{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. Code is available at https://github.com/Aalto-QuML/DiffAlign.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730634999624

thread_type:
official_review

<thread_object>
creator: reviewer_hvws

<content>
summary:
This paper proposes an aligned permutation equivariant denoiser for a discrete diffusion model to address the graph translation problem. To align the initial state and target, several techniques, including nodal positional encoding, skip connections, and input alignment, are introduced. The model was validated on retrosynthesis, guided generation, and inpainting tasks.

strengths:
Aligning initial and target to generate efficiently makes sense and the results are promising compared to unaligned version.

weaknesses:
Since alignment itself is an empirical approach that has been used previously, the contribution may be limited to explaining why alignment is necessary specifically in a discrete diffusion model with an absorbing state as the prior. When compared to RetroBridge in the retrosynthesis task, it is difficult to find a significant performance gap between the proposed model and RetroBridge. In the cases of guided diffusion and inpainting, only a few figures are provided without numerical comparisons, making precise validation difficult.

questions:
1. The paper assumes that the prior distribution is absorbing state. Does the proposed model still effective even on the uniform distribution prior?

2. What happens if the node-mapping matrix has uncertainty and is not a binary matrix?

3. How does the performance compare when using only direct skip connections without positional encoding?

4. How does the performance differ from graph translation methods in Related Works, such as G2gt?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730647506358

thread_type:
official_review

<thread_object>
creator: reviewer_gptc

<content>
summary:
The paper presents a novel method to perform graph-to-graph translation using graph diffusion models. The authors show why equivariant models are inherently limited in their ability to perform this task when nodes from the prior distribution need to be aligned with nodes from the output. To that purpose, they demonstrate that an ideal equivariant denoiser cannot correctly copy a graph. Then, they propose several ways to break equivariance for the nodes that should be aligned between the input and the output, and empirically validate their method on the retrosynthesis task.

strengths:
- To the best of my knowledge, the topic of breaking symmetries to perform graph-to-graph translation has not yet been addressed in the graph diffusion literature.
- The problem of using equivariant models is clearly identified and backed by theoritical evidences.
- The paper proposed several simple ways to solve this problem.

weaknesses:
- Why using the aborbing state distribution and not the marginal distributions introduced in DiGress (Vignac et al., 2022) ?

- I would avoid statements like "it is easy to see" l304 or "Clearly" l242 and provide instead clear motivations and explanations

- Some claims or design choices lack motivation or justification :
- For example, the claims in section 3.2 "a model that is capable of copying graphs from one side of the reaction
to the other should also be extendable to modulations of this task". Do you mean that it's a prerequisite or that is implies that the model has the ability to peform "the modoluations". Overall, I fell that this subsection conveys very little information.
- The proposed aligning methods are insufficiently motivated.

questions:
- What is the difference between $D$ and $\tilde p_{\theta}$ ?
- How is $\psi$ generated in section 3.5 ?
- How is Y passed to the denoiser ? I feel it's not clearly outlined anywhere except in the section "Aligning Y in the input"
- I don't get how you build the dataset for retrosynthesis. If it is standard benchmark, could you please provide more explanations or point out to a reference
</content>
</thread_object>

</thread>
<thread>
creation date:
1730651935511

thread_type:
official_review

<thread_object>
creator: reviewer_gt4f

<content>
summary:
The central objective of this paper is to design a discrete denoising diffusion model for graph-to-graph translation, where the task is to train a model to predict a target graph given a source graph and, perhaps surprisingly, a matrix of node mappings between them. The authors use [1] as their design template, though the state transition matrix of the forward process relies on the absorbing state formulation from [2]. In this model, the backward process relies on the denoiser network realized by an equivariant transformer architecture, which, on its own, can take only the target graph as input. The authors' main challenge is augmenting the denoiser by also taking the source graph as input. The authors demonstrate theoretically and empirically that a simple augmentation of the denoiser's input by the source graph is not sufficient even for simply copying graphs, claiming that the equivariance of the denoiser is the main reason for this issue. Therefore, the authors propose to partially relax the equivariance by further extending the denoiser's input with the node mapping matrix, proposing the aligned permutation equivariance. The key idea is to design the underlying transformer architecture such that if the source and target graphs are permuted (by two separate permutation matrices) and the node mapping matrix is permuted accordingly, then the output of the denoiser is equivariant only under the permutations of the target graph. In other words, the output remains sensitive only to the permutation of the target graph. The authors compare the performance of their model with many baseline methods on the task of chemical retrosynthesis, achieving state-of-the-art top-1 accuracy and mean reciprocal rank on the USPTO-50k dataset. Finally, they also show the suitability of their model in the context of the guided and conditional generation of molecular graphs.

strengths:
The paper presents a comprehensive set of experiments that demonstrate the benefits and the practical utility of the proposed model in the context of guided and conditional generation of molecular graphs. The model achieves state-of-the-art results in chemical retrosynthesis.

The paper concerns an interesting problem, and its central message is important to the community.

Nice and clear illustrations accompany the ideas introduced in the paper.

weaknesses:
The writing deserves improvement. Certain parts of the paper need to be communicated more clearly. The notation sometimes needs to be clarified. These concerns are detailed in the following comments.

Major:

The content from the start of Section 3 up to Section 3.2 qualifies more like a background section. The material is heavily inspired by [1]. Please note that [1] also uses conditioning, though it is not presented in the equations in such an explicit way. It takes more work to recognize the details behind the authors' contributions in the current form. Can the authors clearly delineate which parts are background and which are novel contributions?

In line 124 the authors define $X=P^{Y\rightarrow X}Y$ as $(X^N,X^E)=(P^{Y\rightarrow X}Y^N,P^{Y\rightarrow X}Y^E(P^{Y\rightarrow X})^T)$, and in line 151 define $D_\theta(X,Y)=(D_\theta(X,Y)^N,D_\theta(X,Y)^E)$, saying that $D_\theta(X,Y)^N$ and $D_\theta(X,Y)^E$ output a probability vector for each node and edge, respectively. Now, in line 107, the authors say that $X^N$ and $X^E$ are one-hot encoded (again, it is not suitable to use $\mathbb{R}$ in lines 104 and 105 to define this), then how is it possible that $D_\theta(X,Y)=P^{Y\rightarrow X}Y$? There is no mention that the output of $D_\theta(X,Y)$ is implicitly one-hot encoded. This confusion repeats throughout the text, and clear definitions should stabilize it. The authors should

- clarify when and where $X^N$ and $X^E$ are one-hot encoded or real-valued matrices/tensors,
- explicitly define the output space of $D_\theta(X,Y)$,
- explain how the equality $D_\theta(X,Y)=P^{Y\rightarrow X}Y$ is possible given the different encodings,
- use consistent notation throughout the paper for the spaces these variables belong to.

Minor:

The introduction should better motivate why graph-to-graph translation is important in practical applications.

Line 50: *``a solution aligned''* -> *``a solution: aligned''*

Figure 1: *``w.r.t to''* -> *``w.r.t.''*

In Section 3.3, the authors use phrases such as *``should be possible''*, *``should have learned''*, *``should be the same''*. Why do not use a more certain language? Can the authors please clarify what is or is not true instead of something that should be?

The definition of graph-to-graph translation task deserves improvements. Specifically, if the authors define that the node feature matrix and edge feature tensor belong to the set of reals, why do they redefine them as one-hot vectors in line 107? If they are one-hot, then they are only a subset of reals.

There should be $Y$ on the left side of (2). Consequently, how does this density differ from (3)?

Line 153: *``such that we have a probability vector for each node and edge.''* Probability is not the whole set of reals. It only ranges from 0 to 1.

Line 217: The use of *`respectively'* does not fit the sentence's construction.

Line 243: $\mathbf{X}=\mathbf{P}^{\mathbf{Y}\rightarrow \mathbf{X}}Y$ -> $\mathbf{X}=\mathbf{P}^{\mathbf{Y}\rightarrow \mathbf{X}}\mathbf{Y}$

Line 323: What is DPS?

Line 374: What is MRR? It is first defined in line 402.

The explanation of retrosynthesis in lines 337 and 358 is nearly identical.

Line 337: *``3 parts"* -> *``three parts''*

Table A3 should be in the main text (e.g., Figure 8 can be moved to the appendix).

Figure 3 demonstrates the case for $X_T$, yet there is $D_\theta(X_t,Y)$.

questions:
Do the authors need the node mapping matrix for this task?

Can the authors explain why the left part of (6) should hold (in light of the above concerns with the output of the denoising network $D_\theta$ and the one-hot encoding of $X$)?

Why do the authors present Theorem 1 for fixed t=T? Is the result for $X_T$ not rather obvious (by definition)?

If the authors use $P^{Y\rightarrow X}$ and $Y$ as input into the denoiser, then these two terms already form the full information about $X$. The authors seem to train the denoiser that always has direct access to $X\_0$ through $P^{Y\rightarrow X}Y$. The consequence of this is that $\tilde{p}\_\theta$ in (3) contains $P^{Y\rightarrow X}$ in the condition, i.e., $\tilde{p}\_\theta(X\_0|X\_t,Y,P^{Y\rightarrow X})=\tilde{p}_\theta(X\_0|X\_t,X\_0)$. This design choice is unsound. Please, can the authors clarify?

Can the authors please provide a more comprehensive explanation of the following comment (starting in line 227): *``This makes it impossible for the model ... some random permutation of $Y$.''*?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732060492996

thread_type:
official_review

<thread_object>
creator: reviewer_uwes

<content>
summary:
This paper presents a novel approach to graph-to-graph translation using diffusion models. The authors propose methods to address the inherent limitations of equivariant models in aligning nodes between input and output graphs. The paper introduces several strategies for breaking equivariance and empirically validates their approach on retrosynthesis tasks. While the paper is innovative in addressing this important problem, certain theoretical justifications and experimental evaluations could be strengthened.

strengths:
- The paper addresses a novel and relevant problem in graph diffusion literature, particularly in breaking symmetries for graph-to-graph translation, which has not been explored in depth.
- The idea of using equivariant models and tackling their limitations is well-motivated, and the theoretical underpinning is clear and compelling.
- The alignment techniques proposed for mapping nodes between input and output graphs are intuitive and demonstrate promising preliminary results, especially compared to unaligned versions.

weaknesses:
- The justification for certain design choices, such as using the absorbing state distribution instead of marginal distributions (as in DiGress), is insufficiently explained. A more thorough motivation or comparison would be beneficial.
- Some key claims, particularly in Section 3.2, are vague or lack detailed explanation. For example, the statement regarding the model's capability to handle graph transformations requires clearer elaboration.
- The experimental validation, while promising, lacks sufficient quantitative comparisons to existing methods (e.g., RetroBridge). More thorough benchmarking and analysis would strengthen the empirical contribution of the paper.

questions:
Please refer to  the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732186680024

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We thank the reviewers for their detailed feedback. Below we give a summarized response to a few recurrent points in the reviews.

**Choice of transition matrix**   While other transitions, like uniform and marginal [1] are equally possible, the absorbing state model [2] is somewhat simpler and has been empirically found to outperform other transitions [2,3,4,5]. In our early experiments, we found that the uniform transition performed worse while marginal and absorbing performed similarly. Since the main content of the paper is not dependent on the particular choice of transition, we decided to focus the experiments on the somewhat simpler absorbing state model.

**Clarifying the format of the denoiser $D_\theta(X,Y)$ output** Reviewers Gt4F and gPTC expressed concerns about the notation of $D_\theta(X,Y)$, its relation to the one-hot vector $P^{Y\to X}Y$, and to the probabilities $\tilde p(X_0|X_t,Y)$. We have now clarified the connections in the text when introducing $D_\theta(X,Y)$ and other relevant modelling objects. Specifically, we consider the space of one-hot vectors of dimension $K$ as $\text{vert}(\Delta^{K-1})$, where vert(.) denotes the vertices of the probability simplex $\Delta^{K-1}$ with $K-1$ degrees of freedom. The output of $D_\theta(X,Y)$ as well as the one-hot-encoded vectors $X$ and $Y$ are vertices of this probability simplex, making the equality $D_\theta(X,Y) = P^{Y\to X}Y$ sensible mathematically.

**Delineating contributions vs background** We reorganized our sections to differentiate more easily between our contributions vs the framework laid out by [1]. We also added Figure 2 to the introduction to summarize our contributions at a glance. 
* *Sec. 5 The Theoretical Limitations of Equivariant Denoisers*: Copying graphs as a case-study on the limitations of equivariance.(Fig. 4), The optimal permutation equivariant denoiser (Theorem 1).
* *Sec. 6 Solution Aligned Equivariant Denoisers*: Relaxing permutation equivariance through alignment (Fig. 6a), Aligned denoisers induce aligned permutation invariant distributions (Theorem 2).
* *Sec. 7 Suitable Alignment Methods*: Skip-connection, positional encodings, and input-alignment (Fig. A9).
* *Sec. 9 Experiments*: a) Proof of concept through simple graph copying tasks (Fig. 5), b)
Application: Retrosynthesis (Table 2), c) Guided-generation for retrosynthesis (Fig. 8a), d) Ibuprofen synthesis through inpainting (Fig. 7).

We also moved all the background information, including the framework of [1] and formal definitions of permutation equivariance and invariance as they pertain to diffusion model, to a dedicated section (Section 2, "Background"). 

**Results compared to RetroBridge** Following RetroBridge's evaluation procedure, which does not take into account atom charges nor stereochemistry, we outperform their results significantly with an over 5% increase in top-1 accuracy. We share the detailed top-k accuracies and mean reciprocal rank below, and include them in Appendix D, Table A.4. 

|Model|k=1|k=3|k=5|k=10|MRR|
|-----|---|---|---|----|---|
|RetroBridge (T=500)|50.8|74.1|**80.6**|**85.6**|0.622|
|Ours-uncharged (T=100)|**56.3**|**75.2**|80.4|84.0|**0.658**|

Our model achieves this with 100 denoising steps, as opposed to 500 by Retrobridge, and achieves higher top-1 scores with 5 steps than Retrobridge with 100 steps (see their Figure 5 compared to our Figure 5b). We believe that this improvement is due to us using multiple methods for alignment, which were inspired by our theoretical framework.

**Performance compared to other graph-translation baselines** Prompted by reviewer HVWS, we higlight our model results compared to the seminal G2G [2] graph translation-based retrosynthesis method. Top-k accuracies and the mean reciprocal rank are listed below:

| Method | k = 1 | k = 3 | k = 5 | k = 10 | MRR |
|--------|-------|-------|-------|--------|-----|
| G2G | 48.9 | 67.6 | 72.5 | 75.5 | 0.582 |
| DiffAlign-PE+skip |**54.7** | **73.3**| **77.8** | **81.1** | **0.639** |

We note that this is a semi-template-based (or synthon-based) model with strong assumptions in retrosynthesis, namely the existence of a reaction center which is first broken to identify synthons. Such assumptions do not apply to all chemical reactions unfortunately, and thus G2G[2] is inherently limited by its inability to model reactions like the one shown in our Figure A11, Appendix D.

**References**

[1] Vignac et al., "DiGress: Discrete Denoising diffusion for graph generation", ICLR 2023

[2] Austin et al., "Structured Denoising Diffusion Models in Discrete State-Spaces", NeurIPS 2021

[3] Lou et al., "Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution", ICML 2024

[4] Sahoo et al., "Simple and Effective Masked Diffusion Language Models", NeurIPS 2024

[5] Shi et al., "Simplified and Generalized Masked Diffusion for Discrete Data", NeurIPS 2024
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work tackles an interesting problem setting of graph-to-graph translation, utilising discrete denoising diffusion models as well as equivariant structures. As the reviewers reaches a consensus that the paper is clearly written and contribution is substantial, with support of interesting chemical synthesis applications, it would be recommended for the work to appear in ICLR.