Equivariant Denoisers Cannot Copy Graphs: Align Your Graph Diffusion Models

## Abstract

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 1 1 1 accuracy ( 54.7 54.7 54.7 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.

## 1 Introduction

Single-step retrosynthesis plays a pivotal role in chemistry, as it focuses on identifying appropriate precursors for a given target compound which can then be chained to form complex synthesis plans []. Approaches to retrosynthesis have steadily moved from relying on explicit expert-curated rules (template-based []), and using simplifying assumptions like a single reaction center (synthon-based [, ]), to ultimately adopting purely data-driven machine-learning models (template-free [, ]). The data-driven paradigm has the least assumptions about the reactions modeled, and can thus scale well to novel compounds.
Although these methods are often not as accurate on benchmarks as approaches handcrafted with chemical inductive biases, developing more powerful and easily extendable template-free models is crucial for overcoming the inherent limitations of hard-coded chemical rules and heuristic decomposition strategies.

Recently, diffusion models have shown impressive results in conditional molecular generation [], suggesting their potential usefulness in probabilistic retrosynthetic planning. Diffusion-based generation offers several unique advantages. A known one is trading off generation speed and quality on a per sample basis []. Another advantage is the possibility to add additional conditions even post-training [, , ], potentially allowing us to specify substructures in the generated reactants [], enforce desirable properties [] like synthesizability, or sample from nuanced regions of the learned distribution [] in a flexible way.
These advantages would translate into a versatile and adjustable single-step model, which in turn could directly be utilized by a multi-step retrosynthetic planner (e.g., spending less time on small products, allowing more exploration of the search tree in a given time [], or optimizing the atom economy [] of the reactions). Furthermore, a simple adaptation of diffusion models to retrosynthesis holds the promise of all future advances in the rapidly expanding methodology of generative modeling directly being opened up to the retrosynthesis domain as well. To enable all of these advantages, however, we need to fully understand how to properly apply diffusion models in a graph-conditional setting.

In this paper, we show how to apply diffusion models to graph-conditioanl probllems in the most straightforward way possible. We first characterize a serious limitation of neural network expressivity that occurs when using standard permutation equivariant denoiser architectures. To address the issue, we introduce *aligned permutation equivariance*, where the reactants and products are aligned with atom-mapping information, and equivariance to permutations only holds when the alignment through atom mapping remains intact (see Fig. 1 with an example sketch of alignments and atom mappings). We propose multiple methods to achieve this alignment with standard GNN-based denoisers, vastly outperforming the permutation equivariant model. Our method achieves the highest top-$1$ accuracy across template-based and template-free methods, and the highest top-$k>1$ and MRR among synthon-based and template-free methods. Our model is thus competing with models using explicit chemical knowledge while enabling the scalability of template-free methods. We also showcase how to apply post-training conditions with our approach, allowing for interactive applications with inpainting and optimizing properties of the reactions.

We summarize our contributions as follows:

- •
Fundamental:
We demonstrate that permutation equivariance is a severe limitation when applying standard graph diffusion models to retrosynthesis by characterizing the theoretically optimal denoiser output, and propose a new class of *aligned* permutation equivariant denoisers that circumvents the issue with atom mapping. We show that they induce generalized permutation invariant distributions that are tailored to graph translation problems with alignment information, such as retrosynthesis.
- •
Methodological:
We propose multiple methods to turn an equivariant GNN into an aligned equivariant model, pushing the generative quality of a direct application of graph diffusion models from a top-$k$ accuracy range of 5–15% to state-of-the-art on USPTO-50k among non-template-based methods.
- •
Topical:
We show how to apply different post-training conditions, like inpainting or controlling the synthesisability and atom economy of the generated reactions, allowing for interactive applications and additional controls for multi-step retrosynthesis.

## 2 Related Work

#### Retrosynthesis

There are three main types of retrosynthesis models []. Template-based models depend on the availability and quality of hard-coded chemical rules [, ]. Synthon-based models are limited by their definition of a reaction center, which does not necessarily hold for complex reactions [, , ]. Template-free methods are the most scalable since they do not use any chemical assumptions in their design but perform suboptimally compared to template-based methods on benchmarks like UPSTO-50k [, ]. Efforts to bridge the gap between the template-based and template-free paradigms include methods investigating pretraining [, ]. Recently, Igashov et al. [] proposed a Schrödinger-bridge inspired method for retrosynthesis sharing similarities with diffusion models. In another concurrent work [], the authors propose to use a synthon-based hierarchical diffusion model for multi-stage generation. In contrast, we focus on a direct application of conditional graph diffusion models to retrosynthesis, and the importance of alignment when doing so.

#### Product-reactant node alignment in retrosynthesis

The idea of aligning the product and reactant atoms based on atom mapping has been explored by Zhong et al. []. They show that matching the order of atoms between the product and reactant sequences boosts the performance of sequence-to-sequence translation models applied to retrosynthesis. We share their observation regarding the chemical inductive bias brought by the alignment, but the similarity between our methods ends here. Instead of using the atom mapping directly to align the atoms on the reactants and product side as Zhong et al. [] did, we explore different methods to enforce this alignment on molecular graphs to remedy the limitations of permutation equivariant denoisers.

#### Permutation equivariance in graph diffusion

A fundamental inductive bias with neural networks operating on graphs is *permutation equivariance* [], which ensures that the network only uses topological connectivity information and does not process the order of nodes in an input graph $\mathbf{X}$. It is widely used in graph diffusion models to parameterize the reverse process [, , , ]. One motivation is that, in diffusion models, permutation equivariant neural nets induce permutation invariant distributions [], meaning that different permutations of the same graph have the same probabilities under the model.
Recently, Yan et al. [] showed that relaxing permutation equivariance in graph diffusion using absolute positional encodings can improve performance empirically.

## 3 Methods

Figure: Figure 2: An overview of our approach to retrosynthetic graph diffusion. We adopt *absorbing state diffusion* [], and use atom mapping information (highlighted with the matching colours) as an input to the neural network denoiser to align the product condition with the reactant data.
Refer to caption: x4.png

Consider a database of $N_{\mathrm{obs}}$ known chemical reactions $\mathcal{D}=\{(\mathbf{X}_{n},\mathbf{Y}_{n},\mathbf{P}^{\mathbf{Y}\to\mathbf{
X}}_{n})\}_{n=1}^{N_{\mathrm{obs}}}$, where $\mathbf{X}_{n}$ are reactants, $\mathbf{Y}_{n}$ are products and $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}_{n}$ are matrices defining *atom mappings* between products and reactants. The *retrosynthesis* task is: given that the data is sampled from an unknown distribution $p(\mathbf{X},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$, predict the valid reactant molecules $\mathbf{X}\sim p(\mathbf{X}\,|\,\mathbf{Y})$ for a given product molecule $\mathbf{Y}$. In our data encoding (also illustrated in Fig. A6 in App. A)
we model the reactants and products as molecular graphs with nodes and edges representing atoms and bonds, respectively. Formally, we define the reactant graph $\mathbf{X}$ as a tuple $(\mathbf{X}^{\mathcal{N}},\mathbf{X}^{E})$ of a node feature matrix $\mathbf{X}^{\mathcal{N}}\in\mathbf{R}^{N_{X}\times K_{a}}$ and an edge feature matrix $\mathbf{X}^{E}\in\mathbf{R}^{N_{X}\times N_{X}\times K_{b}}$, s.t. $K_{a}$ and $K_{b}$ are atom and edge feature dimensions respectively. The product graph $\mathbf{Y}$ is similarly defined as $(\mathbf{Y}^{\mathcal{N}},\mathbf{Y}^{E})$ with a potentially different number of nodes $N_{Y}$. The node features are one-hot vectors over atom types $(\perp,\textrm{C},\textrm{N},\textrm{O},\textrm{P},\ldots)$ and edges are one-hot vectors over bond types $(\perp,1,2,3)$.
Both include an empty value $\perp$ that represents a missing node or edge. Usually, $N_{X}\geq N_{Y}$, as the data is defined to only include the main product molecule.

The atom mapping matrices $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\in\mathbf{R}^{N_{X}\times N_{Y}}$ give us additional information about how atoms that ended up in the product were reconfigured in the reaction, and are defined such that $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}_{i,j}=1$ if the $i$^th atom of the reactant corresponds with the $j$^th atom of the product, and zero otherwise. Thus, $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{\mathcal{N}}$ equals $\mathbf{X}^{\mathcal{N}}$ with the non-atom-mapped atoms zeroed out. Correspondingly, $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{E}(\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}})^{\top}$ equals $\mathbf{X}^{E}$ with edges to non-atom-mapped atoms zeroed out.

We propose to model the problem with a discrete variant of the standard denoising diffusion probabilistic model (DDPM, []) where the reactants $\mathbf{X}_{0}$ are diffused in a Markov chain $q(\mathbf{X}_{t}\,|\,\mathbf{X}_{t-1})$ for $t\in[1,T]$, and we learn how to transform noise back to reactants $p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$, while conditioning on the fixed product $\mathbf{Y}$. Specifically, we adapt discrete diffusion models [, , ] to graphs, following Vignac et al. [].

We assume a Markov *forward* process

$$ $\textstyle q(\mathbf{X}_{t+1}\,|\,\mathbf{X}_{t})=\prod^{N_{X}}_{i=1}q(\mathbf {X}^{\mathcal{N},i}_{t+1}\,|\,\mathbf{X}^{\mathcal{N},i}_{t})\prod_{i,j=1}^{N_ {X}}q(\mathbf{X}^{E,ij}_{t+1}\,|\,\mathbf{X}^{E,ij}_{t}),$ (1) $$

to diffuse the reactant to noise, and a reverse process

$$ $\textstyle p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t})=\prod_{i=1}^{N_{X}} p_{\theta}(\mathbf{X}^{\mathcal{N},i}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y}) \prod_{i,j}^{N_{X}}p_{\theta}(\mathbf{X}^{E,ij}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y}),$ (2) $$

defining our generative model. Here, $\theta$ represents neural network parameters. Note we implicitly always have time conditioning $p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y},t)$, but we drop the explicit $t$ for notational convenience. We will also condition on the atom mapping $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$ in Sec. 3.5, but we will not include it in the notation until then. The full generative distribution is $p_{\theta}(\mathbf{X}_{0:T}\,|\,\mathbf{Y})=p(\mathbf{X}_{T})\prod_{t=1}^{T}p_
{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$, where $p(\mathbf{X}_{T})$ is a predefined prior such that $p(\mathbf{X}_{T})=q(\mathbf{X}_{T}\,|\,\mathbf{X}_{0})$. Following [, ], we use the neural network specifically to predict ground truth labels from noised samples, so that the neural net directly outputs a distribution $\tilde{p}_{\theta}\big{(}\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y}\big{)}$.
The reverse process is then given by

$$ $\textstyle p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})=\sum_{ \mathbf{X}_{0}}q\big{(}\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{X}_{0}\big{ )}\tilde{p}_{\theta}\big{(}\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y}\big{)}.$ (3) $$

Throughout the paper, we denote the direct output of the neural network as $D_{\theta}(\mathbf{X}_{t},\mathbf{Y})=(D_{\theta}(\mathbf{X}_{t},\mathbf{Y})^{
\mathcal{N}},D_{\theta}(\mathbf{X}_{t},\mathbf{Y})^{E})$, where $D_{\theta}(\mathbf{X}_{t},\mathbf{Y})^{\mathcal{N}}\in\mathbf{R}^{N_{X}\times K
_{a}}$ and $D_{\theta}(\mathbf{X}_{t},\mathbf{Y})^{E}\in\mathbf{R}^{N_{X}\times N_{X}
\times K_{b}}$, s.t. we have a probability vector for each node and edge.

The single-step transition for nodes (resp. for edges) is defined with a transition matrix $\mathbf{Q}_{t}^{\mathcal{N}}$ as

$$ $q(\mathbf{X}_{t}^{\mathcal{N},i}\,|\,\mathbf{X}_{t-1}^{\mathcal{N},i})=\textrm {Cat}(\mathbf{X}_{t}^{\mathcal{N},i};\mathbf{p}=\mathbf{X}_{t-1}^{\mathcal{N}, i}\mathbf{Q}_{t}^{\mathcal{N}}).$ (4) $$

To define $\mathbf{Q}_{t}^{\mathcal{N}}$ (resp. $\mathbf{Q}_{t}^{E}$), we adopt the absorbing-state formulation from Austin et al. [], where nodes and edges gradually transfer to the *absorbing state*, defined as the empty state $\perp$. Formally, $\mathbf{Q}_{t}=(1-\beta_{t})I+\beta_{t}\mathbbm{1}e_{\perp}^{\top}$, where $\beta_{t}$ defines the *diffusion schedule* and $e_{\perp}$ is one-hot on the absorbing state. Then, the marginal $q(\mathbf{X}_{t}\,|\,\mathbf{X}_{0})$ and conditional posterior $q(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{X}_{0})$ also have a closed form for the absorbing state transitions.
The prior $p(\mathbf{X}_{T})$ is correspondingly chosen to be a delta distribution at a graph with no edges and nodes set to the $\perp$ state. The noise schedule $\beta_{t}$ is defined using the mutual information criterion proposed in Austin et al. [].

We use the cross-entropy loss, as discussed in Austin et al. [] and Vignac et al. []:

$$ $-\mathbb{E}_{q(\mathbf{X}_{0},\mathbf{Y})q(t)q(\mathbf{X}_{t}\,|\,\mathbf{X}_{ 0})}[\log\tilde{p}_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})],$ (5) $$

where $q(t)$ is a uniform distribution over $t\in\{1\dots T\}$.
Fig. 2 provides an overview of the conditional graph diffusion framework we adopted. For completeness, a comprehensive definition of the model’s components for nodes and edges is given in App. A. The training and sampling procedures with graph diffusion models are presented in Alg. 1 and Alg. 2, along with optional conditioning by $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$, as described in Sec. 3.5.

Figure: Algorithm 1 Loss calculation

### 3.1 Permutation Equivariance and Invariance

Permutation equivariance for the denoiser is defined as $D_{\theta}(\mathbf{P}\mathbf{X})=\mathbf{P}D_{\theta}(\mathbf{X})$ where $\mathbf{P}$ is an arbitrary permutation matrix, and $\mathbf{P}\mathbf{X}=(\mathbf{P}\mathbf{X}^{\mathcal{N}},\mathbf{P}\mathbf{X}^
{E}\mathbf{P}^{\top})$. This makes single-step reverse transitions equivariant as well, in the sense that $p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t})=p_{\theta}(\mathbf{P}\mathbf{X
}_{t-1}\,|\,\mathbf{P}\mathbf{X}_{t})$. For our conditional setting, permutation equivariance can be written as $D_{\theta}(\mathbf{P}\mathbf{X},\mathbf{Y})=\mathbf{P}D_{\theta}(\mathbf{X},
\mathbf{Y})$. The prior in graph diffusion models is also usually permutation invariant so that $p(\mathbf{P}\mathbf{X}_{T})=p(\mathbf{X}_{T})$ [, , ].

### 3.2 The Identity Reaction as a Subproblem of Retrosynthesis

Chemical reactions often induce limited changes to the precursor molecules []. Previous approaches to retrosynthesis capitalize on this observation and propose similarity [], local graph edits [], or product and reactants alignment [] as ways to induce chemical inductive biases.
Motivated by the same observation, in this section, we consider a data set $\mathcal{D}=\{(\mathbf{X}_{n},\mathbf{Y}_{n},\mathbf{P}_{n}^{\mathbf{Y}\to
\mathbf{X}})\}_{n=1}^{N_{\mathrm{obs}}}$, dubbed the ‘identity reaction data’, where for all data points $\mathbf{X}_{n}=\mathbf{P}_{n}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}_{n}$.
In other words, both sides of the reaction are equivalent, up to some permutation, as defined in the atom mapping matrix $\mathbf{P}_{n}^{\mathbf{Y}\to\mathbf{X}}$. While this may seem like a very simplistic scenario, we can use it to reason about the expressiveness of our denoiser neural network. In addition, a model that is capable of copying molecules from one side of the reaction to the other should also be extendable to modulations of this task, e.g., splitting bonds or adding new atoms to the reactant side. Thus, finding a model that can solve this task easily should also be in a good position for more realistic retrosynthesis tasks.

### 3.3 Permutation Equivariant Denoisers Fail at the Identity Reaction

Here, we show that standard graph diffusion models are very ineffective at modelling identity reactions.
In essence, while the copy-paste task should be solvable by always predicting $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}
\mathbf{Y}$, we show that permutation equivariant denoisers $D_{\theta}$ are constrained to output ‘mean’ solutions, such that the outputs for all nodes and all edges are equal in the early stages of generation.

Let us consider the case of a highly noisy input, $t=T$. Given the definition of the identity reaction task and permutation equivariance, it should be possible for the following two to hold at the same time:

$$ $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{P}^{\mathbf{Y}\to\mathbf{X}} \mathbf{Y}\quad\text{and}\quad D_{\theta}(\mathbf{R}\mathbf{X}_{T},\mathbf{Y}) =\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y},$ (6) $$

where $\mathbf{R}$ is an arbitrary permutation matrix. We have a problem: Since $\mathbf{X}_{T}$ does not contain any information about $\mathbf{Y}$, the model should have learned to ignore the $\mathbf{X}$ input. Intuitively, this means that having $\mathbf{R}\mathbf{X}_{T}$ or $\mathbf{X}_{T}$ as input should not make a difference for the output, which combined with Eq. 6 results in the conclusion that the output should be the same for all nodes and edges. The following theorem characterizes this rigorously for the denoiser at high noise levels $t=T$, showing that the optimal solution is one where the value of each row is, in fact, the marginal of atom/bond labels respectively.

###### Theorem 1 .

(The optimal permutation equivariant denoiser)
Let $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})$ be permutation equivariant s.t. $D_{\theta}(\mathbf{P}\mathbf{X}_{T},\mathbf{Y})=\mathbf{P}D_{\theta}(\mathbf{X
}_{T},\mathbf{Y})$, and let $q(\mathbf{X}_{T})$ be permutation invariant. The optimal solution with respect to the cross-entropy loss with the identity reaction data is, for all nodes $i$ and $j$

$$ $\begin{cases}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{\mathcal{N}}_{i,:}=\hat{ \mathbf{y}}^{\mathcal{N}},\hat{\mathbf{y}}^{\mathcal{N}}_{k}=\sum_{i}\mathbf{Y }_{i,k}/\sum_{i,k}\mathbf{Y}_{i,k},\\ D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{E}_{i,j,:}=\hat{\mathbf{y}}^{E},\hat{ \mathbf{y}}^{E}_{k}=\sum_{i,j}\mathbf{Y}_{i,j,k}/\sum_{i,j,k}\mathbf{Y}_{i,j,k },\end{cases}$ (7) $$

where $\hat{\mathbf{y}}^{\mathcal{N}}_{k}$ and $\hat{\mathbf{y}}^{E}_{k}$ are the marginal distributions of node and edge values in $\mathbf{Y}$.

The proof is given in Sec. B.1, and the Theorem is illustrated in Fig. 3. This makes it impossible for the model to solve the task in a single step with $T=1$, clearly limiting the performance at this basic task. With enough steps, it is still possible, albeit slow. This is because the neural network can, in principle, start matching substructures of the partially generated $\mathbf{X}_{t}$ to the condition $\mathbf{Y}$, and ultimately generate some random permutation of $\mathbf{Y}$.

Figure: Figure 3: Comparing the optimal permutation equivariant denoiser to an aligned denoiser. As per Theorem 1, the permutation equivariant model (*top*) outputs the marginal distribution over node types whereas the aligned model (*bottom*) reconstructs the correct reactant nodes in a specific permutation $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$.
Refer to caption: x5.png

###### Theorem 1 .

### 3.4 Aligned Permutation Equivariance

Given the failure of permutation equivariant denoisers to effectively model the identity reaction, we propose to relax the permutation equivariance assumption just enough to give the network the capacity to easily copy the product conditioning $\mathbf{Y}$ to the denoiser output. The key observation we make is that it is enough to have permutation equivariance in the sense that *if we permute $\mathbf{X}$ and/or $\mathbf{Y}$, and accordingly permute the atom mapping $P^{\mathbf{Y}\to\mathbf{X}}$ such that the matching between $\mathbf{Y}$ and $\mathbf{X}$ remains, the model output should be the same*. We call this *aligned permutation equivariance*. Formally, we use the atom-mapping permutation matrix $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$ as an input to the denoiser and consider denoisers that satisfy the following constraint: $D_{\theta}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{
\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=\mathbf{R}D_{\theta}(\mathbf{X},
\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$, where $\mathbf{R}$ and $\mathbf{Q}$ are permutation matrices of shapes $(N_{\mathbf{X}}\times N_{\mathbf{X}})$ and $(N_{\mathbf{Y}}\times N_{\mathbf{Y}})$, respectively.

Clearly, with $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$ and $\mathbf{Y}$ as an input, an unconstrained denoiser can output the ground-truth permutation $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}Y$ for the identity reaction task. Restricting the function to the aligned equivariant class does not clash with this, as can be seen by writing out the optimal solution and the equivariance condition for a permuted input $\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}}\mathbf{Q}^{\top}$:

$$ $\displaystyle(\text{\small Identity reaction})\quad D_{\theta}(\mathbf{R} \mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}} \mathbf{Q}^{\top})$ $\displaystyle=\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top} \mathbf{Q}\mathbf{Y}=\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y},$ (8) $\displaystyle(\text{\lx@cref{creftype~refnum}{eq:aligned_equivariance}})\quad D _{\theta}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{ \mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=\mathbf{R}D_{\theta}(\mathbf{X},\mathbf{Y},\mathbf{P}^{\mathbf{Y }\to\mathbf{X}})=\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}.$ (9) $$

Note that the permutation invariance of $p_{\theta}(\mathbf{X}_{0})$ does not apply as it does for fully permutation equivariant models. However, we can show that a generalized form of distribution invariance holds.

###### Theorem 2 .

(Aligned denoisers induce aligned permutation invariant distributions)
If the denoiser function $D_{\theta}$ has the aligned equivariance property and the prior $p(\mathbf{X}_{T})$ is permutation invariant, then the generative distribution $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$ has the corresponding property for any permutation matrices $\mathbf{R}$ and $\mathbf{Q}$:

$$ $p_{\theta}(\mathbf{R}\mathbf{X}_{0}\,|\,\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf {P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=p_{\theta}(\mathbf{X}_{0}\,|\, \mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}).$ (10) $$

A proof is given in Sec. B.2. Informally, the theorem states that the chemical reaction has the same probability for all isomorphisms of the reactant and product graphs, as long as the atom mapping is not reassigned to different atoms.

Thus, during training, we can use the permutations present in the data, and be sure that the model generalizes to other permutations. Initially, during sampling, we only have access to $\mathbf{Y}$ without atom mapping information. The theorem ensures we can assign atom mappings arbitrarily during sampling and still obtain effectively the same distribution over reactant graphs. To see this, we note that all atom mapping matrices can be characterized by a base permutation matrix $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$ left multiplied by different permutations $\mathbf{R}$. Applying the theorem: $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}})=p_{\theta}(\mathbf{R}^{-1}\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{P
}^{\mathbf{Y}\to\mathbf{X}})$, where $\mathbf{R}^{-1}$ is just another permutation matrix, shows that the distribution equals $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$, up to some permutation. In fact, by sampling the atom mapping matrix randomly, the distribution $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y})=\sum_{\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}}}p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}})p(\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$ becomes invariant to any permutation of $\mathbf{X}_{0}$ and $\mathbf{Y}$, although this is not strictly necessary in our context. Figure 3 showcases the denoising output of an aligned equivariant model with a given $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$.

###### Theorem 2 .

### 3.5 Methods for Alignment in GNNs

We have established how aligning the product $\mathbf{Y}$ and the reactants $\mathbf{X}$ is a key component for the success of the denoising model. In this section, we discuss multiple simple methods to turn a regular, permutation equivariant denoiser into an aligned permutation equivariant denoiser. Here, we present the implementation details, and in Sec. B.3 we prove for each method that they belong to the class of aligned equivariant models. Fig. A7 visualizes the methods. Note that our use of atom mapping here is not a constraint: it is merely an additional input which allows the denoiser to use the structure of the conditional graph efficiently, and as such we benefit even from a partially correct atom mapping.

#### Atom-mapped positional encodings

We consider uniquely identifying pairs of atoms matched via the atom-mapping matrix in both $\mathbf{X}_{t}$ and $\mathbf{Y}$ by adding a positional encoding vector to each unique atom pair. In practice, we generate a set of distinct vectors $\boldsymbol{\varphi}\in\mathbf{R}^{N_{Y}\times d_{\boldsymbol{\varphi}}}$ for each of the product atoms, and then map those vectors to the corresponding inputs for the reactant nodes $\mathbf{X}_{t}$, s.t. $\boldsymbol{\varphi}=g(\mathbf{Y})$ and $D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})=f_{
\theta}([\mathbf{X}^{\mathcal{N}}_{t}\;\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}
\boldsymbol{\varphi}],\mathbf{X}^{E}_{t},[\mathbf{Y}^{\mathcal{N}}\;
\boldsymbol{\varphi}],\mathbf{Y}^{E})$,
where $g$ is a function that generates the encodings based on $\mathbf{Y}$, $f_{\theta}$ is the neural network that takes as inputs the augmented node features and regular edge features and $[\mathbf{X}^{\mathcal{N}}_{t}\;\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\boldsymbol
{\varphi}]\in\mathbf{R}^{N_{X}\times(K_{a}+d_{\boldsymbol{\varphi}})}$ corresponds to concatenation along the feature dimension. The only change required for the neural network is to increase the initial linear layer size for the node inputs.

#### Direct skip connection

Here we modify the network to include a direct connection from the product to the reactant output: $D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})=
\text{softmax}(f^{\text{logit}}_{\theta}(\mathbf{X}_{t},\mathbf{Y})+\lambda
\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y})$
where $f^{\text{logit}}_{\theta}(\mathbf{X}_{t},\mathbf{Y})$ are the logits at the last layer of the neural network for the nodes and edges, $\lambda$ is a learnable parameter and $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}=(\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}}\mathbf{Y}^{\mathcal{N}},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}
\mathbf{Y}^{E}(\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\top})$. The sum is possible because $\mathbf{Y}$ is in one-hot format and the dimensionalities of the denoiser output and $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}$ are the same. It is easy to see that when $\lambda\to\infty$, then $D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})\to
\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}$ because the direct connection from $\mathbf{Y}$ dominates in the softmax.

#### Aligning 𝐘 𝐘 \mathbf{Y} bold_Y in the input

Instead of aligning $\mathbf{Y}$ to $\mathbf{X}$ using atom mapping and adding it to the *output*, we align $\mathbf{Y}$ in the *input* by concatenating $\mathbf{X}$ and $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}$ along the feature dimension before passing it to the neural network. This presents an opportunity for computational savings: Since $\mathbf{Y}$ is already concatenated to $\mathbf{X}$ along the feature dimension, $\mathbf{Y}$ doesn’t need to be processed as a separate graph anymore. This is useful, e.g., with the graph transformer model, since computation scales quadratically with the total amount of nodes, which is reduced from $(N_{X}+N_{Y})$ to $N_{X}$. Thus, we can drop out the explicit $\mathbf{Y}$ graph entirely, and only process the $\mathbf{X}$ graph augmented with $\mathbf{Y}$: $D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})=f_{
\theta}([\mathbf{X}_{t}^{\mathcal{N}}\;\;\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}
\mathbf{Y}^{\mathcal{N}}],[\mathbf{X}_{t}^{E}\;\;\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}}\mathbf{Y}^{E}(\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\top}])$,
where $[\cdot]$ means concatenation along the feature dimension. In fact, while preparing the manuscript of the paper, we noticed that a similar method was used implicitly in the code of the recent/concurrent work of Igashov et al. [], although atom mapping is not mentioned in the paper and the focus is on the Schrödinger-bridge like probabilistic framework in retrosynthesis.

### 3.6 Adding Post-training Conditioning for Discrete Diffusion Models

One of the advantages of the diffusion framework is its support for flexible conditioning during generation. This means either fixing part of the generated samples (i.e. inpainting []) or enforcing a property of the generated sample as an extra condition (i.e. guided inference []).
We implement inpainting similarly to [], by fixing a portion of the generated reactants to a desired structure.
For property-guided inference, we derive a DPS-like [] method for discrete models where we use the denoiser output to evaluate the likelihood $p(y\,|\,\mathbf{X}_{0})$ of the condition $y$ (e.g., probability of synthesisability) given the reactants $\mathbf{X}_{0}$, and backpropagate to get an adjustment of the backward update step:

$$ $\displaystyle\vspace*{-1em}\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\, \mathbf{X}_{t},\mathbf{Y},y)\propto\nabla_{\mathbf{X}_{t}}\log(\mathbb{E}_{p_{ \theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})}p(y\,|\,\mathbf{X}_{0})) \mathbf{X}_{t-1}{+}\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t} ,\mathbf{Y}).$ (11) $$

More details in App. D. The method can be implemented with a few lines of code and works as long as $p(y\,|\,\mathbf{X}_{0})$ is differentiable. In particular, we experiment with a toy model of synthesisability that is based on the count of atoms in the reaction [], a concept also known as atom economy [].

Figure: (a) Aligned vs unaligned denoiser
Refer to caption: x6.png

## 4 Experiments

#### Data and preprocessing

We use the benchmark data set USPTO-50k for our experiments. We use the same data split as Dai et al. []. Sec. C.1 includes a review of the main subsets of the USPTO data sets used in the literature.
To generate precursor sets of various sizes, we set the reactant blocks to a fixed size of $N_{Y}+d$, where $d$ is a hyperparameter we set empirically to 15. In practice, a number $N_{X}\leq N_{Y}+d$ will be turned into atoms constituting valid molecules, while $(N_{Y}+d)-N_{X}$ of the nodes will be identified as ‘blank nodes’ by the neural network and given a corresponding value. The atom features are one-hot encoded atom labels (28, including 26 formally charged atoms, a node label used as the absorbing state, and a ‘blank’ node label used for the nodes not turned into atoms in the precursor set). The bond features are covalent bond types (single, double, triple and an additional ‘none’ type). We include stereochemistry in the generated samples by utilizing the atom mapping implicitly attached to the samples to transfer the chiral tags and cis/trans isomerism from the product to the reactants. More details on our handling of stereochemistry in Sec. C.3.

#### Neural network

We use the graph transformer architecture introduced by [] and used by []. The network architecture and hyperparameters features are detailed in Sec. C.5.

#### Sampling and ranking

For each product, we generate 100 reactant sets. We remove duplicates from the generated sets per product, and count the times that they appear. We then rank the unique reactants from highest-to-least likely, as judged by the model, using the likelihood lower bound and duplicate counts as a proxy. For details on our ranking scheme, see Sec. C.2.

#### Our models

Based on the types of alignment discussed in Sec. 3.5, we experiment with (1) a permutation equivariant model (*Unaligned*), (2) a model augmented with atom-mapped positional encodings, calculated from the graph Laplacian eigendecomposition (*DiffAlign-PE*), (3) A model with Laplacian positional encodings as well as the skip connection (*DiffAlign-PE+skip*), (4) A model where $\mathbf{X}$ and $\mathbf{Y}$ are aligned at the input as well as in the output (*DiffAlign-input alignment*). We use $T=100$. We trained the models for 400–600 epochs and chose the best checkpoint based on the MRR score with $T=10$ of the validation set. For more details, see Sec. C.5.

#### Quantitative evaluation

We follow previous practices and use the accuracy of obtaining the ground-truth reactants as our main metric. This is measured by top-$k$ accuracy, which counts the number of ground-truth matches in the deduplicated and ranked samples among the top-$k$ generated reactions. However, an issue with top-$k$ accuracy metrics is that most models perform best on a subset of them, and it is thus difficult to determine the best model across all ranks. As a remedy, we also report Mean Reciprocal Rank (MRR, []) as used by Maziarz et al. []. To compare to methods which do not report MRR, we use a method that estimates the MRR from the reported top-$k$ scores by assuming $p(r)$ to be uniformly distributed between the reported values of $k$. More information about MRR and the exact details of the estimation procedure are given in Sec. C.4.

#### Baselines

We choose representative methods from each category for comparison to our models. Methods using pre-training [, ], or a more relaxed evaluation procedure [] are not directly comparable to other baselines including our models. We therefore exclude them from our evaluation. See Sec. C.6 for an overview of these methods and Table A3 for a more comprehensive results table.

### 4.1 Results

#### Denoiser visualization

We begin by visualizing the aligned denoiser and the permutation equivariant denoiser in Fig. 4(a). As predicted by Theorem 1, a sample from the permutation equivariant denoising distribution at high levels of noise $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})$ has no information about the structure of the product. In contrast, an aligned denoiser is able to copy the product structure, and the initial denoising output is of much higher quality.

**Table 1: Top-$k$ accuracy and MRR on the USPTO-50k test data set. Our model outperforms existing template-free baselines on all top-$k$ and MRR. Template-based models use explicit chemical rules and cannot scale beyond the templates available to them. For comparison to models pre-trained on larger data sets, including RSMILES [] and PMSR [], or using non-standard evaluation procedures, see Sec. C.6.**
|  | Method | $\mathbf{k=1}\uparrow$ | $\mathbf{k=3}\uparrow$ | $\mathbf{k=5}\uparrow$ | $\mathbf{k=10}\uparrow$ | $\widehat{\mathrm{\textbf{MRR}}}\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Temp. | Retrosym [] | 37.3 | 54.7 | 63.3 | 74.1 | 0.480 |
| GLN [] | 52.5 | 74.7 | 81.2 | 87.9 | 0.641 |  |
| LocalRetro [] | 52.6 | 76.0 | 84.4 | 90.6 | 0.650 |  |
| Synthon | G2G [] | 48.9 | 67.6 | 72.5 | 75.5 | 0.582 |
| GraphRetro [] | 53.7 | 68.3 | 72.2 | 75.5 | 0.611 |  |
| MEGAN [] | 48.0 | 70.9 | 78.1 | 85.4 | 0.601 |  |
| RetroDiff [] | 52.6 | 71.2 | 81.0 | 83.3 | 0.629 |  |
|  | SCROP [] | 43.7 | 60.0 | 65.2 | 68.7 | 0.521 |
|  | Tied Transformer [] | 47.1 | 67.1 | 73.1 | 76.3 | 0.572 |
|  | Aug. Transformer [] | 48.3 | — | 73.4 | 77.4 | 0.569 |
|  | GTA_aug [] | 51.1 | 67.6 | 74.8 | 81.6 | 0.605 |
|  | Graph2SMILES [] | 52.9 | 66.5 | 70.0 | 72.9 | 0.597 |
|  | Retroformer [] | 52.9 | 68.2 | 72.5 | 76.4 | 0.608 |
|  | DualTF_aug [] | 53.6 | 70.7 | 74.6 | 77.0 | 0.619 |
|  | Unaligned | 4.1 | 6.5 | 7.8 | 9.8 | 0.056 |
|  | DiffAlign-input | 44.1 | 65.9 | 72.2 | 78.7 | 0.554 |
|  | DiffAlign-PE | 49.0 | 70.7 | 76.6 | 81.8 | 0.601 |
| Template-free<br>Ours | DiffAlign-PE+skip | 54.7 | 73.3 | 77.8 | 81.1 | 0.639 |

#### Top- k 𝑘 k italic_k scores and MRR

We report results for the different models and baselines for top-$k$ and the combined MRR score on the USPTO-50k test set in LABEL:tab:50k. The model without alignment performs worst, although it does recover some reactants. The model with input and output alignment (DiffAlign-input) performs much better, but is overshadowed by a model with atom-mapped positional encodings (DiffAlign-PE). Adding the output alignment pushes the top-$k$ up, raising top-$k=1$ to state-of-the-art across all model types, and top-$k>1$ and MRR to state-of-the-art among template-free methods.

#### The effect of sampling steps

While we use a relatively large value of $T$ for our best models, we also highlight that the performance of our model does not degrade significantly even when reducing the count of sampling steps to a fraction of $T=100$. See Fig. 4 for an ablation study with the number of sampling steps for our best model with positional encodings and skip connections. We also compare to a permutation equivariant model and show that the top-$k$ scores go to near zero at 10 steps, while the aligned model sometimes recovers the ground truth even with a single denoising step.

#### Controlled generation

We illustrate examples of inpainting in Fig. 5(b). The substructures highlighted in yellow lead to the model proposing a different set of valid, complementary precursors each time. Property guided generation is illustrated in Fig. 5(a), where we nudge the model towards precursors with lower atom count, a property known as *atom economy* []. See Sec. D.2 and Sec. D.1 for more elaborate case studies using both types of conditioning, and discussion on interpreting atom count as a toy synthesisability model.

Figure: (a) Atom guidance
Refer to caption: x7.png

## 5 Conclusion

In this work, we studied an important aspect of the design space of conditional graph diffusion models: the equivariance of the denoiser. We showed that a permutation equivariant model converges to a ‘mean’ distribution for all components of the graph. We discuss how this result limits the expressivity of the denoising process, which ultimately impedes the performance of graph diffusion models. We propose aligned permutation equivariance to force the model to only consider permutations which maintain the intrinsic alignment between the conditioning and generated graphs. Our aligned denoisers achieve state-of-the-art results among template-free methods, reaching a top-$1$ accuracy beyond that of template-based methods. Furthermore, aligned permutation equivariance opens up the benefits of graph diffusion in retrosynthesis, including flexible post-training conditioning mechanisms, and the ability to adjust the amount of sampling steps during inference freely. These properties could be used, e.g., in interactive applications or to add more controls for multi-step retrosynthesis planners. Our work also opens up all the future advances in the diffusion paradigm itself to retrosynthesis. For instance, for discrete diffusion, ARDM[] or MaskGIT [] types of sampling would be straightforward to incorporate into our model.

A limitation of the method is the requirement of some atom-mapping information, although not fully mapped reactions. The model can handle errors in atom mapping with additional denoising steps, and recent atom-mapping tools are improving in efficiency []. Another limitation is the high computational resource demand, making the model likely slower than template-based models but similar to autoregressive models. Advances in accelerated diffusion sampling methods [, , , , , ] are likely to improve this aspect.

This study enhances graph diffusion models for retrosynthesis, with similar ethical concerns as other chemistry ML models. While improving synthetic route predictions can expedite discoveries, it also risks creating harmful substances []. Responsible usage guidelines are crucial. No specific harmful impacts are noted from this work.

## Appendices

This appendix is organized as follows.
App. A provides additional details on the setup for conditional graph diffusion, including the transition matrices, noise schedule, and data encoding as graphs. App. B presents our theoretical results on aligned permutation equivariance and accompanying proofs. App. C includes additional details to replicate our experimental setup. App. D develops a method to apply arbitrary post-training conditioning with discrete diffusion models, and presents case studies showcasing the usefulness of post-training conditional inference in applications relevant to retrosynthesis. This includes generating samples with desired properties and refining the generation interactively through inpainting.

[appendices]
\printcontents[appendices]l1

## Appendix A Details on Conditional Graph Diffusion

#### Our transition matrices

To define $\mathbf{Q}_{t}^{\mathcal{N}}$ and  $\mathbf{Q}_{t}^{E}$, we adopt the absorbing-state formulation from Austin et al. [], where nodes and edges gradually transfer to the *absorbing state* $\perp$. Formally, we give the generic form of the transition matrix $\mathbf{Q}_{t}$ for node input $\mathbf{X}^{\mathcal{N}}\in\mathbb{R}^{N_{X}\times N_{X}}$ (^1^11The only difference between $\mathbf{Q}_{t}^{\mathcal{N}}$ and $\mathbf{Q}_{t}^{E}$ for the absorbing-state and uniform transitions is the dimensions of $I$, $e$, $\mathbbm{1}$ and the value of $K$. We therefore give a generic form for both and imply choosing the right dimensions for each case.)

$$ $\displaystyle\mathbf{Q}_{t}$ $\displaystyle=(1-\beta_{t})\mathbf{I}+\beta_{t}\mathbbm{1}e_{\perp}^{\top},$ (12) $$

where $\beta_{t}$ defines the *diffusion schedule* and $e_{\perp}$ is one-hot on the absorbing state $\perp$. For completeness, we list the other two common transitions relevant to our application. The first is the uniform transition as proposed by Hoogeboom et al. []

$$ $\displaystyle\mathbf{Q}_{t}=(1-\beta_{t})\mathbf{I}+\beta_{t}\frac{\mathbbm{1} \mathbbm{1}^{\top}}{K}$ (13) $$

where $\beta_{t}$, $I$ are as before and $K$ is the number of element (edge or node) types, i.e. the number of input features for both nodes and edges. Vignac et al. [] also proposed a marginal transition matrix

$$ $\displaystyle\mathbf{Q}_{t}^{\mathcal{N}}=(1-\beta^{t})\mathbf{I}+\beta^{t} \mathbbm{1}\big{(}m^{\mathcal{N}}\big{)}^{\top}\quad\text{and}\quad\mathbf{Q}^ {E}_{t}=(1-\beta^{t})\mathbf{I}+\beta^{t}\mathbbm{1}\big{(}m^{E}\big{)}^{\top}$ (14) $$

which they argued leads to faster convergence. In this case, $m^{\mathcal{N}}\in\mathbb{R}^{K_{a}}$ and $m^{E}\in\mathbb{R}^{K_{b}}$ are row vectors representing the marginal distributions for node and edge types respectively. We tested all three types of transition matrices in early experiments and noted the absorbing state model to be slightly better than the others. The marginal $q(\mathbf{X}_{t}\,|\,\mathbf{X}_{0})$ and conditional posterior $q(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{X}_{0})$ also have a closed form for all of these transition matrices.

#### Noise schedule

We use the mutual information noise schedule proposed by Austin et al. [], which leads to

$$ $\frac{t}{T}=1-\frac{I(\mathbf{X}_{t};\mathbf{X}_{0})}{H(\mathbf{X}_{0})}=\frac {H(\mathbf{X}_{0},\mathbf{X}_{t})-H(\mathbf{X}_{t})}{H(\mathbf{X}_{0})}=\frac{ \sum_{\mathbf{X}_{0},\mathbf{X}_{t}}p(\mathbf{X}_{0})q(\mathbf{X}_{t}\,|\, \mathbf{X}_{0})\log\frac{q(\mathbf{X}_{t}\,|\,\mathbf{X}_{0})}{\sum_{\mathbf{X }^{\prime}_{0}}p(\mathbf{X}^{\prime}_{0})q(\mathbf{X}_{t}\,|\,\mathbf{X}^{ \prime}_{0})}}{\sum_{\mathbf{X}_{0}}p(\mathbf{X}_{0})\log p(\mathbf{X}_{0})}$ (15) $$

For absorbing state diffusion, these equations lead to $\beta_{t}=\frac{1}{T-t+1}$. Similarly, the total transition probability to the absorbing state at time $t$ has a simple form: $q(\mathbf{X}_{t}=\perp\,|\,\mathbf{X}_{0})=\frac{t}{T}$.

#### Forward process posterior

For transition matrices that factorize over dimensions, we have

$$ $\displaystyle q(\mathbf{X}_{t-1,i,:}\,|\,\mathbf{X}_{t,i,:},\mathbf{X}_{0,i,:} )\sim\frac{\mathbf{X}_{t}\mathbf{Q}_{t}^{\top}\cdot\mathbf{X}_{0,i,:}\bar{ \mathbf{Q}}_{t-1}}{\mathbf{X}_{0,i,:}\bar{\mathbf{Q}}_{t}\mathbf{X}^{\top}}$ (16) $$

where $\mathbf{X}_{i,:}$ is the one-hot encoding of i^th node/edge of the graph, in row vector format. The formula outputs the correct distribution

#### Variational lower-bound loss

Diffusion models are commonly trained by minimizing the negative variational lower-bound on the model’s likelihood []. Austin et al. [] discuss the difference between optimizing the ELBO and cross-entropy losses and show that the two losses are equivalent for the absorbing-state transition. We choose to use cross-entropy, similar to Vignac et al. [], due to faster convergence during training. We include the formula for the ELBO in App. A for completeness.

$$ $\displaystyle\mathcal{L}_{\mathrm{vb}}$ $\displaystyle=\mathbb{E}_{\mathbf{X}_{0}\sim q(\mathbf{X}_{0})}\bigg{[} \underbrace{\mathrm{KL}\big{[}q(\mathbf{X}_{T}\,|\,\mathbf{X}_{0})\,\|\,p( \mathbf{X}_{T})\big{]}}_{\mathcal{L}_{T}}$ $\displaystyle\qquad+\underbrace{\sum_{t=2}^{T}\mathbb{E}_{\mathbf{X}_{t}\sim q (\mathbf{X}_{t}\,|\,\mathbf{X}_{0})}\mathrm{KL}\big{[}q(\mathbf{X}_{t-1}\,|\, \mathbf{X}_{t},\mathbf{X}_{0})\,\|\,p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X} _{t})\big{]}}_{\mathcal{L}_{t-1}}$ $\displaystyle\qquad-\underbrace{\mathbb{E}_{\mathbf{X}_{1}\sim q(\mathbf{X}_{1 }\,|\,\mathbf{X}_{0})}\log p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{1})}_{ \mathcal{L}_{0}}\bigg{]}$ (17) $$

We also note that we use this quantity as part of the scoring function mentioned in Sec. 4.

#### Data encoding and atom-mapping

We illustrate our graph encoding using atom-mapping and permutation matrices.

Figure: Figure A6: Illustrating atom mapping-based permutations and the tensor-based encoding of the reaction graph. The colours highlight the atom-mapped atoms and the elements that link them in the atom-mapping matrix $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$. $\mathbf{Y}^{\mathcal{N}}$ and $\mathbf{X}^{\mathcal{N}}$ are the one-hot encoded products and reactants, respectively. $\mathbf{Y}^{E}$ and $\mathbf{X}^{E}$ show the edge tensors, where 1=single bond, 2=double bond, 3=triple bond and 0=no bond. The grey boxes highlight links that connect to atom-mapped atoms in the visualization.
Refer to caption: x9.png

## Appendix B Theoretical Results on Aligned Permutation Equivariance

### B.1 Proof that Permutation Equivariant Denoisers Do Not Implement the Identity Reaction

#### Definitions

Let us consider a data set $\mathcal{D}=\{\mathbf{X}_{n},\mathbf{Y}_{n},\mathbf{P}_{n}^{\mathbf{Y}\to
\mathbf{X}}\}_{n=1}^{N_{obs}}$, where for all data points, $\mathbf{X}_{n}=\mathbf{P}_{n}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}_{n}$, that is, both sides of the reactions are equivalent, up to some permutation, as defined in the atom mapping matrix $\mathbf{P}_{n}^{\mathbf{Y}\to\mathbf{X}}$. It is always possible to preprocess the data such that the rows of $\mathbf{Y}_{n}$ are permuted with $\mathbf{Y}_{n}\leftarrow\mathbf{P}_{n}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}_{n}$ so that the resulting atom mapping between $\mathbf{Y}_{n}$ and $\mathbf{X}_{n}$ is always identity. For simplicity, we assume such a preprocessed data set in this section.

Let us assume that the one-step denoiser probability, $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})$, is parameterized by the neural network $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})\in\mathbb{R}^{N\times K}$ such that the probability factorises for the individual nodes and edges (so there is one output in the network for each node and each edge): $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})=\prod_{i}\sum_{k}
\mathbf{X}^{N}_{0,i,k}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{\mathcal{N}}_{i,k
}\prod_{i,j}\sum_{k}\mathbf{X}_{0,i,j,k}^{E}D_{\theta}(\mathbf{X}_{T},\mathbf{
Y})^{E}_{i,j,k}$.

The ideal denoiser Clearly, the correct one-step denoiser $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{Y}$. This can be shown with the Bayes’ rule by $q(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})=\frac{q(\mathbf{X}_{T}\,|\,
\mathbf{X}_{0},\mathbf{Y})q(\mathbf{X}_{0}\,|\,\mathbf{Y})}{q(\mathbf{X}_{T}\,
|\,\mathbf{Y})}=\frac{q(\mathbf{X}_{T})q(\mathbf{X}_{0}\,|\,\mathbf{Y})}{q(
\mathbf{X}_{T})}=q(\mathbf{X}_{0}\,|\,\mathbf{Y})$. Because one $\mathbf{X}$ always matches with exactly one $\mathbf{Y}$ in the data, this is a delta distribution $q(\mathbf{X}_{0}\,|\,\mathbf{Y})=\prod_{i}\delta_{x_{0,i},y_{i}}$, where we define $x_{0,i}$ and $y_{i}$ as the value of the $i$:th node / edge. It is also easy to see that $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{Y}$ is the optimal solution for the cross-entropy loss:

$$ $\displaystyle-\sum_{(\mathbf{X}_{0},\mathbf{Y})}q(\mathbf{X}_{0})\log p_{ \theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})\propto-\sum_{\mathbf{Y}} \log p_{\theta}(\mathbf{Y}\,|\,\mathbf{X}_{T},\mathbf{Y})$ $\displaystyle\qquad=-\sum_{\mathbf{Y}}\log\left[\prod_{i}\sum_{k}\mathbf{Y}^{N }_{0,i,k}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{\mathcal{N}}_{i,k}\prod_{i,j} \sum_{k}\mathbf{Y}_{0,i,j,k}^{E}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{E}_{i,j ,k}\right].$ (18) $$

All of the sums $\sum_{k}\mathbf{Y}^{N}_{0,i,k}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{\mathcal{
N}}_{i,k}$ and $\sum_{k}\mathbf{Y}_{0,i,j,k}^{E}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{E}_{i,j
,k}$ are maximized for each $i,j$ and $\mathbf{Y}$ if $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{Y}$. In this case, the loss goes to zero.

The following theorem states that if the neural net is permutation equivariant, it will converge to a ‘mean’ solution, where the output for each node and each edge is the marginal distribution of nodes and edges in the conditioning product molecule, instead of the global optimum $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})=\mathbf{Y}$.

###### Theorem 3 .

The optimal permutation equivariant denoiser
Let $D_{\theta}(\mathbf{X}_{T},\mathbf{Y})$ be permutation equivariant s.t. $D_{\theta}(\mathbf{P}\mathbf{X}_{T},\mathbf{Y})=\mathbf{P}D_{\theta}(\mathbf{X
}_{T},\mathbf{Y})$, and let $q(\mathbf{X}_{T})$ be permutation invariant. The optimal solution with respect to the cross-entropy loss with the identity reaction data is, for all nodes $i$ and $j$

$$ $\begin{cases}D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{\mathcal{N}}_{i,:}=\hat{ \mathbf{y}}^{\mathcal{N}},\hat{\mathbf{y}}^{\mathcal{N}}_{k}=\sum_{i}\mathbf{Y }_{i,k}/\sum_{i,k}\mathbf{Y}_{i,k},\\ D_{\theta}(\mathbf{X}_{T},\mathbf{Y})^{E}_{i,j,:}=\hat{\mathbf{y}}^{E},\hat{ \mathbf{y}}^{E}_{k}=\sum_{i,j}\mathbf{Y}_{i,j,k}/\sum_{i,j,k}\mathbf{Y}_{i,j,k },\end{cases}$ (19) $$

where $\hat{\mathbf{y}}^{\mathcal{N}}_{k}$ and $\hat{\mathbf{y}}^{E}_{k}$ are the marginal distributions of node and edge values in $\mathbf{Y}$.

###### Proof.

Nodes. The cross-entropy denoising loss for the nodes can be written as

$$ $\displaystyle CE$ $\displaystyle=-\sum_{(\mathbf{X}_{0},\mathbf{Y})}\mathbb{E}_{q(\mathbf{X}_{T} \,|\,\mathbf{X}_{0})}\sum_{i,k}\mathbf{X}^{\mathcal{N}}_{0,i,k}\log D_{\theta} (\mathbf{X}_{T},\mathbf{Y})_{i,k}$ (20) $\displaystyle=-\sum_{(\mathbf{X}_{0},\mathbf{Y})}\mathbb{E}_{q(\mathbf{X}_{T}) }\sum_{i,k}\mathbf{X}^{\mathcal{N}}_{0,i,k}\log D_{\theta}(\mathbf{X}_{T}, \mathbf{Y})_{i,k}$ (21) $\displaystyle=-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T})}\sum_{i,k} \mathbf{Y}^{\mathcal{N}}_{i,k}\log D_{\theta}(\mathbf{X}_{T},\mathbf{Y})_{i,k},$ (22) $$

where the first equality is due to $q(\mathbf{X}_{T}\,|\,\mathbf{X}_{0})$ containing no information about $\mathbf{X}_{0}$ at the end of the forward process, and the second equality is due to $\mathbf{X}_{0}=\mathbf{Y}$ in the data. Since $q(\mathbf{X}_{T})$ is permutation invariant, that is, all permuted versions $\mathbf{P}\mathbf{X}_{T}$ of $\mathbf{X}_{T}$ are equally probable, we can split the expectation into two parts $\mathbb{E}_{q(\mathbf{X}_{T})}[\cdot]\propto\mathbb{E}_{q(\mathbf{X}_{T}^{
\prime})}\sum_{\mathbf{P}}[\cdot]$, where $\mathbf{X}_{T}^{\prime}$ contain only graphs in distinct isomorphism classes, and $\sum_{\mathbf{P}}$ sums over all permutation matrices of size $N\times N$,

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{\mathbf{P}}\sum_{i,k}\mathbf{Y}^{\mathcal{N}}_{i,k}\log D_{\theta}( \mathbf{P}\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{i,k}.$ (23) $$

Due to the permutation equivariance, $D_{\theta}(\mathbf{P}\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}=\mathbf
{P}D_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}$, and $D_{\theta}(\mathbf{P}\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{i,k}=D
_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{\pi(i),k}$, where $\pi(i)$ denotes the index the index $i$ is mapped to in the permutation $\mathbf{P}$. Thus,

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{\pi}\sum_{i,k}\mathbf{Y}^{\mathcal{N}}_{i,k}\log D_{\theta}(\mathbf{X}_{ T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{\pi(i),k}$ (24) $\displaystyle=-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})}\sum_{i ,k}\sum_{\pi}\mathbf{Y}^{\mathcal{N}}_{\pi^{-1}(i),k}\log D_{\theta}(\mathbf{X }_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{i,k},$ (25) $$

where the equality is due to all permutations being in a symmetric position: What matters is the relative permutation between $\mathbf{Y}^{\mathcal{N}}$ and $D_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})$. Now, $\sum_{\pi}\mathbf{Y}^{\mathcal{N}}_{\pi^{-1}(i),k}=\sum_{\pi}\mathbf{Y}^{
\mathcal{N}}_{\pi(i),k}=C\sum_{i}\mathbf{Y}^{\mathcal{N}}_{i,k}$, because for each node index $i$, all the nodes in $\mathbf{Y}$ are included equally often due to symmetry. This is proportional to the marginal distribution $\hat{\mathbf{y}}^{\mathcal{N}}$ up to some constant, and thus we have:

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{i,k}\hat{\mathbf{y}}^{\mathcal{N}}_{k}\log D_{\theta}(\mathbf{X}_{T}^{ \prime},\mathbf{Y})_{i,k}.$ (26) $$

Clearly, the optimal value for each node output $i$ is the empirical marginal distribution $D_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{N}_{i,:}=(\hat{\mathbf{y}}^{
\mathcal{N}})^{\top}$.

Edges With the exact same steps, we can get the equivalent of Eq. 23 for the edges:

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{P}\sum_{i,j,k}\mathbf{Y}^{E}_{i,j,k}\log D_{\theta}(P\mathbf{X}_{T}^{ \prime},\mathbf{Y})^{E}_{i,j,k}.$ (27) $$

The permutation equivariance property for the edges is now written as $D_{\theta}(\mathbf{P}\mathbf{X}_{T}^{\prime},\mathbf{Y})^{E}=\mathbf{P}D_{
\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{E}\mathbf{P}^{\top}$, and $D_{\theta}(\mathbf{P}\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{i,jk}=
D_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{\mathcal{N}}_{\pi(i),\pi(j),k}$. Thus,

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{\pi}\sum_{i,j,k}\mathbf{Y}^{E}_{i,j,k}\log D_{\theta}(\mathbf{X}_{T}^{ \prime},\mathbf{Y})^{E}_{\pi(i),\pi(j),k}$ (28) $\displaystyle=-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})}\sum_{i ,j,k}\sum_{\pi}\mathbf{Y}^{E}_{\pi^{-1}(i),\pi^{-1}(j),k}\log D_{\theta}( \mathbf{X}_{T}^{\prime},\mathbf{Y})^{E}_{i,j,k},$ (29) $$

with the equality holding again due to symmetry. Now, for any pair of node indices $i$ and $j$, the set of all permutations contains all pairs of node indices $(\pi(i),\pi(j))$ equally often due to symmetry. These pairs correspond to edges in $\mathbf{Y}^{E}$, and thus $\sum_{\pi}\mathbf{Y}^{E}_{\pi^{-1}(i),\pi^{-1}(j),k}=\sum_{\pi}\mathbf{Y}^{E}_
{\pi^{(}i),\pi(j),k}=D\sum_{i,j}\mathbf{Y}^{E}_{i,j,k}$, where $D$ is a constant that counts how many times each edge pair appeared in the set of all permutations. This is again proportional to the marginal distribution over the edges $\hat{\mathbf{y}}^{E}$

$$ $\displaystyle CE$ $\displaystyle\propto-\sum_{\mathbf{Y}}\mathbb{E}_{q(\mathbf{X}_{T}^{\prime})} \sum_{i,j,k}\hat{\mathbf{y}}^{E}_{k}\log D_{\theta}(\mathbf{X}_{T}^{\prime}, \mathbf{Y})_{i,j,k}.$ (30) $$

Again, the optimal value for each edge output $(i,j)$ is $D_{\theta}(\mathbf{X}_{T}^{\prime},\mathbf{Y})^{E}_{i,j,:}=(\hat{\mathbf{y}}^{
E})^{\top}$.

∎

###### Theorem 3 .

###### Proof.

### B.2 Proof of the Generalized Distributional Invariance with Aligned Equivariance

We start by proving a useful lemma, and then continue and continue to the proof of the main theorem.

###### Lemma 1 .

(An aligned denoiser induces aligned distribution equivariance for a single reverse step)

If the denoiser function $D_{\theta}$ has the aligned equivariance property $D_{\theta}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{
\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=\mathbf{R}D_{\theta}(\mathbf{X},
\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$, then the conditional reverse distribution $p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{
Y}\to\mathbf{X}})$ has the property $p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\,\mathbf{R}\mathbf{X}_{t},\mathbf{Q}
\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=p_
{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}
\to\mathbf{X}})$.

###### Proof.

First, let us denote the transition probabilities from $t$ to $t-1$ with $F_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$, where formally $F_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{
\mathcal{N}}_{i,k}=\sum_{k^{\prime}}q(\mathbf{X}^{\mathcal{N}}_{t-1,i,k}\,|\,
\mathbf{X}^{\mathcal{N}}_{t,i,k},\mathbf{X}^{\mathcal{N}}_{0,i,k})D_{\theta}(
\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\mathcal{N}}_
{i,k^{\prime}}$ and $F_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{E}
_{i,j,k}=\sum_{k^{\prime}}q(\mathbf{X}^{E}_{t-1,i,j,k}\,|\,\mathbf{X}^{E}_{t,i
,j,k},\mathbf{X}^{E}_{0,i,j,k})D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}
^{\mathbf{Y}\to\mathbf{X}})^{E}_{i,j,k^{\prime}}$. Clearly, since the values of $F_{\theta}$ depend only pointwise on the values of $D_{\theta}$, $F_{\theta}$ is aligned permutation equivariant as well.

We continue by directly deriving the connection:

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\,\mathbf{R}\mathbf{X}_{ t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{ Q}^{\top})$ (31) $\displaystyle=\prod_{i}\sum_{k}(\mathbf{R}\mathbf{X}_{t-1})^{\mathcal{N}}_{i,k }F_{\theta}(\mathbf{R}\mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P} ^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})_{i,k}$ $\displaystyle\qquad\times\prod_{i,j}\sum_{k}(\mathbf{R}\mathbf{X}_{t-1})^{E}_{ i,j,k}F_{\theta}(\mathbf{R}\mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf{R} \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})^{E}_{i,j,k}$ (32) $\displaystyle=\prod_{i}\sum_{k}(\mathbf{X}_{t-1})^{\mathcal{N}}_{\pi(i),k}F_{ \theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})_{\pi(i ),k}$ $\displaystyle\qquad\times\prod_{i,j}\sum_{k}(\mathbf{X}_{t-1})^{E}_{\pi(i),\pi (j),k}F_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X} })^{E}_{\pi(i),\pi(j),k},$ (33) $$

where in the last line we used the aligned permutation equivariance definition, and the the effect of the permutation matrix $\mathbf{R}$ on index $i$ was denoted as $\pi(i)$. Now, regardless of the permutation, the products contain all possible values $i$ and pairs $i,j$ exactly once. Thus, the expression remains equal if we replace $\pi(i)$ with just $i$:

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\,\mathbf{R}\mathbf{X}_{ t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{ Q}^{\top})$ (34) $\displaystyle=\prod_{i}\sum_{k}(\mathbf{X}_{t-1})^{\mathcal{N}}_{i,k}F_{\theta }(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})_{i,k}\prod_{ i,j}\sum_{k}(\mathbf{X}_{t-1})^{E}_{i,j,k}F_{\theta}(\mathbf{X}_{t},\mathbf{Y} ,\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{E}_{i,j,k}$ (35) $\displaystyle=p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y}, \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}),$ (36) $$

which concludes the proof.
∎

###### Theorem 4 .

Aligned denoisers induce aligned permutation invariant distributions
If the denoiser function $D_{\theta}$ has the aligned equivariance property and the prior $p(\mathbf{X}_{T})$ is permutation invariant, then the generative distribution $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$ has the corresponding property for any permutation matrices $\mathbf{R}$ and $\mathbf{Q}$:

$$ $p_{\theta}(\mathbf{R}\mathbf{X}_{0}\,|\,\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf {P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=p_{\theta}(\mathbf{X}_{0}\,|\, \mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$ (37) $$

###### Proof.

Let us assume that the result holds for some noisy data level $t$: $p_{\theta}(\mathbf{R}\mathbf{X}_{t}\,|\,\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf
{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=p_{\theta}(\mathbf{X}_{t}\,|\,
\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$. We will then show that the same will hold for $\mathbf{X}_{t-1}$, which we can use to inductively show that the property holds for $\mathbf{X}_{0}$. We begin as follows:

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\,\mathbf{Q}\mathbf{Y}, \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=\sum_{\mathbf{X}_{t}}p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\, \mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf {X}}\mathbf{Q}^{\top})p_{\theta}(\mathbf{X}_{t}\,|\,\mathbf{Q}\mathbf{Y}, \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ (38) $\displaystyle=\sum_{\mathbf{X}_{t}}p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{R}^ {-1}\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})p_{\theta}( \mathbf{R}^{-1}\mathbf{X}_{t}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{ X}}).$ (39) $$

where on the second line we used Lem. 1 and the assumption that the result holds for noise level $t$.
The sum over $\mathbf{X}_{t}$ contains all possible graphs and all of their permutations. Thus, the exact value of $\mathbf{R}^{-1}$ does not affect the value of the final sum, as we simply go through the same permutations in a different order, and aggregate the permutations with the sum. Thus,

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{t-1}\,|\,\mathbf{Q}\mathbf{Y}, \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=\sum_{\mathbf{X}_{t}}p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_ {t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})p_{\theta}(\mathbf{X}_{t} \,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$ (40) $\displaystyle=p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y }\to\mathbf{X}})$ (41) $$

showing that if the result holds for level $t$, then it also holds for level $t-1$. We only need to show that it holds for level $\mathbf{X}_{T-1}$ to start the inductive chain:

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{T-1}\,|\,\mathbf{Q}\mathbf{Y}, \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=\sum_{\mathbf{X}_{T}}p_{\theta}(\mathbf{R}\mathbf{X}_{T-1}\,|\, \mathbf{X}_{T},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf {X}}\mathbf{Q}^{\top})p(\mathbf{X}_{T})$ (42) $\displaystyle=\sum_{\mathbf{X}_{T}}p_{\theta}(\mathbf{X}_{T-1}\,|\,\mathbf{R}^ {-1}\mathbf{X}_{T},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})p(\mathbf{R }^{-1}\mathbf{X}_{T}),$ (43) $$

where on the second line we again used Lem. 1 and the permutation invariance of $p(\mathbf{X}_{T})$. Again, the exact value of $\mathbf{R}^{-1}$ does not matter for the sum, since the sum goes through all possible permutations in any case. Thus we have

$$ $\displaystyle p_{\theta}(\mathbf{R}\mathbf{X}_{T-1}\,|\,\mathbf{Q}\mathbf{Y}, \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=\sum_{\mathbf{X}_{T}}p_{\theta}(\mathbf{X}_{T-1}\,|\,\mathbf{X}_ {T},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})p(\mathbf{X}_{T})$ (44) $\displaystyle=p_{\theta}(\mathbf{X}_{T-1}\,|\,\mathbf{Y},\mathbf{P}^{\mathbf{Y }\to\mathbf{X}}).$ (45) $$

Thus, since the property holds for $\mathbf{X}_{T-1}$, it also holds for $\mathbf{X}_{T-2},\dots$, until $\mathbf{X}_{0}$. This concludes the proof.
∎

###### Lemma 1 .

###### Proof.

###### Theorem 4 .

###### Proof.

### B.3 Proofs that Our Denoisers Are Aligned Permutation Equivariant

In this section, we show for each of the three alignment methods that the corresponding denoisers do indeed fall within the aligned permutation equivariance function class. Fig. A7 summarizes the different alignment methods.

Figure: Figure A7: Different ways to align the graphs in the architecture. All of them can be combined. The $\oplus$ sign means concatenation along the feature dimension, and $+$ is the standard addition. All of the methods can be combined together.
Refer to caption: x10.png

#### Atom-mapped positional encodings

We start by writing out one side of the aligned permutation equivariance condition, $D_{\theta}(\mathbf{R}\mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^
{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ for this particular function class, and directly show that it equals $\mathbf{P}D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to
\mathbf{X}})$.

$$ $\displaystyle D_{\theta}(\mathbf{R}\mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf {R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=f_{\theta}^{\mathbf{X}}(\begin{bmatrix}\mathbf{R}\mathbf{X}_{t}^ {N}&\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top}\mathbf{Q} \varphi\end{bmatrix},\mathbf{R}\mathbf{X}^{E}\mathbf{R}^{\top},\begin{bmatrix} \mathbf{Q}\mathbf{Y}^{N}&\mathbf{Q}\varphi\end{bmatrix},\mathbf{Q}\mathbf{Y}^{ E}\mathbf{Q}^{\top})$ (46) $\displaystyle=f_{\theta}^{\mathbf{X}}(\begin{bmatrix}\mathbf{R}\mathbf{X}_{t}^ {N}&\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\varphi\end{bmatrix},\mathbf {R}\mathbf{X}^{E}\mathbf{R}^{\top},\begin{bmatrix}\mathbf{Q}\mathbf{Y}^{N}& \mathbf{Q}\varphi\end{bmatrix},\mathbf{Q}\mathbf{Y}^{E}\mathbf{Q}^{\top}),$ (47) $$

where $f_{\theta}$ itself is a function that is permutation equivariant for the combined $\mathbf{X}$ and $\mathbf{Y}$ graph as input. This means that the neural net itself gives an output for the entire combined graph, but we only consider the $\mathbf{X}$ subgraph as the denoiser output, denoted here as $f_{\theta}^{\mathbf{X}}$. For clarity, we can combine the reactant and product node features and adjacency matrices in the notation, and use the permutation equivariance property of our GNN:

$$ $\displaystyle f_{\theta}$ $\displaystyle(\begin{bmatrix}\mathbf{R}\mathbf{X}_{t}^{\mathcal{N}}&\mathbf{R} \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\varphi\\ \mathbf{Q}\mathbf{Y}^{\mathcal{N}}&\mathbf{Q}\varphi\end{bmatrix},\begin{ bmatrix}\mathbf{R}\mathbf{X}_{t}^{E}\mathbf{R}^{\top}&0\\ 0&\mathbf{Q}\mathbf{Y}^{E}\mathbf{Q}^{\top}\end{bmatrix})$ (48) $\displaystyle=\begin{bmatrix}\mathbf{R}&0\\ 0&\mathbf{Q}\end{bmatrix}f_{\theta}(\begin{bmatrix}\mathbf{X}_{t}^{\mathcal{N} }&\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\varphi\\ \mathbf{Y}^{\mathcal{N}}&\varphi\end{bmatrix},\begin{bmatrix}\mathbf{X}_{t}^{E }&0\\ 0&\mathbf{Y}^{E}\end{bmatrix})\quad\textrm{(permutation equivariance of base NN)}$ (49) $$

Taking only the $\mathbf{X}$ part of the output and reverting to $D_{\theta}$ notation, we directly arrive at the result that $D_{\theta}(\mathbf{R}\mathbf{X}_{t},\mathbf{Q}\mathbf{Y},\mathbf{R}\mathbf{P}^
{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})=\mathbf{R}D_{\theta}(\mathbf{X}_{t
},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$.

#### Directly adding 𝐘 𝐘 \mathbf{Y} bold_Y to the output

We again start by writing out an aligned equivariant input to the denoiser, with $f_{\theta}$ again denoting a network that is permutation equivariant with respect to the combined $\mathbf{X}$ and $\mathbf{Y}$ graph:

$$ $\displaystyle D_{\theta}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},$ $\displaystyle\mathbf{R}\mathbf{P}^{\mathbf{X}\to\mathbf{Y}}\mathbf{Q}^{\top})= \text{softmax}(f_{\theta}^{\mathbf{X}}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{ Y})+\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top}\mathbf{Q} \mathbf{Y})$ (50) $\displaystyle=\text{softmax}(\mathbf{R}f_{\theta}^{\mathbf{X}}(\mathbf{X}, \mathbf{Y})+\mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y})\quad \text{(permutation equivariance of base denoiser)}$ (51) $\displaystyle=\mathbf{R}\ \text{softmax}(f_{\theta}^{\mathbf{X}}(\mathbf{X}, \mathbf{Y})+\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y})$ (52) $\displaystyle=\mathbf{R}D_{\theta}(\mathbf{X},\mathbf{Y},\mathbf{P}^{\mathbf{X }\to\mathbf{Y}})$ (53) $$

where we were able to move the permutation outside the softmax since the softmax is applied on each node and edge separately.

#### Aligning 𝐘 𝐘 \mathbf{Y} bold_Y and 𝐗 𝐗 \mathbf{X} bold_X at the input to the model

Let’s denote by $[\mathbf{X}\ \ \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}]$ concatenation along the feature dimension for both the nodes and edges of the graphs. Recall then that the definition of aligning the graphs in the input is $D_{\theta}(\mathbf{X},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})=f_{
\theta}([\mathbf{X}\ \ \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}])$, where $f_{\theta}$ is a permutation equivariant denoiser. Writing out the aligned equivariance condition

$$ $\displaystyle D_{\theta}(\mathbf{R}\mathbf{X},\mathbf{Q}\mathbf{Y},\mathbf{R} \mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top})$ $\displaystyle=f_{\theta}([\mathbf{R}\mathbf{X}\ \ \mathbf{R}\mathbf{P}^{ \mathbf{Y}\to\mathbf{X}}\mathbf{Q}^{\top}\mathbf{Q}\mathbf{Y}])=f_{\theta}([ \mathbf{R}\mathbf{X}\ \ \mathbf{R}\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{ Y}])$ (54) $\displaystyle=f_{\theta}(\mathbf{R}[\mathbf{X}\ \ \mathbf{P}^{\mathbf{Y}\to \mathbf{X}}\mathbf{Y}])=\mathbf{R}f_{\theta}([\mathbf{X}\ \ \mathbf{P}^{ \mathbf{Y}\to\mathbf{X}}\mathbf{Y}])$ (55) $\displaystyle=D_{\theta}(\mathbf{X},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to \mathbf{X}})$ (56) $$

which shows that this method results in aligned equivariance as well.

### B.4 A Single-layer Graph Transformer with Orthogonal Atom-mapped Positional Encodings is Able to Implement the Identity Reaction Solution for Nodes

Here, we show that a single-layer Graph Transformer neural net can model the identity reaction for the nodes given orthogonal atom-mapped positional encodings (e.g., the graph Laplacian eigenvector-based ones). In particular it is possible to find a $\theta$ such that $D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{
\mathcal{N}}=\mathbf{Y}^{\mathcal{N}}$. We hope that this section can serve as an intuitive motivation for why matched positional encodings help in copying over the structure from the product to the reactant side.

Recall that we have $N$ atoms on both sides of the reaction. Let us map the atom mapping indices to basis vectors in an orthogonal basis $\varphi=[\varphi_{1},\varphi_{2},\dots,\varphi_{N}]^{\top}$. In practice, the node input to the neural net on the reactant side is now $\mathbf{X}_{t}^{*}=[\mathbf{X}_{t}^{\mathcal{N}},\varphi]$, and the node input on the product side is $\mathbf{Y}^{*}=[\mathbf{Y}^{\mathcal{N}},\varphi]$.

Now, a single-layer Graph Transformer looks as follows (as defined in the Vignac et al. [] codebase):

$$ $\displaystyle\mathbf{N}$ $\displaystyle=[(\mathbf{X}_{t}^{*\mathcal{N}})^{\top},(\mathbf{Y}^{*\mathcal{N }})^{\top}]^{\top}\quad\text{(Concatenate rows)}$ (57) $\displaystyle\mathbf{E}$ $\displaystyle=\begin{bmatrix}\mathbf{X}_{t}^{E}&0\\ 0&\mathbf{Y}^{E}\end{bmatrix}\quad\text{(Create joined graph)}$ (58) $\displaystyle\mathbf{N}_{1}$ $\displaystyle=MLP_{N}(\mathbf{N})\quad\text{(Applied with respect to last dimension)}$ (59) $\displaystyle\mathbf{E}_{1}$ $\displaystyle=\frac{1}{2}(MLP_{E}(\mathbf{E})+MLP_{E}(\mathbf{E})^{\top})\quad \text{(Symmetrize the input)}$ (60) $\displaystyle t_{1}$ $\displaystyle=MLP_{t}(t)$ (61) $\displaystyle\mathbf{Q},\mathbf{K},\mathbf{V}$ $\displaystyle=\mathbf{W}_{Q}\mathbf{N}_{1},\mathbf{W}_{K}\mathbf{N}_{1}, \mathbf{W}_{V}\mathbf{N}_{1}\quad\text{(One attention head for simplicity)}$ (62) $\displaystyle\mathbf{A}_{1}$ $\displaystyle=\mathbf{Q}\mathbf{K}^{\top}$ (63) $\displaystyle\mathbf{A}_{2}$ $\displaystyle=\mathbf{A}_{1}*(\mathbf{W}_{\mathbf{E}_{2}}\mathbf{E}_{1}+1)+ \mathbf{W}_{\mathbf{E}_{3}}\mathbf{E}_{1}$ (64) $\displaystyle\mathbf{A}_{3}$ $\displaystyle=\text{softmax}(\mathbf{A}_{3})$ (65) $\displaystyle\mathbf{N}_{2}$ $\displaystyle=\mathbf{A}\mathbf{V}/\sqrt{d_{f}}\quad\text{($d_{f}$ is the embedding dim of V)}$ (66) $\displaystyle\mathbf{N}_{3}$ $\displaystyle=\mathbf{N}_{2}*(\mathbf{W}_{t_{2}}t_{1}+1)+\mathbf{W}_{t_{3}}t_{1}$ (67) $\displaystyle\mathbf{E}_{2}$ $\displaystyle=\mathbf{W}(\mathbf{A}_{2}*(\mathbf{W}_{t_{4}}t_{1}+1)+W_{t_{5}}t _{1})$ (68) $\displaystyle t_{4}$ $\displaystyle=MLP_{t_{4}}(\mathbf{W}t_{2}+\mathbf{W}[\mathrm{min}(\mathbf{N}_{ 3}),\mathrm{max}(\mathbf{N}_{3}),\mathrm{mean}(\mathbf{N}_{3}),\mathrm{std}( \mathbf{N}_{3})]$ $\displaystyle+\mathbf{W}[\mathrm{min}(\mathbf{N}_{3}),\mathrm{max}(\mathbf{N}_ {3}),\mathrm{mean}(\mathbf{N}_{3}),\mathrm{std}(\mathbf{N}_{3})]$ (69) $\displaystyle\mathbf{E}_{3}$ $\displaystyle=MLP_{\mathbf{E}_{2}}(\mathbf{E}_{2})+\mathbf{E}$ (70) $\displaystyle\mathbf{E}_{out}$ $\displaystyle=\frac{1}{2}(\mathbf{E}_{3}+\mathbf{E}_{3}^{\top})\quad\text{( Symmetrize the output)}$ (71) $\displaystyle\mathbf{N}_{out}$ $\displaystyle=MLP_{N_{3}}(\mathbf{N}_{3})+\mathbf{N}$ (72) $\displaystyle t_{out}$ $\displaystyle=t_{4}+t$ (73) $$

Now, for purposes of illustration, we can define most linear layers to be zero layers: $MLP_{E},MLP_{t},\mathbf{W}_{E_{2}},\mathbf{W}_{E_{3}},\mathbf{W}_{t_{2}},
\mathbf{W}_{t_{3}},\mathbf{W}_{t_{4}},\mathbf{W}_{t_{5}}=0$. In addition to this, we define $MLP_{N}$ to be an identity transform. $\mathbf{W}_{Q}$ and $\mathbf{W}_{K}$ are both chosen as picking out the $\mathbf{U}$ columns of $\mathbf{N}_{1}$, with additional overall scaling by some constant $\alpha$. $\mathbf{W}_{V}$ is chosen to pick out the product node labels: $\mathbf{W}_{V}\mathbf{N}_{1}=\begin{bmatrix}0\\
\mathbf{Y}^{\mathcal{N}}\end{bmatrix}$. Now, we can easily see how the Graph Transformer can obtain the optimal denoising solution for the nodes. Consider an input $\mathbf{N}=[(\mathbf{P}\mathbf{X}_{t}^{*})^{\top},(\mathbf{Y})^{\top}]^{\top}$, where the reactant side is permuted. The output of the network should be $\mathbf{P}\mathbf{Y}$. Focusing on the parts of the network that compute the node features:

$$ $\displaystyle\mathbf{A}_{1}$ $\displaystyle=\alpha^{2}\begin{bmatrix}\mathbf{P}\varphi\varphi^{\top}\mathbf{ P}^{\top}&\mathbf{P}\varphi\varphi^{\top}\\ \varphi\varphi^{\top}\mathbf{P}^{\top}&\varphi\varphi^{\top}\end{bmatrix}= \alpha^{2}\begin{bmatrix}\mathbf{I}&\mathbf{P}\\ \mathbf{P}^{\top}&\mathbf{I}\end{bmatrix},$ (74) $\displaystyle\mathbf{A}_{2}$ $\displaystyle=\mathbf{A}_{1},$ (75) $\displaystyle\mathbf{A}_{3}$ $\displaystyle\approx\text{softmax}(\mathbf{A}_{2})=\frac{1}{2}\begin{bmatrix} \mathbf{I}&\mathbf{P}\\ \mathbf{P}^{\top}&\mathbf{I}\end{bmatrix}\quad\text{(If $\alpha\gg 1$)},$ (76) $\displaystyle\mathbf{N}_{2}$ $\displaystyle=\frac{\mathbf{A}_{3}\mathbf{V}}{\sqrt{d_{F}}}=\begin{bmatrix} \mathbf{I}&\mathbf{P}\\ \mathbf{P}^{\top}&\mathbf{I}\end{bmatrix}\begin{bmatrix}0\\ Y_{N}\end{bmatrix}/(2\sqrt{d_{f}})=\begin{bmatrix}\mathbf{P}\mathbf{Y}_{N}\\ \mathbf{Y}_{N}\end{bmatrix}/(2\sqrt{d_{f}}).$ (77) $$

Here, we used the fact that we chose the positional embeddings to be an orthogonal basis, and $\varphi\varphi^{\top}=\mathbf{I}$, as well as the fact that $\mathbf{P}\mathbf{P}^{\top}=\mathbf{I}$ for any permutation matrices. The term $\frac{1}{2}$ in the third equation came from the fact that each row of $\mathbf{A}_{1}$ contains two non-zero values that are also equal. The probability gets divided between the two of them in the softmax if the logits are scaled large enough, and the approximation becomes arbitrarily accurate.

From now on, since we are interested in the denoising output only for the reactant side, we drop out the reactant side $\mathbf{Y}^{\mathcal{N}}$ and only focus on the $\mathbf{P}\mathbf{Y}^{\mathcal{N}}$ part. We choose the final $MLP_{\mathbf{N}_{2}}$ to scale the output by some factor $\beta\gg 1$:

$$ $\displaystyle\mathbf{N}_{out}$ $\displaystyle=\beta\mathbf{N}_{2}+\mathbf{N}$ (78) $\displaystyle n_{\theta}(\mathbf{P}\mathbf{X}_{t}^{*},\mathbf{Y}^{*})^{ \mathcal{N}}$ $\displaystyle=\text{softmax}(\beta\mathbf{P}\mathbf{Y}^{\mathcal{N}}+\mathbf{N })\approx\mathbf{P}\mathbf{Y}^{\mathcal{N}}.$ (79) $$

Here, the approximation can be made arbitrarily accurate by scaling $\beta$ to a higher value, since the logits will become more and more peaked towards the values where $P\mathbf{Y}^{\mathcal{N}}$ equals one instead of zero. This showcases how the attention mechanism in the Graph Transformer pairs with atom mapping-based orthogonal positional encodings to achieve the identity function from products to reactants.

## Appendix C Experimental Setup

### C.1 Data: USPTO Data Sets

All open-source data sets available for reaction modelling are derived in some form from the patent mining work of Lowe []. We distinguish 5 subsets used in previous work: 15k, 50k, MIT, Stereo, and full (original data set). Table A2 provides key information about the subsets.

**Table A2: UPSTO-50K subsets used in retrosynthesis**
| Subset | Introduced by | # of reactions | Preprocessiong & Data split (script) |
| --- | --- | --- | --- |
| Full | Lowe [] | 1 808 938 | Dai et al. [] |
| Stereo | Schwaller et al. [] | 1 002 970 | Schwaller et al. [] |
| MIT | Jin et al. [] | 479 035 | - |
| 50k | Schneider et al. [] | 50 016 | Dai et al. [] |

#### 15k

is proposed by Coley et al. []. The subset includes reactions covered by the 1.7k most common templates. All molecules appearing in the reaction are included to model the involvement of reagents and solvents despite not contributing with atoms to the product.

#### 50k

is preprocessed by Schneider et al. []. The goal of the analysis is to assign roles (reactant, reagent, solvent) to different participants in a reaction through atom mapping. This effort led to the creation of an atom-mapped and classified subset of around 50k reactions, which is used nowadays as a benchmark for retrosynthesis tasks. It is not clear how said subset was selected.

#### MIT

is used by Jin et al. []. The preprocessing is described as ‘removing duplicate and erroneous reactions’ with no further explanation of what qualifies as an erroneous reaction. The output of this filtering is a data set of 49k reactions (from an original set of 1.8M reactions).

#### Stereo

is proposed by Schwaller et al. []. The authors apply a more flexible filtering strategy compared to USPTO-MIT. Their data set only discards   800k reactions from the original data set because they are duplicates or they could not be canonicalized by RDKit. In addition, the data set only considers single-product reactions (92% of the full data set), as opposed to splitting multi-product reactions. The preprocessing steps include removing reagents (molecules with no atoms appearing in the product), removing hydrogen atoms from molecules, discarding atom-mapping information and canonicalizing molecules. In addition, since the original method applied to this subset is a language model, tokenization is performed on the atoms.

#### Full

is preprocessed by Dai et al. []. The processing includes removing duplicate reactions, splitting reactions with multiple products into multiple reactions with one product, removing reactant molecules appearing unchanged on the product side, removing all reactions with bad atom-mapping (i.e., when the sorted mapping between products and reactants is not one-to-one), and removing bad products (missing mapping, or not parsed by Rdkit).

#### Our choice

Similar to many other works on retrosynthesis, we use 50k as the main data set to evaluate our method.

### C.2 Notes on Our Sampling and Ranking Procedures

#### Duplicate removal

Removing duplicates from the set of generated precursors is a common methodology in retrosynthesis, albeit often not discussed explicitly in papers. The benefit of duplicate removal is to ensure that an incorrect molecule that is nevertheless judged as the best one according to the ranking scheme does not fill up all of the top-$k$ positions after ranking. While this does not affect top-1 scores, not removing the duplicates would degrade the other top-$k$ scores significantly.

#### Choice of scoring function

Specifically, we use the following formula for approximating the likelihood of the sample under the model.

$$ $s(\mathbf{X})=(1-\lambda)\frac{\mathrm{count}(\mathbf{X})}{\sum_{\mathbf{X}^{ \prime}}\mathrm{count}(\mathbf{X}^{\prime})}+\lambda\cdot\frac{e^{\mathrm{elbo }(\mathbf{X})}}{\sum_{\mathbf{X}^{\prime}}e^{\mathrm{elbo}(\mathbf{X}^{\prime} )}},$ (80) $$

where count(.) returns the number of occurrences of sample X in the set of generated samples (by default, 100), elbo(.) computes the variational lower bound of the specific sample under the model, and $\lambda$ is a weighting hyperparameter. The sums are taken over the set of generated samples. Intuitively, the idea is to provide an estimate of the likelihood from two routes. The ELBO is an estimate (a lower bound) of $\log p_{\theta}(\mathbf{X})$, and exponentiating and normalizing gives an estimate of the probability distribution. Since the same reactants are often repeated in a set of 100 samples, the counts can be used as a more direct proxy, although they inherently require a relatively large amount of samples to limit the variance of this estimate.

We set the value of $\lambda$ to be $0.9$, although we find that the top-$k$ scores are not at all sensitive to variation in the exact value, as long as it is below 1, and the count information is used. Thus, the counts seem more important than the ELBO, which may be due to the lower bound nature of the ELBO or stochasticity in estimating its value. More accurate likelihood estimation schemes for diffusion models, such as exact likelihood values using the probability flow ODE [], could be a valuable direction for future research in the context of retrosynthetic diffusion models.

### C.3 Details on stereochemistry

Our model does not explicitly consider changes in stereochemistry in the reaction, but instead, we use the atom mappings implicitly assigned to the samples by the model to transfer the chiral tags from the products to the reactants. The initial choice of $\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$ at the start of sampling can be considered to be the atom mapping of the generated reactants, given that the model has been trained on correct atom mappings.

For the chiral tags, we take the ground-truth SMILES for the product molecules from the dataset and assign the corresponding chiral tag to the corresponding atom mapping on the generated reactants. For cis/trans isomerism, we use the Chem.rdchem.BondDir bond field in rdkit molecules and transfer them to the reactant side based on the atom mapping of the pair of atoms at the start and end of the bond.

Note that when using rdkit, transferring chirality requires some special care: The chiral tags Chem.ChiralType.CHI_TETRAHEDRAL_CCW and Chem.ChiralType.CHI_TETRAHEDRAL_CW are defined in the context of the order in which the bonds are attached to the chiral atom in the molecule data structure. Thus, the chiral tag sometimes has to be flipped to retain the correct stereochemistry, based on whether the order of the bonds is different on the reactant molecule data structure and the product molecule data structure.

### C.4 Details of the Evaluation Procedure

#### Top- k 𝑘 k italic_k scores

We evaluate the top-$k$ scores by ranking the list of generated and deduplicated samples and calculating the percentage of products for which the ground-truth reactants are in the first k elements in the list.

#### Mean reciprocal rank

We formally define the MRR as $\textrm{MRR}=\mathbb{E}_{p(r)}[r^{-1}]$. Verbally, it is the expected value of the inverse of the amount of reactant suggestions that the model makes before encountering the ground truth, and as such measures how early on is the correct reactant encountered in the ranked samples. It also incorporates the intuition that the difference between obtaining the correct reactants in, say, the 9th or the 10th position, is not as significant as the difference between the 1st and 2nd positions.

While we do not have direct access to the entire $p(r)$ just based on the top-$k$ scores, we can estimate it with a uniform distribution assumption on $r$ within the different top-$k$ ranges. Formally, we define four sets $S_{1},S_{2},S_{3},S_{4}=\{1\},\{2,3\},\{4,5\},\{6,7,8,9,10\}$ in which $p(r)$ is assumed to be uniform, and $s(r)\in\{1,2,3,4\}$ denotes the group that rank $r$ belongs to. Top-$k$ is denoted as $\text{top}_{k}$, where $k\in\{1,3,5,10\}$ in our case. Note that they are also equal to the cumulative distribution of $p(r)$ until $k$. We thus define $\hat{p}(r)=\frac{\text{top}_{\max(G_{s(r)})}-\text{top}_{\max(G_{s(r)-1})}}{|G
_{s(r)}|}$ and $\hat{p}(1)=\text{top}_{1}$. For the case where the ground truth was not in the top-10, we assume it is not recovered and place the rest of the probability mass on $p(\infty)$. Our MRR estimate is then defined as

$$ $\widehat{\textrm{MRR}}=\sum_{r=1}^{10}\hat{p}(r)\frac{1}{r}.$ (81) $$

In cases where we do not have top-$k$ values for all {1,3,5,10} (such as the Augmented Transformer in LABEL:tab:50k for top-3), we assume that $\hat{p}(r)$ is constant in the wider interval between the preceding and following top-$k$s (2–5 in the case of the top-3 missing).

### C.5 Neural Network Architecture, Hyperparameters, and Compute Resources

We discuss here the neural network architecture and hyper-parameters we choose. Our denoiser is implemented as a Graph Transformer [], based on the implementation of Vignac et al. [] with additional graph-level level features added to the input of the model. See Vignac et al. [] for an in-depth discussion of the neural network.

In all of our models, we use 9 Graph Transformer layers. When using Laplacian positional encodings, we get the 20 eigenvectors of the Graph Laplacian matrix with the largest eigenvalues and assign to each node a 20-dimensional feature vector.

We use a maximum of 15 ’blank’ nodes, in practice meaning that the models have the capacity to add 15 additional atoms on the reactant side. In another detail, following Vignac et al. [], we weigh the edge components in the cross-entropy loss by a factor of 5 compared to the node components.

We used a batch size of 16 for the models where the expanded graph containing X and Y as subgraphs is given as input. These models were trained for approximately 600 epochs with a single A100/V100/AMD MI250x GPU. For the model where alignment is done by concatenating $\mathbf{Y}$ along the feature dimension in the input, the attention map sizes were smaller and we could fit a larger batch of 32 with a single V100 GPU. This model was trained for 600 epochs. The training time for all of our models was approximately three days. In early experiments and developing the model, we trained or partially trained multiple models that did not make it to the main paper. Sampling 100 samples for one product with $T=100$ from the model takes roughly 60 seconds with the current version of our code with an AMD MI250x GPU, and 100 samples with $T=10$ takes correspondingly about 6 seconds. It is likely that the inference could be optimized, increasing the sample throughput.

The reported models were chosen based on evaluating different checkpoints with 10 diffusion steps on the validation set for different checkpoints and chose the best checkpoint based on the MRR score.

### C.6 Comparison to Other Baselines

**Table A3: Top-$k$ accuracy and MRR on the USPTO-50k test data set—extended comparison, including models with pretraining on larger data sets, and Retrobridge [], where the evaluation is done with a relaxed metric that does not consider charges or stereochemistry.**
|  | Method | $\mathbf{k=1}\uparrow$ | $\mathbf{k=3}\uparrow$ | $\mathbf{k=5}\uparrow$ | $\mathbf{k=10}\uparrow$ | $\widehat{\mathbf{MRR}}\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Pre-<br>trained | RSMILES [] | 56.3 | 79.2 | 86.2 | 91.0 | 0.680 |
| PMSR [] | 62.0 | 78.4 | 82.9 | 86.8 | 0.704 |  |
| Temp. | Retrosym [] | 37.3 | 54.7 | 63.3 | 74.1 | 0.480 |
| GLN [] | 52.5 | 74.7 | 81.2 | 87.9 | 0.641 |  |
| LocalRetro [] | 52.6 | 76.0 | 84.4 | 90.6 | 0.650 |  |
| Synthon | GraphRetro [] | 53.7 | 68.3 | 72.2 | 75.5 | 0.611 |
| RetroDiff [] | 52.6 | 71.2 | 81.0 | 83.3 | 0.629 |  |
| MEGAN [] | 48.0 | 70.9 | 78.1 | 85.4 | 0.601 |  |
| G2G [] | 48.9 | 67.6 | 72.5 | 75.5 | 0.582 |  |
| Template-free | SCROP [] | 43.7 | 60.0 | 65.2 | 68.7 | 0.521 |
| Tied Transformer [] | 47.1 | 67.1 | 73.1 | 76.3 | 0.572 |  |
| Aug. Transformer [] | 48.3 | - | 73.4 | 77.4 | 0.569 |  |
| Retrobridge (*) [] | 50.3 | 74.0 | 80.3 | 85.1 | 0.622 |  |
| GTA_aug [] | 51.1 | 67.6 | 74.8 | 81.6 | 0.605 |  |
| Graph2SMILES [] | 52.9 | 66.5 | 70.0 | 72.9 | 0.597 |  |
| Retroformer [] | 52.9 | 68.2 | 72.5 | 76.4 | 0.608 |  |
| DualTF_aug [] | 53.6 | 70.7 | 74.6 | 77.0 | 0.619 |  |
|  | Unaligned | 4.1 | 6.5 | 7.8 | 9.8 | 0.056 |
|  | DiffAlign-input | 44.1 | 65.9 | 72.2 | 78.7 | 0.554 |
|  | DiffAlign-PE | 49.0 | 70.7 | 76.6 | 81.8 | 0.601 |
| Ours | DiffAlign-PE+skip | 54.7 | 73.3 | 77.8 | 81.1 | 0.639 |

We present other results in the literature not directly comparable to our model.

#### Methods with a different evaluation procedure

Despite Igashov et al. []’s RetroBridge being closely related to a diffusion model, we cannot include it in a straightforward comparison because it discards atom charges from the ground truth smiles during evaluation. Specifically, the model uses only atom types as node features and compares the generated samples to the smiles reconstructed from the ground truth data through the same encoding (i.e. without charges too) (^2^22This can be seen in the code shared by Igashov et al. []: [https://github.com/igashov/RetroBridge](https://github.com/igashov/RetroBridge)).

#### Methods with pretraining

Zhong et al. [] and Jiang et al. [] pre-train their models on the USPTO-Full and Pistachio data sets, respectively, and as such the results are not directly comparable to models trained on the standard USPTO-50k benchmark. Pretraining with diffusion models is an interesting direction for future research, but we consider it outside the scope of our work. Furthermore, comparison between models with different pretraining datasets and pretraining strategies has the danger of complicating comparisons, given that relative increases in performance could be explained by the model, the pretraining strategy, or the pretraining dataset. As such, we believe that standardized benchmarks like USPTO-50k are necessary when researching modelling strategies.

## Appendix D Adding Post-Training Conditioning to Discrete Diffusion Models

In this section, we show a method to add additional controls and conditions on the used discrete diffusion model post-training. While the notation is from the point of view of our retrosynthesis model, the method here applies in general to any discrete diffusion model. We start by writing the Bayes’ rule for an additional condition $y$ (e.g., a specified level of drug-likeness or synthesizability, or an inpainting mask)

$$ $\displaystyle p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y},y)$ $\displaystyle\propto p(y\,|\,\mathbf{X}_{t-1},\mathbf{X}_{t},\mathbf{Y})p_{ \theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$ (82) $\displaystyle=p(y\,|\,\mathbf{X}_{t-1},\mathbf{Y})p_{\theta}(\mathbf{X}_{t-1} \,|\,\mathbf{X}_{t},\mathbf{Y})$ (83) $$

where the second equation was due to the Markovian structure of the generative process ($\mathbf{X}_{t-1}$ d-separates $y$ and $\mathbf{X}_{t}$). Now, we can take the log and interpret the probabilities as tensors $\mathbf{P}_{\theta}(y\,|\,\mathbf{X}_{t-1},\mathbf{Y})$ and $\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$ defined in the same space as the one-hot valued tensors $\mathbf{X}_{t-1}$ and $\mathbf{X}_{t}$. We get:

$$ $\displaystyle\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y},y)\propto\log\mathbf{P}(y\,|\,\mathbf{X}_{t-1},\mathbf{Y})+\log \mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$ (84) $$

Similarly to [], we can now Taylor expand $\log\mathbf{P}_{\theta}(y\,|\,\mathbf{X}_{t-1},\mathbf{Y})$ around $\mathbf{X}_{t}$ with

$$ $\displaystyle\log\mathbf{P}(y\,|\,\mathbf{X}_{t-1},\mathbf{Y})\approx\log \mathbf{P}(y\,|\,\mathbf{X}_{t},\mathbf{Y})+\nabla_{\mathbf{X}_{t}^{\prime}} \log\mathbf{P}(y\,|\,\mathbf{X}_{t}^{\prime},\mathbf{Y})|_{\mathbf{X}_{t}^{ \prime}=\mathbf{X}_{t}}(\mathbf{X}_{t-1}-\mathbf{X}_{t})$ (85) $$

Given that we are interested in the distribution w.r.t. $\mathbf{X}_{t-1}$, the $\mathbf{X}_{t}$ terms are constant when we plug them in to Eq. 84, resulting in

$$ $\displaystyle\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y},y)\propto\nabla_{\mathbf{X}_{t}^{\prime}}\log\mathbf{P}(y\,|\, \mathbf{X}_{t}^{\prime},\mathbf{Y})|_{\mathbf{X}_{t}^{\prime}=\mathbf{X}_{t}} \mathbf{X}_{t-1}+\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}|\mathbf{X}_{t}, \mathbf{Y})$ (86) $$

Simplifying the notation to assume taking the gradient at $\mathbf{X}_{t}$, we can also write

$$ $\displaystyle\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y},y)\propto\nabla_{\mathbf{X}_{t}}\log\mathbf{P}_{\theta}(y\,|\, \mathbf{X}_{t},\mathbf{Y})\mathbf{X}_{t-1}+\log\mathbf{P}_{\theta}(\mathbf{X}_ {t-1}\,|\,\mathbf{X}_{t},\mathbf{Y})$ (87) $$

In practice, the equation means that given $\log p_{\theta}(y\,|\,\mathbf{X}_{t},\mathbf{Y})$, we get $p_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t},\mathbf{Y},y)$ by adding the input gradient of $\log p_{\theta}(y\,|\,\mathbf{X}_{t},\mathbf{Y})$ to the logits given by the regular reverse transition and re-normalizing.

It would be possible to approximate $\log p_{\theta}(y\,|\,\mathbf{X}_{t},\mathbf{Y})$ by training an additional classifier, leading to so-called classifier guidance []. However, if we have direct access to $p(y\,|\,\mathbf{X}_{0})$ (e.g., a synthesizability model), there exists a wide range of methods developed for continuous diffusion models that provide different levels of approximations of the true conditional distribution [, , , , , ]. Here, we show an approach similar to [] for discrete diffusion models by using the gradients with respect to the denoiser $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$

$$ $\displaystyle p(y\,|\,\mathbf{X}_{t},\mathbf{Y})$ $\displaystyle=\sum_{\mathbf{X}_{0}}q(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf {Y})p(y\,|\,\mathbf{X}_{0})$ (88) $\displaystyle\approx\sum_{\mathbf{X}_{0}}p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf {X}_{t},\mathbf{Y})p(y\,|\,\mathbf{X}_{0}),$ (89) $$

where $q(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$ is the true denoising distribution. This results in the following update step:

$$ $\displaystyle\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y},y)\propto\nabla_{\mathbf{X}_{t}}\log\left(\mathbb{E}_{p_{\theta}( \mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})}p(y\,|\,\mathbf{X}_{0})\right) \mathbf{X}_{t-1}+\log\mathbf{P}_{\theta}(\mathbf{X}_{t-1}\,|\,\mathbf{X}_{t}, \mathbf{Y}).$ (90) $$

Summing over all possible graphs $\mathbf{X}_{0}$ is prohibitive, however. Instead, we could sample $\mathbf{X}_{0}$ from $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$ with the Gumbel-Softmax trick [] and evaluate $\log p(y\,|\,\mathbf{X}_{0})$. As long as $\log p(y\,|\,\mathbf{X}_{0})$ is differentiable, we can then just use automatic differentiation to get our estimate of $\nabla_{\mathbf{X}_{t}}\log\mathbf{P}_{\theta}(y\,|\,\mathbf{X}_{t},\mathbf{Y})$. Another, more simplified approach that avoids sampling from $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$ is to relax the definition of the likelihood function to directly condition on the continuous-valued probability vector $\mathbf{P}_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$ instead of the discrete-valued $\mathbf{X}_{0}$. For simplicity, in the next section, we adopt this approach, but the full method with Gumbel-Softmax is not significantly more difficult to implement.

### D.1 Toy Synthesisability Model: Controlling Atom Economy

One obvious use-case for posthoc conditioning in the retrosynthesis context is to increase the probability of the generated reactants being synthesisable, using some pre-trained synthesisability model. Synthesisability of the generated reactants is crucial in the multi-step retrosynthesis context since otherwise, the generated branch of the search path is a dead end. The idea is that if we can improve the probability of synthesisability of the generated reactants, this should result in a much smaller search space for the multi-step model. To showcase the idea, we use a toy synthesisability model based on total the count of atoms in the reactants. The size of the reactants is indeed a feature used in synthesisability models [], with the motivation that smaller reactants are more likely to be synthesisable than large and complex ones. The total count of atoms in a reaction is also referred to as the *atom economy* [] of the reaction. It is the efficiency of a chemical process in converting all involved atoms into the desired products and is a key factor in creating synthesis routes that are efficient with respect to the raw materials. The ability to control the atom economy of a single-step retrosynthesis model thus has the potential to improve the efficiency of synthesis routes, in contrast to simply finding correct synthesis routes.

The model of synthesizability that we use is

$$ $\displaystyle\tilde{\mathbf{X}}_{0}$ $\displaystyle=p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y}),$ (91) $\displaystyle p(y=\textrm{synthesizable}\,|\,\tilde{\mathbf{X}}_{0})$ $\displaystyle=\sigma(\frac{\sum_{i\in S}\tilde{\mathbf{X}}_{0,i,d}-a}{b})^{ \gamma},$ (92) $$

where $S$ is the set of non-atom-mapped nodes, $a$, $b$ and $\gamma$ are constants that define the synthesizability model and $d$ refers to the dummy node index. The intuition is that the more nodes are classified as dummy nodes (non-atoms), the fewer atoms we have in total, leaving the atom economy higher. Note that $\sum_{i\in S}\tilde{\mathbf{X}}_{0,i,d}$ is the expected amount of dummy nodes from $p_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{t},\mathbf{Y})$. We set $a$ to half the amount of dummy nodes and $b$ to one-quarter of the amount of dummy nodes. It turns out that this leaves $\gamma$ as a useful parameter to tune the sharpness of the conditioning. The gradient estimate is then given by

$$ $\displaystyle\nabla_{\mathbf{X}_{t}}\log\mathbf{P}_{\theta}(y\,|\,\mathbf{X}_{ t},\mathbf{Y})=\gamma\nabla_{\mathbf{X}_{t}}\log\sigma(\frac{\sum_{i\in S} \tilde{\mathbf{X}}_{0,i,d}-a}{b}),$ (93) $$

which can be directly calculated with automatic differentiation. The full algorithm is detailed in Alg. 3. Figure A8 shows the effect of increasing $\gamma$ on an example product.

Figure: Algorithm 3 Sampling with atom-count guidance

Figure: Figure A8: The effect of increasing the $\gamma$ parameter. With negative $\gamma$, the model outputs large reactants, and with positive $\gamma$, the model outputs smaller ones, allowing for fine-grained control over the *atom economy* of the generated step.
Refer to caption: x11.png

### D.2 A Case Study on Interactive Generation With Inpainting

We replicate BHC’s green synthesis of Ibuprofen hypothetically using our model interactively with inpainting. Fig. A9 compares the output of our model to the ground truth synthesis. Below we explain the synthesis steps in detail:

- 1.
The retrosynthesis path starts with carbonylation, a simple and well-known synthetic reaction which adds a ‘CO’ structure to a compound. The practitioner tries basic generation but then notices that the suggested reactant is not promising. They then suggest a partial structure of the reactants which could lead to a more sensible path. Our retrosynthetic model can complete the reactants with this information.
- 2.
Next, our model proposes hydrogenation, which is a sensible suggestion in this case. The data does not consider explicitly individual H2 molecules, meaning that they are inferred from the context.
- 3.
In the third step, regular generation again seems off, but the practitioner notices that an acylation reaction (a reaction with the C(=O)C group) might lead to reactants that are readily available. The model is able to complete the rest of the reaction after knowing that C(=O)C is present. These steps match the synthesis plan proposed by BHC.

Figure: Figure A9: Replicating BHC’s green synthesis of Ibuprofen using our model interactively, and comparing to the known synthesis path.
Refer to caption: x12.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to  the weaknesses.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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$.''*?
soundness: 3
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.
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.
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)$.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
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?
soundness: 3
strengths: Aligning initial and target to generate efficiently makes sense and the results are promising compared to unaligned version.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)