# Alignment is Key for Applying
Diffusion Models to Retrosynthesis

Najwa Laabid*
Aalto University
<EMAIL_ADDRESS>
&Severi Rissanen*
Aalto University
<EMAIL_ADDRESS>
Markus Heinonen
Aalto University
&Arno Solin
Aalto University
&Vikas Garg
Aalto University, YaiYai, Ltd.

###### Abstract

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

ReactantsProduct Atom-mapping-aligned denoisers induce aligned permutation
equivariance

Aligned permutation equivariance

Figure 1:  We use atom-mapping information (visualized by highlighting atoms
with matching colours) as an input to the neural network denoiser to
generalize permutation equivariance to _aligned_ permutation equivariance
between reactant and product graphs.

##  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 [8]. Approaches to retrosynthesis
have steadily moved from relying on explicit expert-curated rules (template-
based [9]), and using simplifying assumptions like a single reaction center
(synthon-based [45, 48]), to ultimately adopting purely data-driven machine-
learning models (template-free [57, 21]). 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 [22], 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 [49]. Another advantage is the possibility to
add additional conditions even post-training [49, 5, 28], potentially allowing
us to specify substructures in the generated reactants [33], enforce desirable
properties [5] like synthesizability, or sample from nuanced regions of the
learned distribution [15] 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 [34], or optimizing the atom economy [53] 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 [29]. Template-based
models depend on the availability and quality of hard-coded chemical rules
[43, 61]. Synthon-based models are limited by their definition of a reaction
center, which does not necessarily hold for complex reactions [62, 45, 58].
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 [57, 44]. Efforts to bridge the gap
between the template-based and template-free paradigms include methods
investigating pretraining [65, 24]. Recently, Igashov et al. [21] proposed a
Schrödinger-bridge inspired method for retrosynthesis sharing similarities
with diffusion models. In another concurrent work [58], 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. [65]. 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. [65] 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_ [60], 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 [35, 56, 19, 20]. One motivation is that, in
diffusion models, permutation equivariant neural nets induce permutation
invariant distributions [35], meaning that different permutations of the same
graph have the same probabilities under the model. Recently, Yan et al. [63]
showed that relaxing permutation equivariance in graph diffusion using
absolute positional encodings can improve performance empirically.

##  3 Methods

ReactantsProduct$\mathbf{X}_{0}$$\mathbf{X}_{t}$$\mathbf{X}_{T}$$\mathbf{Y}$

$q(\mathbf{X}_{1}\\!\,|\,\\!\mathbf{X}_{0})$

$p_{\theta}(\mathbf{X}_{0}\\!\,|\,\\!\mathbf{X}_{1},\\!\mathbf{Y})$

$q(\mathbf{X}_{t}\\!\,|\,\\!\mathbf{X}_{t-1})$

$p_{\theta}(\mathbf{X}_{t-1}\\!\,|\,\\!\mathbf{X}_{t},\\!\mathbf{Y})$

$q(\mathbf{X}_{t+1}\\!\,|\,\\!\mathbf{X}_{t})$

$p_{\theta}(\mathbf{X}_{t}\\!\,|\,\\!\mathbf{X}_{t+1},\\!\mathbf{Y})$

$q(\mathbf{X}_{T}\\!\,|\,\\!\mathbf{X}_{T-1})$

$p_{\theta}(\mathbf{X}_{T-1}\\!\,|\,\\!\mathbf{X}_{T},\\!\mathbf{Y})$

$=\sum_{\mathbf{X}_{0}}q(\mathbf{X}_{0}\\!\,|\,\\!\mathbf{X}_{T},\\!\mathbf{X}_{0})p_{\theta}(\mathbf{X}_{0}\\!\,|\,\\!\mathbf{X}_{T},\\!\mathbf{Y},\\!\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$

Figure 2: An overview of our approach to retrosynthetic graph diffusion. We
adopt _absorbing state diffusion_ [1], 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.

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, [16]) 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 [47, 19, 1] to graphs, following Vignac et al. [56].

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 [17, 1],
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. [1], 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. [1].

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

$-\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.

Algorithm 1 Loss calculation

Input: product $\mathbf{Y}$, reactant $\mathbf{X}_{0}$, and optional
permutation matrix $\mathbf{P}^{\mathbf{X}\rightarrow\mathbf{Y}}$ for
alignment

$t\sim\mathrm{Uniform}(\\{0,\dots,T\\})$

$\mathbf{X}_{t}\sim q(\mathbf{X}_{t}\,|\,\mathbf{X}_{0})$

$\tilde{\mathbf{X}}_{0}=D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$

Return Cross-Entropy($\mathbf{X}_{0},\tilde{\mathbf{X}}_{0}$)

Algorithm 2 Sampling

Input: product $\mathbf{Y}$

Choose (for alignment):
$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\in\mathbb{R}^{N_{\mathbf{X}}\times
N_{\mathbf{Y}}}$

$\mathbf{X}_{T}\propto p(\mathbf{X}_{T})$

for $t=T$ to $1$ do

$\tilde{\mathbf{X}}_{0}=D_{\theta}(\mathbf{X}_{t},\mathbf{Y},\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})$

$\mathbf{X}_{t-1}^{i}\sim\sum_{k}q(\mathbf{X}_{t-1}^{i}\,|\,\mathbf{X}_{t}^{i},\mathbf{X}_{0}^{i})\tilde{\mathbf{X}}_{0}^{i}$

Return $\mathbf{X}_{0}$

###  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})$ [35, 56, 19].

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

Chemical reactions often induce limited changes to the precursor molecules
[65]. Previous approaches to retrosynthesis capitalize on this observation and
propose similarity [7], local graph edits [4], or product and reactants
alignment [65] 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}$.

Reactant

Product

$D_{\theta}(\mathbf{X}_{t},\\!\mathbf{Y})^{\mathcal{N}}$

C

O

N

Perm.Eq.

GNN

$\mathbf{X}_{T}^{\mathcal{N}}$

C

O

N

$\perp$

$\mathbf{Y}^{\mathcal{N}}$

C

O

N

$D_{\theta}(\mathbf{X}_{t},\\!\mathbf{Y},\\!\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\mathcal{N}}$

C

O

N

Aligned

GNN

$\mathbf{X}_{T}^{\mathcal{N}}$

C

O

N

$\perp$

$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$

$\mathbf{Y}^{\mathcal{N}}$

C

O

N

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}}$.

###  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}}$.

###  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}$ 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. [21], 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 [33]) or enforcing a property of the generated sample
as an extra condition (i.e. guided inference [5]). We implement inpainting
similarly to [56], by fixing a portion of the generated reactants to a desired
structure. For property-guided inference, we derive a DPS-like [5] 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 [13], a concept also known as atom
economy [53].

Aligned DenoiserUnaligned DenoiserSample with AMProductSampleProductSample
from
$\tilde{p}_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y},P^{\mathbf{Y}\to\mathbf{X}})$Sample
from $\tilde{p}_{\theta}(\mathbf{X}_{0}\,|\,\mathbf{X}_{T},\mathbf{Y})$= blank
= absorbing = atom mapping

(a) Aligned vs unaligned denoiser

1251020501000%20%40%60%80%100%top-1top-3top-5top-100Sampling step
countTop-$k$Aligned modelPermutation equivariant model
(b) Quality of samples vs. sampling speed

Figure 4: Top-$k$ scores for sampling step counts $T$ for our PE-skip model
and a model with a standard permutation equivariant denoiser, using the first
10% of validation set reactions.

##  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. [10]. 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 [12] and used by [56].
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, [30]) as used by Maziarz et al.
[34]. 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 [65, 24], or a more relaxed evaluation
procedure [21] 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 [65] and PMSR [24], 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 [7] | 37.3 | 54.7 | 63.3 | 74.1 | 0.480
GLN [10] | 52.5 | 74.7 | 81.2 | 87.9 | 0.641
LocalRetro [4] | 52.6 | 76.0 | 84.4 | 90.6 | 0.650
Synthon |  G2G [45] | 48.9 | 67.6 | 72.5 | 75.5 | 0.582
GraphRetro [48] | 53.7 | 68.3 | 72.2 | 75.5 | 0.611
MEGAN [37] | 48.0 | 70.9 | 78.1 | 85.4 | 0.601
RetroDiff [58] | 52.6 | 71.2 | 81.0 | 83.3 | 0.629
|  SCROP [64] | 43.7 | 60.0 | 65.2 | 68.7 | 0.521
|  Tied Transformer [27] | 47.1 | 67.1 | 73.1 | 76.3 | 0.572
|  Aug. Transformer [52] | 48.3 | — | 73.4 | 77.4 | 0.569
|  GTA_aug [44] | 51.1 | 67.6 | 74.8 | 81.6 | 0.605
|  Graph2SMILES [54] | 52.9 | 66.5 | 70.0 | 72.9 | 0.597
|  Retroformer [57] | 52.9 | 68.2 | 72.5 | 76.4 | 0.608
|  DualTF_aug [51] | 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 Ours | DiffAlign-PE+skip | 54.7 | 73.3 | 77.8 | 81.1 | 0.639

#### Top-$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_ [53]. 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.

ProductHigh atom countguidanceLow atom countguidance (a) Atom guidance

ProductGenerate one reactant conditional on anotherGenerate conditional on
desired substructure

(b) Inpainting

Figure 5: (a) Atom count guidance lets us specify if reactants should have
many or few atoms, controlling the _atom economy_. (b) Examples of inpainting
with our model. Parts highlighted in yellow are fixed by a practitioner with
desired characteristics, and the diffusion model completes the reaction.

##  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[18] or
MaskGIT [3] 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 [42]. 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 [18, 26, 32, 50, 38, 46] 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 [55]. Responsible usage guidelines are crucial. No specific harmful
impacts are noted from this work.

## References

  * Austin et al. [2021] J. Austin, D. D. Johnson, J. Ho, D. Tarlow, and R. van den Berg.  Structured denoising diffusion models in discrete state-spaces.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 17981–17993. Curran Associates, Inc., 2021.
  * Boys et al. [2023] B. Boys, M. Girolami, J. Pidstrigach, S. Reich, A. Mosca, and O. D. Akyildiz.  Tweedie moment projected diffusions for inverse problems.  _arXiv preprint arXiv:2310.06721_ , 2023.
  * Chang et al. [2022] H. Chang, H. Zhang, L. Jiang, C. Liu, and W. T. Freeman.  Maskgit: Masked generative image transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11315–11325, 2022.
  * Chen and Jung [2021] S. Chen and Y. Jung.  Deep retrosynthetic reaction prediction using local reactivity and global attention.  _JACS Au_ , 1(10):1612–1620, 2021.
  * Chung et al. [2023] H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye.  Diffusion posterior sampling for general noisy inverse problems.  In _International Conference on Learning Representations_ , 2023.
  * Coley et al. [2017a] C. W. Coley, R. Barzilay, T. S. Jaakkola, W. H. Green, and K. F. Jensen.  Prediction of organic reaction outcomes using machine learning.  _ACS Central Science_ , 3(5):434–443, 2017a.
  * Coley et al. [2017b] C. W. Coley, L. Rogers, W. H. Green, and K. F. Jensen.  Computer-assisted retrosynthesis based on molecular similarity.  _ACS Central Science_ , 3(12):1237–1245, 2017b.
  * Corey and Cheng [1996] E. Corey and X.-M. Cheng.  The logic of chemical synthesis.  _Journal of the American Chemical Society_ , 118(43):10678–10678, 1996.
  * Corey and Wipke [1969] E. J. Corey and W. T. Wipke.  Computer-assisted design of complex organic syntheses: Pathways for molecular synthesis can be devised with a computer and equipment for graphical communication.  _Science_ , 166(3902):178–192, 1969.
  * Dai et al. [2019] H. Dai, C. Li, C. Coley, B. Dai, and L. Song.  Retrosynthesis prediction with conditional graph logic network.  In _Advances in Neural Information Processing Systems_ , volume 32, pages 8872–8882. Curran Associates, Inc., 2019.
  * Dou and Song [2024] Z. Dou and Y. Song.  Diffusion posterior sampling for linear inverse problem solving: A filtering perspective.  In _International Conference on Learning Representations_ , 2024.
  * Dwivedi and Bresson [2021] V. P. Dwivedi and X. Bresson.  A generalization of transformer networks to graphs.  In _AAAI Workshop on Deep Learning on Graphs_ , 2021.
  * Ertl and Schuffenhauer [2009] P. Ertl and A. Schuffenhauer.  Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions.  _Journal of Cheminformatics_ , 1:1–11, 2009.
  * Finzi et al. [2023] M. A. Finzi, A. Boral, A. G. Wilson, F. Sha, and L. Zepeda-Nunez.  User-defined event sampling and uncertainty quantification in diffusion models for physical dynamical systems.  In _International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 10136–10152. PMLR, 2023\.
  * Garipov et al. [2023] T. Garipov, S. De Peuter, G. Yang, V. Garg, S. Kaski, and T. Jaakkola.  Compositional sculpting of iterative generative processes.  In _Advances in Neural Information Processing Systems_ , volume 36, pages 12665–12702. Curran Associates, Inc., 2023.
  * Ho et al. [2020] J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 6840–6851. Curran Associates, Inc., 2020.
  * Hoogeboom et al. [2021] E. Hoogeboom, D. Nielsen, P. Jaini, P. Forré, and M. Welling.  Argmax flows and multinomial diffusion: Learning categorical distributions.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 12454–12465. Curran Associates, Inc., 2021.
  * Hoogeboom et al. [2022a] E. Hoogeboom, A. A. Gritsenko, J. Bastings, B. Poole, R. van den Berg, and T. Salimans.  Autoregressive diffusion models.  In _International Conference on Learning Representations_ , 2022a.
  * Hoogeboom et al. [2022b] E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling.  Equivariant diffusion for molecule generation in 3D.  In _International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 8867–8887. PMLR, 2022b.
  * Huang et al. [2022] H. Huang, L. Sun, B. Du, Y. Fu, and W. Lv.  Graphgdp: Generative diffusion processes for permutation invariant graph generation.  In _2022 IEEE International Conference on Data Mining (ICDM)_ , pages 201–210. IEEE Computer Society, 2022.
  * Igashov et al. [2024a] I. Igashov, A. Schneuing, M. Segler, M. M. Bronstein, and B. Correia.  RetroBridge: Modeling retrosynthesis with Markov bridges.  In _Conference on Learning Representations_ , 2024a.
  * Igashov et al. [2024b] I. Igashov, H. Stärk, C. Vignac, A. Schneuing, V. G. Satorras, P. Frossard, M. Welling, M. Bronstein, and B. Correia.  Equivariant 3d-conditional diffusion model for molecular linker design.  _Nature Machine Intelligence_ , pages 1–11, 2024b.
  * Jang et al. [2016] E. Jang, S. Gu, and B. Poole.  Categorical reparameterization with gumbel-softmax.  In _International Conference on Learning Representations_ , 2016.
  * Jiang et al. [2023] Y. Jiang, W. Ying, F. Wu, Z. Huang, K. Kuang, and Z. Wang.  Learning chemical rules of retrosynthesis with pre-training.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, pages 5113–5121, 2023.
  * Jin et al. [2017] W. Jin, C. Coley, R. Barzilay, and T. Jaakkola.  Predicting organic reaction outcomes with Weisfeiler-Lehman network.  In _Advances in Neural Information Processing Systems_ , volume 30, pages 2607–2616. Curran Associates, Inc., 2017.
  * Karras et al. [2022] T. Karras, M. Aittala, T. Aila, and S. Laine.  Elucidating the design space of diffusion-based generative models.  In _Advances in Neural Information Processing Systems_ , volume 35, pages 26565–26577. Curran Associates, Inc., 2022.
  * Kim et al. [2021] E. Kim, D. Lee, Y. Kwon, M. S. Park, and Y.-S. Choi.  Valid, plausible, and diverse retrosynthesis using tied two-way transformers with latent variables.  _Journal of Chemical Information and Modeling_ , 61(1):123–133, 2021.
  * Li et al. [2022] X. Li, J. Thickstun, I. Gulrajani, P. S. Liang, and T. B. Hashimoto.  Diffusion-lm improves controllable text generation.  In _Advances in Neural Information Processing Systems_ , volume 35, pages 4328–4343. Curran Associates, Inc., 2022.
  * Liu et al. [2023] S. Liu, Z. Tu, M. Xu, Z. Zhang, L. Lin, R. Ying, J. Tang, P. Zhao, and D. Wu.  FusionRetro: Molecule representation fusion via in-context learning for retrosynthetic planning.  In _International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 22028–22041. PMLR, 2023\.
  * Liu et al. [2009] T.-Y. Liu et al.  Learning to rank for information retrieval.  _Foundations and Trends® in Information Retrieval_ , 3(3):225–331, 2009.
  * Lowe [2012] D. M. Lowe.  _Extraction of Chemical Structures and Reactions from the Literature_.  PhD thesis, University of Cambridge, 2012.
  * Lu et al. [2022] C. Lu, Y. Zhou, F. Bao, J. Chen, C. LI, and J. Zhu.  Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.  In _Advances in Neural Information Processing Systems_ , volume 35, pages 5775–5787. Curran Associates, Inc., 2022.
  * Lugmayr et al. [2022] A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, and L. Van Gool.  Repaint: Inpainting using denoising diffusion probabilistic models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11461–11471, 2022.
  * Maziarz et al. [2023] K. Maziarz, A. Tripp, G. Liu, M. Stanley, S. Xie, P. Gaiński, P. Seidl, and M. Segler.  Re-evaluating retrosynthesis algorithms with syntheseus.  In _NeurIPS 2023 AI for Science Workshop_ , 2023.
  * Niu et al. [2020] C. Niu, Y. Song, J. Song, S. Zhao, A. Grover, and S. Ermon.  Permutation invariant graph generation via score-based generative modeling.  In _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , volume 108 of _Proceedings of Machine Learning Research_ , pages 4474–4484. PMLR, 2020.
  * Peng et al. [2024] X. Peng, Z. Zheng, W. Dai, N. Xiao, C. Li, J. Zou, and H. Xiong.  Improving diffusion models for inverse problems using optimal posterior covariance.  _arXiv preprint arXiv:2402.02149_ , 2024.
  * Sacha et al. [2021] M. Sacha, M. Błaz, P. Byrski, P. Dabrowski-Tumanski, M. Chrominski, R. Loska, P. Włodarczyk-Pruszynski, and S. Jastrzebski.  Molecule edit graph attention network: modeling chemical reactions as sequences of graph edits.  _Journal of Chemical Information and Modeling_ , 61(7):3273–3284, 2021.
  * Sauer et al. [2023] A. Sauer, D. Lorenz, A. Blattmann, and R. Rombach.  Adversarial diffusion distillation.  _arXiv preprint arXiv:2311.17042_ , 2023.
  * Schneider et al. [2016] N. Schneider, N. Stiefl, and G. A. Landrum.  What’s what: The (nearly) definitive guide to reaction role assignment.  _Journal of Chemical Information and Modeling_ , 56(12):2336–2346, 2016.
  * Schwaller et al. [2018] P. Schwaller, T. Gaudin, D. Lanyi, C. Bekas, and T. Laino.  “Found in Translation”: predicting outcomes of complex organic chemistry reactions using neural sequence-to-sequence models.  _Chemical Science_ , 9(28):6091–6098, 2018.
  * Schwaller et al. [2019] P. Schwaller, T. Laino, T. Gaudin, P. Bolgar, C. A. Hunter, C. Bekas, and A. A. Lee.  Molecular transformer: A model for uncertainty-calibrated chemical reaction prediction.  _ACS Central Science_ , 5(9):1572–1583, 2019\.
  * Schwaller et al. [2021] P. Schwaller, B. Hoover, J.-L. Reymond, H. Strobelt, and T. Laino.  Extraction of organic chemistry grammar from unsupervised learning of chemical reactions.  _Science Advances_ , 7(15), 2021.
  * Segler and Waller [2017] M. H. Segler and M. P. Waller.  Neural-symbolic machine learning for retrosynthesis and reaction prediction.  _Chemistry–A European Journal_ , 23(25):5966–5971, 2017.
  * Seo et al. [2021] S.-W. Seo, Y. Y. Song, J. Y. Yang, S. Bae, H. Lee, J. Shin, S. J. Hwang, and E. Yang.  GTA: Graph truncated attention for retrosynthesis.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 35(1):531–539, 2021.
  * Shi et al. [2020] C. Shi, M. Xu, H. Guo, M. Zhang, and J. Tang.  A graph to graphs framework for retrosynthesis prediction.  In _International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 8818–8827. PMLR, 2020\.
  * Shih et al. [2023] A. Shih, S. Belkhale, S. Ermon, D. Sadigh, and N. Anari.  Parallel sampling of diffusion models.  In _Advances in Neural Information Processing Systems_ , pages 4263–4276. Curran Associates, Inc., 2023.
  * Sohl-Dickstein et al. [2015] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pages 2256–2265. PMLR, 2015\.
  * Somnath et al. [2021] V. R. Somnath, C. Bunne, C. Coley, A. Krause, and R. Barzilay.  Learning graph models for retrosynthesis prediction.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 9405–9415. Curran Associates, Inc., 2021.
  * Song et al. [2021] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2021.
  * Song et al. [2023] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever.  Consistency models.  In _International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 32211–32252. PMLR, 2023\.
  * Sun et al. [2021] R. Sun, H. Dai, L. Li, S. Kearnes, and B. Dai.  Towards understanding retrosynthesis by energy-based models.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 10186–10194. Curran Associates, Inc., 2021.
  * Tetko et al. [2020] I. V. Tetko, P. Karpov, R. Van Deursen, and G. Godin.  State-of-the-art augmented NLP transformer models for direct and single-step retrosynthesis.  _Nature Communications_ , 11(1):5575, 2020.
  * Trost [1991] B. M. Trost.  The atom economy—a search for synthetic efficiency.  _Science_ , 254(5037):1471–1477, 1991.
  * Tu and Coley [2022] Z. Tu and C. W. Coley.  Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction.  _Journal of Chemical Information and Modeling_ , 62(15):3503–3513, 2022.
  * Urbina et al. [2022] F. Urbina, F. Lentzos, C. Invernizzi, and S. Ekins.  Dual use of artificial-intelligence-powered drug discovery.  _Nature Machine Intelligence_ , 4(3):189–191, 2022.
  * Vignac et al. [2023] C. Vignac, I. Krawczuk, A. Siraudin, B. Wang, V. Cevher, and P. Frossard.  DiGress: Discrete denoising diffusion for graph generation.  In _International Conference on Learning Representations_ , 2023.
  * Wan et al. [2022] Y. Wan, C.-Y. Hsieh, B. Liao, and S. Zhang.  Retroformer: Pushing the limits of end-to-end retrosynthesis transformer.  In _International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 22475–22490. PMLR, 2022\.
  * Wang et al. [2023] Y. Wang, Y. Song, M. Xu, R. Wang, H. Zhou, and W. Ma.  RetroDiff: Retrosynthesis as multi-stage distribution interpolation.  _arXiv preprint arXiv:2311.14077_ , 2023.
  * Wu et al. [2023] L. Wu, B. Trippe, C. Naesseth, D. Blei, and J. P. Cunningham.  Practical and asymptotically exact conditional sampling in diffusion models.  In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, _Advances in Neural Information Processing Systems_ , volume 36, pages 31372–31403. Curran Associates, Inc., 2023.
  * Wu et al. [2021] Z. Wu, S. Pan, F. Chen, G. Long, C. Zhang, and P. S. Yu.  A comprehensive survey on graph neural networks.  _IEEE Transactions on Neural Networks and Learning Systems_ , 32(1):4–24, 2021.
  * Xie et al. [2023] S. Xie, R. Yan, J. Guo, Y. Xia, L. Wu, and T. Qin.  Retrosynthesis prediction with local template retrieval.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, pages 5330–5338, 2023.
  * Yan et al. [2020] C. Yan, Q. Ding, P. Zhao, S. Zheng, J. YANG, Y. Yu, and J. Huang.  RetroXpert: Decompose retrosynthesis prediction like a chemist.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 11248–11258. Curran Associates, Inc., 2020.
  * Yan et al. [2023] Q. Yan, Z. Liang, Y. Song, R. Liao, and L. Wang.  SwinGNN: Rethinking permutation invariance in diffusion models for graph generation.  _arXiv preprint arXiv:2307.01646_ , 2023.
  * Zheng et al. [2019] S. Zheng, J. Rao, Z. Zhang, J. Xu, and Y. Yang.  Predicting retrosynthetic reactions using self-corrected transformer neural networks.  _Journal of Chemical Information and Modeling_ , 60(1):47–55, 2019.
  * Zhong et al. [2022] Z. Zhong, J. Song, Z. Feng, T. Liu, L. Jia, S. Yao, M. Wu, T. Hou, and M. Song.  Root-aligned SMILES: a tight representation for chemical reaction prediction.  _Chemical Science_ , 13(31):9023–9034, 2022.

## 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] [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. [1], 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}}$ 111The 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. [17]

$\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. [56] 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. [1],
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 ith 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 [16]. Austin et al. [1] 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. [56], 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.

ReactantsProductAtom mapping from the nodes of $\mathbf{Y}$ to the nodes of
$\mathbf{X}$Atom mapping from the edges of $\mathbf{Y}$ to the edges of
$\mathbf{X}$

C

O

N

F

C

O

N

F

C

O

N

F

$\mathbf{X}^{\mathcal{N}}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{\mathcal{N}}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$$\mathbf{Y}^{\mathcal{N}}$
Non-atom-
---
mapped atoms
$N_{\mathbf{X}}$$N_{\mathbf{Y}}$$\mathbf{X}^{E}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{E}(\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\top}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}$$\mathbf{Y}^{E}$$(\mathbf{P}^{\mathbf{Y}\to\mathbf{X}})^{\top}$
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.

##  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}$.

∎

###  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. ∎

###  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.

Matched positional
---
encodings
Neural net predicts
---
the difference to
the product graph
Align the graphs
---
in the input
$\mathbf{X}^{E}$$\mathbf{Y}^{E}$$\mathbf{X}^{\mathcal{N}}$$\mathbf{Y}^{\mathcal{N}}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\boldsymbol{\varphi}$$\boldsymbol{\varphi}$GNN$D_{\theta}(\mathbf{X},\mathbf{Y})^{\mathcal{N}}$$D_{\theta}(\mathbf{X},\mathbf{Y})^{E}$$\mathbf{X}^{E}$$\mathbf{Y}^{E}$$\mathbf{X}^{\mathcal{N}}$$\mathbf{Y}^{\mathcal{N}}$GNN$D_{\theta}(\mathbf{X},\mathbf{Y})^{\mathcal{N}}$$D_{\theta}(\mathbf{X},\mathbf{Y})^{E}$$\mathbf{P}^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{\mathcal{N}}$$P^{\mathbf{Y}\to\mathbf{X}}\mathbf{Y}^{E}(P^{\mathbf{Y}\to\mathbf{X}})^{\top}$$\mathbf{X}^{E}$$\mathbf{Y}^{E}$$\mathbf{X}^{\mathcal{N}}$$\mathbf{Y}^{\mathcal{N}}$GNN$D_{\theta}(\mathbf{X},\mathbf{Y})^{\mathcal{N}}$$D_{\theta}(\mathbf{X},\mathbf{Y})^{E}$
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.

#### 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}$ 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}$ and $\mathbf{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. [56] 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 [31]. 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 [31] | 1 808 938 | Dai et al. [10]
Stereo | Schwaller et al. [40] | 1 002 970 | Schwaller et al. [41]
MIT | Jin et al. [25] | 479 035 | -
50k | Schneider et al. [39] | 50 016 | Dai et al. [10]

#### 15k

is proposed by Coley et al. [6]. 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. [39]. 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. [25]. 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. [40]. 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. [10]. 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 [49], 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$ 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 [12], based on the
implementation of Vignac et al. [56] with additional graph-level level
features added to the input of the model. See Vignac et al. [56] 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. [56], 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 [21], 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- trained |  RSMILES [65] | 56.3 | 79.2 | 86.2 | 91.0 | 0.680
PMSR [24] | 62.0 | 78.4 | 82.9 | 86.8 | 0.704
Temp. |  Retrosym [7] | 37.3 | 54.7 | 63.3 | 74.1 | 0.480
GLN [10] | 52.5 | 74.7 | 81.2 | 87.9 | 0.641
LocalRetro [4] | 52.6 | 76.0 | 84.4 | 90.6 | 0.650
Synthon |  GraphRetro [48] | 53.7 | 68.3 | 72.2 | 75.5 | 0.611
RetroDiff [58] | 52.6 | 71.2 | 81.0 | 83.3 | 0.629
MEGAN [37] | 48.0 | 70.9 | 78.1 | 85.4 | 0.601
G2G [45] | 48.9 | 67.6 | 72.5 | 75.5 | 0.582
Template-free |  SCROP [64] | 43.7 | 60.0 | 65.2 | 68.7 | 0.521
Tied Transformer [27] | 47.1 | 67.1 | 73.1 | 76.3 | 0.572
Aug. Transformer [52] | 48.3 | - | 73.4 | 77.4 | 0.569
Retrobridge (*) [21] | 50.3 | 74.0 | 80.3 | 85.1 | 0.622
GTA_aug [44] | 51.1 | 67.6 | 74.8 | 81.6 | 0.605
Graph2SMILES [54] | 52.9 | 66.5 | 70.0 | 72.9 | 0.597
Retroformer [57] | 52.9 | 68.2 | 72.5 | 76.4 | 0.608
DualTF_aug [51] | 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. [21]’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) 222This can be seen in
the code shared by Igashov et al. [21]:
https://github.com/igashov/RetroBridge.

#### Methods with pretraining

Zhong et al. [65] and Jiang et al. [24] 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 [56], 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 [49]. 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 [5, 11, 14, 2, 59, 36]. Here, we show an approach similar to [5]
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 [23] 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 [13], 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_ [53] 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.

Algorithm 3 Sampling with atom-count guidance

Input: Product $\mathbf{Y}$

Choose: $P^{\mathbf{Y}\to\mathbf{X}}\in\mathbb{R}^{N_{\mathbf{X}}\times
N_{\mathbf{Y}}}$

$\mathbf{X}_{T}\propto p(\mathbf{X}_{T})$

for $t=T$ to $1$ do

$\tilde{\mathbf{X}}_{0}=D_{\theta}(\mathbf{X}_{t},\mathbf{Y},P^{\mathbf{Y}\to\mathbf{X}})$
$\triangleright$ Denoising output

$\tilde{\mathbf{X}}_{t-1}^{i}=\sum_{k}q(\mathbf{X}_{t-1}^{i}\,|\,\mathbf{X}_{t}^{i},\mathbf{X}_{0}^{i})\tilde{\mathbf{X}}_{0}^{i}$
$\triangleright$ Regular reverse transition probabilities

$\mathbf{X}_{t-1}^{i}\sim\textrm{softmax}(\log\tilde{\mathbf{X}}_{t-1}+\gamma\nabla_{\mathbf{X}_{t}}\log\sigma(\frac{\sum_{i\in
S}\tilde{\mathbf{X}}_{0,i,d}-a}{b}))$ $\triangleright$ Renormalize

Return $\mathbf{X}_{0}$

ProductGenerated ReactantsIncreasing atom economyHigh negative guidance (gamma
< 0)High positive guidance (gamma > 0) 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.

###  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. 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. 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. 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.

Objective: synthesize= does not correspond to ground truth= the correct
reaction (possibly up to implicit$H_{2}$ not present in the data)= inpainted
formReal (retro)synthesis pathBasic generationWith inpainting1\.
Carbonylation2\. Hydrogenation3\. Friedel-Crafts AcylationKnowledge about the
desired main reactantand a guess about side reactantsKnowledge about the type
ofreaction (acylation) Figure A9: Replicating BHC’s green synthesis of
Ibuprofen using our model interactively, and comparing to the known synthesis
path.