SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The backbone-first strategy simplifies the generation problem by splitting it into two steps: the generation of the protein backbone and of the protein sequence given the backbone structure. A backbone can be designed by refining a rationally constructed template [24] or by a structure-based generative model [90; 32; 15; 87; 50; 92; 47]. On the other hand, the determination of the residues that fit the backbone can be cast within the self-supervised learning paradigm [3; 11] or by minimizing some score that is a proxy for energy [60]. Despite the promise of backbone-first strategies, in the context of deep learning, it is generally preferable to solve problems in an end-to-end fashion. Decoupling the problem into separate steps can also prohibit the modeling of complex dependencies between backbone and sequence, as early decisions are not updated in light of new evidence. As discussed by Harteveld et al. [24], one failure mode of the backbone-first strategy is obtaining a secondary structure that is not physically realizable with natural amino acids. It is also worth mentioning that structure diffusion generative models can also be used for protein-protein docking, as shown by Ketata et al. [40].
Sequence and structure co-design entails learning to generate a full protein in an end-to-end fashion [51; 89]. In this paper we also adopt this strategy. As highlighted in the introduction, AbDiffuser differs from these previous works in the neural architecture and protein representation (APMixer), as well as by shaping the diffusion process to better suit the problem domain.
A relevant line of works has also focused on the re-generation of the complementarity-determining regions (CDRs) of antibodies, a subset of the heavy and light chains containing typically between 48 and 82 residues [10] of particular relevance to binding [34; 20; 43; 33; 53; 78; 44]. These constitute meaningful first steps towards the goal of full-atom antibody design that we consider in this work. Having the ability to regenerate the full antibody sequence and structure significantly enlarges the design space while also enabling the probing of properties beyond affinity, such as stability and immunogenicity. Furthermore, even in the context of affinity optimization, it can be beneficial to also redesign the framework and vernier regions, as the latter sometimes interact with the antigen and affect the CDR conformation [1; 70; 79].
B Antibody Diffusion
This section complements the discussion in Section 2 by describing in greater detail how we frame denoising diffusion.
B.1 Additional Background on Denoising Diffusion
Recall that the forward process  $q(X_{1},\dots,X_{T}|X_{0}) = q(X_{1}|X_{0})\Pi_{t=2}^{T}q(X_{t}|X_{t-1})$ , is Markovian. This Markovianity allows for the recovery of the reverse process
$$
q \left(X _ {t - 1} \mid X _ {t}, X _ {0}\right) = \frac {q \left(X _ {t} \mid X _ {t - 1}\right) q \left(X _ {t - 1} \mid X _ {0}\right)}{q \left(X _ {t} \mid X _ {0}\right)}.
$$
Thus, to generate new data, a neural network learns to approximate the true denoising process  $\hat{X}_0 = \phi(X_t, t)$ , which can be accomplished by minimizing the variational upper bound on the negative loglikelihood [84; 26; 5]:
$$
\begin{array}{l} L _ {\text {e l b o}} \left(X _ {0}\right) := \mathbb {E} _ {q \left(X _ {0}\right)} \left[ \underbrace {D _ {K L} \left(q \left(X _ {T} \mid X _ {0}\right) \| p \left(X _ {T}\right)\right)} _ {L _ {T} \left(X _ {0}\right)} \right. \tag {3} \\ + \sum_ {t = 2} ^ {T} \mathbb {E} _ {q (X _ {t} | X _ {0})} \underbrace {D _ {K L} (q (X _ {t - 1} | X _ {t} , X _ {0}) \| p _ {\theta} (X _ {t - 1} | X _ {t}))} _ {L _ {t} (X _ {0})} \\ \left. - \mathbb {E} _ {q (X _ {1} | X _ {0})} \underbrace {\log \left(p _ {\theta} \left(X _ {0} \mid X _ {1}\right)\right)} _ {L _ {1} (X _ {0})} \right] \\ \end{array}
$$
In practice, the model is trained to optimize each of the  $L_{t}(X_{0})$  loss terms and the  $L_{0}(X_{0})$  loss term individually, by sampling an arbitrary time step  $t \in \{1, T\}$ . Note that  $L_{T}(X_{0})$  approaches 0 for any data distribution if sufficiently many time steps are used. The latent variable distribution  $q$  is usually defined so that one can directly sample any step in the trajectory efficiently.
As discussed in Section 2, in our diffusion process, we factorize the posterior probability distribution over the atom positions and residue types:
$$
q \big (X _ {t} | X _ {t - 1} \big) = q \big (X _ {t} ^ {\mathrm {p o s}} | X _ {t - 1} ^ {\mathrm {p o s}} \big) q \big (X _ {t} ^ {\mathrm {r e s}} | X _ {t - 1} ^ {\mathrm {r e s}} \big).
$$
We discuss the constituent processes in the following two sections.
B.2 Equivariant Diffusion of Atom Positions
One of the most important features in the training of a diffusion model is the ability to efficiently sample any step in the forward process  $q(X_{t}^{\mathrm{pos}}|X_{0}^{\mathrm{pos}})$ . This is easily achieved if we use a Gaussian distribution  $q(X_{t}^{\mathrm{pos}}|X_{0}^{\mathrm{pos}}) = \mathcal{N}(\alpha_{t}X_{0}^{\mathrm{pos}},\sigma_{t}^{2}\Sigma)$ , where  $\alpha_{t}$  and  $\sigma_t^2$  are positive, scalar-valued functions of  $t$ , with  $\alpha_{1}\approx 1$  and  $\alpha_{T}\approx 0$ . A common way to define the noisig process is the variance-preserving parameterization [84; 26] with  $\alpha_{t} = \sqrt{1 - \sigma_{t}^{2}}$ . A forward process step in this general diffusion formulation can be expressed as  $q(X_{t}^{\mathrm{pos}}|X_{t - 1}^{\mathrm{pos}}) = \mathcal{N}(\alpha_{t|t - 1}X_0^{\mathrm{pos}},\sigma_{t|t - 1}^2\Sigma)$ , where  $\alpha_{t|t - 1} = \alpha_t / \alpha_{t - 1}$  and  $\sigma_{t|t - 1}^2 = \sigma_t^2 -\alpha_{t|t - 1}^2\sigma_{t - 1}^2$ . Similarly, a reverse process step  $q(X_{t - 1}^{\mathrm{pos}}|X_{t}^{\mathrm{pos}},X_{0}^{\mathrm{pos}})$  can then be expressed as the following Gaussian:
$$
\mathcal {N} \left(\frac {\alpha_ {t | t - 1} \sigma_ {t - 1} ^ {2} X _ {t} ^ {\mathrm {p o s}} + \alpha_ {t - 1} \sigma_ {t | t - 1} ^ {2} X _ {0} ^ {\mathrm {p o s}}}{\sigma_ {t} ^ {2}}, \frac {\sigma_ {t | t - 1} ^ {2} \sigma_ {t - 1} ^ {2}}{\sigma_ {t} ^ {2}} \Sigma\right).
$$
Similarly to Watson et al. [90], we can reduce the temperature  $\tau$  of the Brownian motion of the backward process by scaling the variance of this Gaussian distribution by  $\tau$ .
Kingma et al. [41] showed that the training objective for the Gaussian distribution can be simplified to
$$
L _ {t} \left(X _ {0} ^ {\text {p o s}}\right) = \frac {1}{2} \left(\operatorname {S N R} (t - 1) - \operatorname {S N R} (t)\right) \| \hat {X} _ {0} ^ {\text {p o s}} - X _ {0} ^ {\text {p o s}} \| _ {\Sigma^ {- 1}} ^ {2},
$$
where the signal-to-noise ratio is  $\mathrm{SNR}(t) = (\alpha_t / \sigma_t)^2$  and  $\| \cdot \|_{\Sigma^{-1}}^2 = \operatorname{tr}\big((\cdot)^T\Sigma^{-1}(\cdot)\big)$ .
It is quite common in the literature to choose a slightly different optimization objective. In particular, it has been found that predicting the noise  $\hat{\epsilon} = \phi(X_t^{\mathrm{pos}}, t)$  that has been added  $X_t^{\mathrm{pos}} = \alpha_t X_0^{\mathrm{pos}} + \sigma_t \epsilon$ ,
$\epsilon \sim \mathcal{N}(0,\Sigma)$  and using an unweighted loss  $L_{t}(X_{0}^{\mathrm{pos}}) = \frac{1}{2}\|\hat{\epsilon} - \epsilon\|_{\Sigma^{-1}}^{2}$  improves model training [26; 41; 28]. A possible explanation is that this parameterization makes it easier for the model to minimize the loss at time steps close to  $T$ . We found that a similar effect can be obtained by re-weighing the mean prediction loss (Appendix C):
$$
L _ {t} \left(X _ {0} ^ {\text {p o s}}\right) = \frac {1}{2} \operatorname {S N R} (t) \| \hat {X} _ {0} ^ {\text {p o s}} - X _ {0} ^ {\text {p o s}} \| _ {\Sigma^ {- 1}} ^ {2}, \tag {4}
$$
which empirically renders the mean  $\hat{X}_0^{\mathrm{pos}} = \phi (X_t^{\mathrm{pos}},t)$  and error  $\hat{\epsilon} = \phi (X_t^{\mathrm{pos}},t)$  prediction alternatives perform comparably. We rely on equation 4 throughout our experiments, since predicting  $\hat{X}_0^{\mathrm{pos}} = \phi (X_t^{\mathrm{pos}},t)$  allows for easier incorporation of known constraints, such as atom bond lengths (see Section 3.4). As shown in Appendix D,  $L_{1}(X_{0}^{\mathrm{pos}})$  can also be optimized using the same objective.
B.3 Discrete Residue-type Diffusion
To define the discrete diffusion process for the residue types we adopt the formulation introduced by Austin et al. [5]. If we have a discrete random variable with  $k$  categories, represented as a one-hot vector  $X^{\mathrm{res}}$ , we define the forward process  $q(X_{t}^{\mathrm{res}}|X_{0}^{\mathrm{res}}) = \mathrm{Cat}(X_{0}^{\mathrm{res}}Q_{t})$ , where the transition matrix  $Q_{t} = \beta_{t}I + (1 - \beta_{t})Q$  is a linear interpolation between identity and target distribution  $Q$  (e.g., uniform distribution). Similarly to atom diffusion, one forward process step is  $q(X_{t}^{\mathrm{res}}|X_{t-1}^{\mathrm{res}}) = \mathrm{Cat}(X_{t}^{\mathrm{res}}Q_{t|t-1})$ , where  $Q_{t|t-1} = Q_{t-1}^{-1}Q_{t}$ . From this, a single reverse process step can be expressed as:
$$
q (X _ {t - 1} ^ {\text {r e s}} | X _ {t} ^ {\text {r e s}}, X _ {0} ^ {\text {r e s}}) = \operatorname {C a t} \left(\frac {X _ {t} ^ {\text {r e s}} Q _ {t | t - 1} ^ {T} \odot X _ {0} ^ {\text {r e s}} Q _ {t - 1}}{\mathcal {Z}}\right),
$$
for an appropriate normalizing constant  $\mathcal{Z}$ . In the discrete case, we implement temperature  $\tau$  scaling as the temperature of the softmax applied to the model output. This serves to sharpen the target categorical distribution  $\sim p^{1 / \tau}$ .
Similarly to the atom diffusion case, we use a simplified objective for the discrete diffusion [22; 7] which has been found to lead to better results:
$$
L _ {t} \left(X _ {0} ^ {\mathrm {r e s}}\right) = - \beta_ {t} \log p _ {\theta} \left(X _ {0} ^ {\mathrm {r e s}} \mid X _ {t} ^ {\mathrm {r e s}}\right)
$$
Hyperparameter  $\beta_{t}$  ensures that the loss weight is proportional to the fraction of unchanged classes and we set  $\beta_{t} = \alpha_{t}^{2}$  to keep the noise schedules for the residue-type diffusion and atom-position diffusion the same. We also note that now the  $L_{1}(X_{0})$  loss term is simply the cross-entropy of the model prediction  $L_{1}(X_{0}) = \log p_{\theta}(X_{0}|X_{1})$ . The derivation of our loss can be found in Appendix E.
C Re-weighted Optimization Objective for Atom Positions
As shown by Kingma et al. [41], the  $L_{t}$  term in the variational upper bound corresponds to
$$
L _ {t} \left(X _ {0} ^ {\text {p o s}}\right) = \frac {1}{2} \left(\operatorname {S N R} (t - 1) - \operatorname {S N R} (t)\right) \| \hat {X} _ {0} ^ {\text {p o s}} - X _ {0} ^ {\text {p o s}} \| _ {\Sigma^ {- 1}} ^ {2}.
$$
If we instead parameterize  $X_{t}^{\mathrm{pos}} = \alpha_{t}X_{0}^{\mathrm{pos}} + \sigma_{t}\epsilon$  we can directly re-write this as
$$
L _ {t} (\epsilon) = \frac {1}{2} \left(\frac {\operatorname {S N R} (t - 1)}{\operatorname {S N R} (t)} - 1\right) \| \hat {\epsilon} - \epsilon \| _ {\Sigma^ {- 1}} ^ {2}.
$$
Meanwhile, the simplified objective of Ho et al. [26] is instead defined as
$$
L _ {t} (\epsilon) = \frac {1}{2} \| \hat {\epsilon} - \epsilon \| _ {\Sigma^ {- 1}} ^ {2}.
$$
If we reverse the parameterization as  $\epsilon = \frac{X_t^{\mathrm{pos}} - \alpha_t X_0^{\mathrm{pos}}}{\sigma_t}$  and using the fact that  $\mathrm{SNR}(t) = \alpha_t^2 / \sigma_t^2$  this becomes
$$
L _ {t} (X _ {0} ^ {\mathrm {p o s}}) = \frac {1}{2} \left\| \frac {X _ {t} ^ {\mathrm {p o s}} - \alpha_ {t} \hat {X} _ {0} ^ {\mathrm {p o s}}}{\sigma_ {t}} - \frac {X _ {t} ^ {\mathrm {p o s}} - \alpha_ {t} X _ {0} ^ {\mathrm {p o s}}}{\sigma_ {t}} \right\| _ {\Sigma^ {- 1}} ^ {2} = \frac {1}{2} \operatorname {S N R} (t) \| \hat {X} _ {0} ^ {\mathrm {p o s}} - X _ {0} ^ {\mathrm {p o s}} \| _ {\Sigma^ {- 1}} ^ {2}.
$$
D Reconstruction Loss for Atom Positions
For the Gaussian atom position diffusion (Section B.2 we need to account for a few considerations in the reconstruction term  $L_{1}(X_{0})$  of the variational upper bound. Assuming that the data probability is constant,  $q(X_0^{\mathrm{pos}}|X_1^{\mathrm{pos}})\approx \mathcal{N}(\frac{X_1^{\mathrm{pos}}}{\alpha_1},\frac{\sigma_1^2}{\alpha_1^2}\Sigma)$ . Similarly to Hoogeboom et al. [28], we parameterize the mean using the model prediction instead:
$$
L _ {1} \left(X _ {0} ^ {\text {p o s}}\right) = \log \mathcal {Z} ^ {- 1} - \frac {1}{2} \operatorname {S N R} (1) \| X _ {0} ^ {\text {p o s}} - \hat {X} _ {0} ^ {\text {p o s}} \| _ {\Sigma^ {- 1}} ^ {2}.
$$
As highlighted by Hoogeboom et al. [28], this increases the quality of the results.
Isotropic distributions, such as Gaussians with appropriate covariance matrices, assign the same probability to every rotation of a vector in 3D, implying that the resulting atom diffusion process is equivariant w.r.t. orthogonal group O3. To additionally impose that the diffusion is translation-equivariant, we follow Hoogeboom et al. [28] and use a re-normalized multivariate Gaussian distribution, where the sample center of mass is fixed to zero:  $\sum_{i=1}^{n} X_i^{\mathrm{pos}} = 0$ . The latter has the same law as an ordinary Gaussian with the main difference that the normalizing constant is now  $\mathcal{Z} = \left(\sqrt{2\pi}\sigma\right)^{3(n-1)}$  for  $n$  atoms in 3D space.
It is clear to see that the O3 property is satisfied if we have an identity covariance. However, it also holds for non-isotropic covariance matrices. Let us define  $Z \sim \mathcal{N}(0,I)$  of dimension  $n \times 3$  and  $X^{\mathrm{pos}} = KZ$ , where  $\Sigma^{-1} = KK^{T}$ . The latter can be re-written as  $x_0 = \operatorname{vec}(X_0) = (I \otimes K)\operatorname{vec}(Z) = \Sigma^{1/2}z$ . We then note that for any 3x3 unitary matrix  $U$ , we have  $x_0^\top(U^\top \otimes I)\Sigma^{-1}(U \otimes I)x_0 = x_0^\top(U^\top \otimes I)(I \otimes L)(U \otimes I)x_0 = x_0^\top(U^\top U \otimes L)x_0 = x_0^\top(I \otimes L)x_0$ . In other words, since the covariance acts on the node dimension and the unitary matrix acts on the spatial dimension, they are orthogonal and do not affect each other.
E Simplified Discrete Diffusion Optimization Objective
We can obtain the simplified training objective by taking a slightly different bound on the negative log-likelihood to the one presented in Equation 3 [22]:
$$
\begin{array}{l} - \log \left(p _ {\theta} \left(X _ {0}\right)\right) \leq \mathbb {E} _ {X _ {1: T} \sim q \left(X _ {1: T} \mid X _ {0}\right)} \left[ - \log \left(\frac {p _ {\theta} \left(X _ {0 : T}\right)}{q \left(X _ {1 : T} \mid X _ {0}\right)}\right) \right] \\ = \mathbb {E} _ {q} \left[ - \log \left(\frac {p _ {\theta} \left(X _ {0 : T}\right)}{q \left(X _ {1 : T} \mid X _ {0}\right)}\right) \right] \\ = \mathbb {E} _ {q} \left[ - \log \left(p _ {\theta} \left(X _ {T}\right)\right) - \sum_ {t = 2} ^ {T} \log \left(\frac {p _ {\theta} \left(X _ {t - 1} \mid X _ {t}\right)}{q \left(X _ {t} \mid X _ {t - 1}\right)}\right) - \log \left(\frac {p _ {\theta} \left(X _ {0} \mid X _ {1}\right)}{q \left(X _ {1} \mid X _ {0}\right)}\right) \right] \\ = \mathbb {E} _ {q} \left[ - \log \left(p _ {\theta} \left(X _ {T}\right)\right) - \log \left(\frac {p _ {\theta} \left(X _ {0} \mid X _ {1}\right)}{q \left(X _ {1} \mid X _ {0}\right)}\right) \right. \\ \left. - \sum_ {t = 2} ^ {T} \log \left(\frac {q \left(X _ {t - 1} \mid X _ {t} , X _ {0}\right) p _ {\theta} \left(X _ {0} \mid X _ {t}\right)}{q \left(X _ {t - 1} \mid X _ {t} , X _ {0}\right)} \cdot \frac {q \left(X _ {t - 1} \mid X _ {0}\right)}{q \left(X _ {t} \mid X _ {0}\right)}\right) \right] \\ = \mathbb {E} _ {q} \left[ - \underbrace {\log \left(\frac {p _ {\theta} \left(X _ {T}\right)}{q \left(X _ {T} \mid X _ {0}\right)}\right)} _ {L _ {T} \left(X _ {0}\right)} - \sum_ {t = 2} ^ {T} \underbrace {\log \left(p _ {\theta} \left(X _ {0} \mid X _ {t}\right)\right)} _ {L _ {t} \left(X _ {0}\right)} - \underbrace {\log \left(p _ {\theta} \left(X _ {0} \mid X _ {1}\right)\right)} _ {L _ {1} \left(X _ {0}\right)} \right] \\ \end{array}
$$
In this formulation, compared to equation 3, the term  $L_{t} = D_{\mathrm{KL}} \left( q(X_{t-1}|X_{t}, X_{0}) \| p_{\theta}(X_{t-1}|X_{t}) \right)$  is replaced by
$$
L _ {t} = - \log \left(p _ {\theta} \left(X _ {0} | X _ {t}\right)\right)
$$
To form the discrete loss  $L_{t}(X_{0}^{\mathrm{res}})$ , we additionally weight each loss term  $L_{t}$  by  $\beta_{t}$  to match the loss weighting used for the atom positions (Appendix C). Note that when this objective is minimized, the original  $D_{\mathrm{KL}}$  term will also be at a minimum.
E.1 Non-Informative Priors Demand Higher Complexity Denoising Models
Denote by  $p$  the data distribution and by  $p_{\theta}$  the learned distribution, the latter of which is obtained by the push-forward  $p_{\theta}(f_{\theta}(Z)) = q(Z)$  of some prior measure  $q$  by function  $f_{\theta}$ . In the context of denoising diffusion,  $Z = X_{T}$ , while  $f_{\theta} = \phi(X_{T}, T)$  is the learned reverse process.
We prove the following general result which relates the informativeness of the prior with the quality of the generative model and the complexity of the function computed:
Theorem E.1. For any  $f_{\theta}$  that is an invertible equivariant function w.r.t. a subgroup  $G$  of the general-linear group  $GL(d,\mathbb{R})$  the following must hold:
$$
c _ {q} (f _ {\theta}) \geq W _ {t} (p, q) - W _ {t} (p, p _ {\theta}),
$$
where  $c_{q}(f) \coloneqq (\min_{g \in \mathbf{G}} \mathbf{E}_{Z \sim q}[\| f(gZ) - Z\|_{t}^{t}])^{1 / t}$  quantifies the expected complexity of the learned model under  $q$ ,  $W_{t}(p,p_{\theta})$  is the Wasserstein  $t$ -distance of our generative model to the data distribution, and  $W_{t}(p,q)$  is the distance between prior and posterior.
The bound asserts that one cannot learn a simple (in terms of  $c_{q}(f)$ ) generative model that fits the data well unless the prior has a small Wasserstein distance to the data distribution.
The complexity measure  $c_{q}(f)$  is particularly intuitive in the context of denoising diffusion. For the diffusion of atom positions with a SE(3) equivariant model  $\phi$ , we have
$$
c _ {q} (\phi) = \min _ {g \in \mathrm {S E} (3)} \mathbf {E} _ {X _ {T}} [ \| \phi (g X _ {T}, T) - X _ {T} \| _ {2} ^ {2} ],
$$
that corresponds to the expected amount of denoising that the model performs throughout the reverse process, discounting for rotations. By selecting an informative prior, we decrease  $W_{t}(p,q)$  and therefore reduce the amount of denoising our model needs to do.
Proof. We recall that the Wasserstein  $t$ -distance between two probability measures  $p, p'$  with finite  $t$ -moments is
$$
W _ {t} (p, p ^ {\prime}) := \left(\inf  _ {\gamma \in \Gamma (p, p ^ {\prime})} \mathbf {E} _ {(X, X ^ {\prime}) \sim \gamma} [ \| X - X ^ {\prime} \| _ {t} ]\right) ^ {1 / t},
$$
where the coupling  $\gamma$  is a joint measure on  $\mathbb{R}^d\times \mathbb{R}^d$  whose marginals satisfy  $\int_{\mathbb{R}^d}\gamma (X,X')dx' = p(X)$  and  $\int_{\mathbb{R}^d}\gamma (X,X')dx = p'(X')$ , whereas  $\Gamma (p,p')$  is the set of all couplings of  $p,p'$ .
We will first lower bound  $W_{1}(p_{\theta},q)$ . To achieve this, let  $X = f(Z)$  and  $X' = f_{\theta}(Z')$  and select  $\gamma^{*}(X,X') = 0$  when  $Z \neq Z'$  and  $\gamma^{*}(X,X') = q(Z)$  otherwise, which is a coupling because it satisfies  $\int_{X'} \gamma^{*}(X,X') dx' = q(Z) = p(X)$  and  $\int_{X} \gamma^{*}(X,X') dx = q(Z') = p_{\theta}(X')$ . Then, exploiting that the Wasserstein distance is a metric (and thus abides by the triangle inequality), we get
(by the definition of  $\gamma^{*}$  and invertibility of  $f$ )
$$
\begin{array}{l} W _ {t} \left(p _ {\theta}, q\right) \geq W _ {t} (p, q) - W _ {t} \left(p, p _ {\theta}\right) \quad (\text {b y}) \\ = W _ {t} (p, q) - \left(\inf  _ {\gamma \in \Gamma (p, p _ {\theta})} \int_ {X, X ^ {\prime}} \gamma (X, X ^ {\prime}) \| X - X ^ {\prime} \| _ {t} ^ {t} d x d x ^ {\prime}\right) ^ {1 / t} \\ \geq W _ {t} (p, q) - \left(\int_ {X, X ^ {\prime}} \gamma^ {*} (X, X ^ {\prime}) \| X - X ^ {\prime} \| _ {t} ^ {t} d x d x ^ {\prime}\right) ^ {1 / t} \\ = W _ {t} (p, q) - \left(\int_ {X} \left(\int_ {X ^ {\prime}} \gamma^ {*} (X, X ^ {\prime}) \| X - X ^ {\prime} \| _ {t} ^ {t} d x ^ {\prime}\right) d x\right) ^ {1 / t} \\ = W _ {t} (p, q) - \left(\int_ {X} p (X) \| X - f _ {\theta} \left(f ^ {- 1} (X)\right) \| _ {t} ^ {t} d x\right) ^ {1 / t} \\ = W _ {t} (p, q) - \left(\mathbf {E} _ {X \sim p} [ \| X - f _ {\theta} (f ^ {- 1} (X)) \| _ {t} ^ {t} ]\right) ^ {1 / t} \\ = W _ {t} (p, q) - \left(\mathbf {E} _ {Z \sim q} [ \| f (Z) - f _ {\theta} (f ^ {- 1} (f (Z))) \| _ {t} ^ {t} ]\right) ^ {1 / t} \\ = W _ {t} (p, q) - \left(\mathbf {E} _ {Z \sim q} [ \| f (Z) - f _ {\theta} (Z) \| _ {t} ^ {t} ]\right) ^ {1 / t} \tag {5} \\ \end{array}
$$
Next, we upper bound  $W_{t}(p_{\theta},q)$ . Let  $G$  be any subgroup of the general linear group  $\mathrm{GL}(d,\mathbb{R})$  for which  $f_{\theta}(Z)$  is equivariant and further denote by  $g\in G$  an element of the said group.
We fix  $\gamma^{\#}(X,Z) = q(Z)$  when  $X = f_{\theta}(gZ)$  and  $\gamma^{\#}(X,Z) = 0$ , otherwise. To see that the latter is a valid coupling notice that  $f_{\theta}g$  is bijective as the composition of two bijective functions, implying:
$$
\int \gamma^ {\#} (X, Z) d z = q \left(z _ {x}\right) \quad \text {a n d} \quad \int \gamma^ {\#} (X, Z) d x = q (Z),
$$
with  $z_{x} = (f_{\theta}\circ g)^{-1}(X)$  . Further,
$$
\begin{array}{l} q \left(z _ {x}\right) = p _ {\theta} \left(f _ {\theta} \left(g ^ {- 1} \left(f _ {\theta} ^ {- 1} (X)\right)\right)\right) \quad \text {(s i n c e} p _ {\theta} (X) = q \left(f _ {\theta} (Z)\right)) \\ = p _ {\theta} \left(g ^ {- 1} f _ {\theta} \left(f _ {\theta} ^ {- 1} (X)\right)\right) \quad (\text {b y}) \\ = p _ {\theta} \left(g ^ {- 1} X\right) \\ = p _ {\theta} (X), \quad (\text {b y}) \\ \end{array}
$$
as required for  $\gamma^{\#}$  to be a coupling. We continue by bounding the Wasserstein metric in light of  $\gamma^{\#}$ :
$$
\begin{array}{l} W _ {t} \left(p _ {\theta}, q\right) ^ {t} = \inf  _ {\gamma \in \Gamma \left(p _ {\theta}, q\right)} \int \gamma (X, Z) \| X - Z \| _ {t} ^ {t} d x d z \quad (\text {b y}) \\ \leq \int \gamma^ {\#} (X, Z) \| X - Z \| _ {t} ^ {t} d x d z \quad \text {(d u e t o t h e i n i f i m u m)} \\ = \int \left(\int \gamma^ {\#} (X, Z) \| X - Z \| _ {t} ^ {t} d x\right) d z \\ = \int q (Z) \| f _ {\theta} (g Z) - Z \| _ {t} ^ {t} d z \quad (\text {b y}) \\ = \mathbf {E} _ {Z \sim q} [ \| f _ {\theta} (g Z) - Z \| _ {t} ^ {t} ] \tag {6} \\ \end{array}
$$
Combining equation 5 and equation 6, we get
$$
\min  _ {g \in G} \mathbf {E} _ {Z \sim q} [ \| f _ {\theta} (g Z) - Z \| _ {t} ^ {t} ] \geq W _ {t} (p, q) - \left(\mathbf {E} _ {Z \sim q} [ \| f (Z) - f _ {\theta} (Z) \| _ {t} ^ {t} ]\right) ^ {1 / t}, \tag {7}
$$
as claimed.
In the context of our paper, we have  $Z = X_{T} \sim \mathcal{N}(0,\Sigma)$ ,  $f_{\theta}(Z) = \hat{X}_0 = \phi(X_T, T)$ ,  $f(Z) = X_0$ ,  $G$  is the standard Euclidean group SE(3), and we select  $t = 2$ . Then, the aforementioned inequality becomes
$$
\min  _ {g \in \mathrm {S E} (3)} \mathbf {E} _ {X _ {T}} [ \| \phi (X _ {T}, T) - g X _ {T} \| _ {2} ^ {2} ] \geq W _ {2} (p, q _ {T}) - \sqrt {\mathbf {E} _ {X _ {T}} [ \| \phi (X _ {T} , T) - \phi^ {*} (X _ {T} , T) \| _ {2} ^ {2} ]}, \tag {8}
$$
where we have changed the position of  $g$  by exploiting that the  $\ell_2$  norm is unitary invariant and that  $\phi$  is SE(3) equivariant. Alternatively, we can obtain the following tighter bound by stopping earlier at the derivation of equation 5:
$$
\mathbf {E} _ {X _ {T}} \left[ \| \phi \left(X _ {T}, T\right) - X _ {T} \| _ {2} ^ {2} \right] \geq W _ {2} (p, q _ {T}) - W _ {2} \left(p _ {\theta}, p\right), \tag {9}
$$
where  $q_{T}$  is the law of  $X_{T}$ .
F Learning the Conditional Dependence Relations Between Atoms