SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mutual Information (MI) Estimation and Maximization. MI can measure both the linear and nonlinear dependency between random variables. Some previous works [7, 35] try to use neural networks to estimate the lower bound of MI, including Donsker-Varadhan representation [19], Jensen-Shannon divergence [21] and Noise-Contrastive Estimation (NCE) [30, 31]. The optimization with InfoNCE loss [57] maximizes a lower bound of MI and is broadly shown to be a superior representation learning strategy [13, 32, 79, 51, 87]. In this work, we adopt the MI lower bound proposed by Liu et al. [52] with two conditional log-likelihoods, and we formulate the learning objective by mutually denoising the multimodal diffusion processes of two correlated proteins.
A.1 Broader Impacts and Limitations
The main objective of this research project is to enhance protein representations by utilizing joint pre-training using a vast collection of unlabeled protein structures. Unlike traditional unimodel pre-training methods, our approach takes advantage of both sequential and structural information, resulting in superior representations. This advantage allows for more comprehensive analysis of protein research and holds potential benefits for various real-world applications, including protein function prediction and sequence design.
Limitations. In this paper, we limit our pre-training dataset to less than 1M protein structures. However, considering the vast coverage of the AlphaFold Protein Structure Database, which includes over 200 million proteins, it becomes feasible to train more advanced and extensive protein encoders on larger datasets in the future. Furthermore, an avenue for future exploration is the incorporation of conformer-related information during pre-training and the development of improved noise schedules for multi-scale denoising pre-training. It is important to acknowledge that powerful pretrained models can potentially be misused for harmful purposes, such as the design of dangerous drugs. We anticipate that future studies will address and mitigate these concerns.
B Details of SiamDiff
In this section, we discuss details of our SiamDiff method. We first describe the parameterization of the generation process  $p_{\theta}(\mathcal{P}_1^{t - 1}|\mathcal{P}_1^t,\mathcal{P}_2^t)$  in Sec. B.1, derive the pre-training objective in Sec. B.2, and discuss some modifications when applied on residue-level models in Sec. 3.4.
B.1 Parameterization of Generation Process
Remember that we use  $\mathcal{P}_1^{0:T}$  and  $\mathcal{P}_2^{0:T}$  to denote the representation of the siamese trajectories  $\mathcal{P}_1^{0:T}$  and  $\mathcal{P}_2^{0:T}$ , respectively. Different from the generation process in traditional diffusion models, the parameterization of  $p_{\theta}(\mathcal{P}_1^{t-1}|\mathcal{P}_1^t,\mathcal{P}_2^t)$  should inject information from  $\mathcal{P}_2^t$ . Therefore, we use the extracted residue and atom representations (denoted as  $\boldsymbol{a}_2^t$  and  $\boldsymbol{h}_2^t$ ) of  $\mathcal{P}_2^t$  for this denoising step. Given the conditional independence in (2), this generation process can be decomposed into that on protein structures and sequences similarly in Sec. 2.3.
Algorithm 1 SiamDiff Pre-Training
Input: training dataset  $\mathcal{D}$  , learning rate  $\alpha$  Output: trained encoder  $\phi_{\theta}$
1: for  $(\mathcal{S},\mathcal{R})$  in  $\mathcal{D}$  do  
2:  $(\mathcal{S}_1,\mathcal{R}_1)\gets (\mathcal{S},\mathcal{R})$   
3:  $(\mathcal{S}_2,\mathcal{R}_2)\gets (\mathcal{S},\mathrm{perturb}(\mathcal{R},\epsilon))$   
4: sample noise scale  $t\in \{1\dots T\}$   
5:  $\mathcal{R}_1^t\sim \mathcal{N}(\sqrt{\bar{\alpha}_t}\mathcal{R}_1,(1 - \bar{\alpha}_t)I)$   
6:  $\mathcal{R}_2^t\sim \mathcal{N}(\sqrt{\bar{\alpha}_t}\mathcal{R}_2,(1 - \bar{\alpha}_t)I)$   
7:  $\mathcal{S}_1^t,\mathcal{S}_2^t\sim \mathrm{random\_mask}(\mathcal{S},t)$   
8:  $\pmb {h}_1\gets \phi_\theta (\pmb {S}_1^t,\pmb {R}_1^t)$   
9:  $\pmb {h}_2\gets \phi_\theta (\pmb {S}_2^t,\pmb {R}_2^t)$   
10:  $\mathcal{L}^{(\mathcal{R})} = \mathcal{L}(\mathcal{R}_1,\pmb {h}_2) + \mathcal{L}(\mathcal{R}_2,\pmb {h}_1)$   
11:  $\mathcal{L}^{(\mathcal{S})} = \mathcal{L}(\mathcal{S}_1,\pmb {h}_2) + \mathcal{L}(\mathcal{S}_2,\pmb {h}_1)$   
12:  $\theta \gets \theta -\alpha \nabla_{\theta}(\mathcal{L}^{(\mathcal{R})} + \mathcal{L}^{(\mathcal{S})})$   
13: end for
Generation process on protein structures. As in (4), modeling the generation process of protein structures is to model the noise on  $\mathcal{R}_1^t$  and gradually decorrupt the noisy structure. This can be parameterized with a noise prediction network  $\epsilon_{\theta}(\mathcal{P}_1^t,\mathcal{P}_2^t,t)$  that is translation-invariant and rotation-equivariant w.r.t.  $\mathcal{R}_1^t$ . Besides, the noise applied on  $\mathcal{R}_1^t$  should not change with transformations on  $\mathcal{R}_2^t$ , so  $\epsilon_{\theta}$  should be SE(3)-invariant w.r.t.  $\mathcal{R}_2^t$ .
To achieve these goals, we build our noise prediction network with atom representations  $\pmb{a}_2^t$  (which is SE(3)-invariant w.r.t.  $\mathcal{R}_2^t$ ) and atom coordinates  $\pmb{r}_1^t$  (which is SE(3)-equivariant w.r.t.  $\mathcal{R}_1^t$ ). We define an equivariant output similarly as in DiffPreT. Specifically, we have
$$
\left[ \epsilon_ {\theta} \left(\mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}, t\right) \right] _ {i} = \sum_ {j \in \mathcal {N} _ {1} ^ {t} (i)} m _ {i, j} \cdot \frac {\boldsymbol {r} _ {1 i} ^ {t} - \boldsymbol {r} _ {1 j} ^ {t}}{\left\| \boldsymbol {r} _ {1 i} ^ {t} - \boldsymbol {r} _ {1 j} ^ {t} \right\| _ {2}}, \text {w i t h} m _ {i, j} = \mathrm {M L P} \left(\boldsymbol {a} _ {2 i} ^ {t}, \boldsymbol {a} _ {2 j} ^ {t}, \mathrm {M L P} \left(\left\| \boldsymbol {r} _ {1 i} ^ {t} - \boldsymbol {r} _ {1 j} ^ {t} \right\| _ {2}\right)\right),
$$
where  $\mathcal{N}_1^t (i)$  denotes the neighbors of the atom  $i$  in the corresponding graph of  $\mathcal{P}_1^t$ . Note that  $\epsilon_{\theta}(\mathcal{P}_1^t,\mathcal{P}_2^t,t)$  achieves the equivariance requirement, as  $m_{i,j}$  is SE(3)-invariant w.r.t.  $\mathcal{R}_1^t$  and  $\mathcal{R}_2^t$  while  $\pmb {r}_{1i}^{t} - \pmb{r}_{1j}^{t}$  is translation-invariant and rotation-equivariant w.r.t.  $\mathcal{R}_1^t$ .
Generation process on protein sequences. As in (7), the generation process on sequences aims to predict masked residue types in  $S_1^0$  with a predictor  $\tilde{p}_{\theta}$ . In our setting of mutual prediction, we define the predictor based on representations of the same residues in  $S_2^t$ , which are also masked. Hence, for each masked residue  $i$  in  $S_2^t$ , we feed its representation  $h_{2i}^{t}$  to an MLP and predict the type of the corresponding residue type  $s_{1i}^0$  in  $S_1^0$ :
$$
\tilde {p} _ {\theta} \left(\mathcal {S} _ {1} ^ {0} | \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}\right) = \prod_ {i} \tilde {p} _ {\theta} \left(s _ {1 i} ^ {0} | \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}\right) = \prod_ {i} \operatorname {S o f t m a x} \left(s _ {1 i} ^ {0} | \mathrm {M L P} \left(h _ {2 i} ^ {t}\right)\right),
$$
B.2 Pre-Training Objective
Given the defined forward and reverse process on two trajectories, we now derive the pre-training objective based on the mutual diffusion loss in (13). We take the term  $\mathcal{L}^{(2\to 1)}$  for example and its counterpart can be derived in the same way. The objective can be decomposed into a structure loss  $\mathcal{L}^{(\mathcal{R},2\to 1)}$  and a sequence loss  $\mathcal{L}^{(\mathcal{S},2\to 1)}$ :
$$
\left. \mathcal {L} ^ {\left(\mathcal {R}, 2 \rightarrow 1\right)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {\mathrm {K L}} \left(q \left(\mathcal {R} _ {1} ^ {t - 1} \mid \mathcal {R} _ {1} ^ {t}, \mathcal {R} _ {1} ^ {0}\right) \mid | p _ {\theta} \left(\mathcal {R} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}\right)\right)\right], \right. \tag {15}
$$
$$
\mathcal {L} ^ {\left(\mathcal {S}, 2 \rightarrow 1\right)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {\mathrm {K L}} \left(q \left(\mathcal {S} _ {1} ^ {t - 1} \mid \mathcal {S} _ {1} ^ {t}, \mathcal {S} _ {1} ^ {0}\right) \| p _ {\theta} \left(\mathcal {S} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}\right)\right)\right]. \tag {16}
$$
Based on the derivation in Sec. 2.3, the structure loss  $\mathcal{L}^{(\mathcal{R},2\to 1)}$  can be simplified as
$$
\mathcal {L} _ {\text {s i m p l e}} ^ {(\mathcal {R}, 2 \rightarrow 1)} = \sum_ {t = 1} ^ {T} \mathbb {E} _ {\epsilon \sim \mathcal {N} (0, I)} \left[ \| \epsilon - \epsilon_ {\theta} \left(\mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}, t\right) \| _ {2} ^ {2} \right], \tag {17}
$$
and the sequence loss  $\mathcal{L}^{(S,2\to 1)}$  can be simplified as
$$
\mathcal {L} _ {\text {s i m p l e}} ^ {(S, 2 \rightarrow 1)} = \sum_ {t = 1} ^ {T} \sum_ {i} \operatorname {C E} \left(s _ {1 i} ^ {0}, \tilde {p} _ {\theta} \left(s _ {1 i} ^ {0} \mid \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {t}\right)\right). \tag {18}
$$
Then, the final objective in Sec. 3.3 can be easily derived.
In this section, we provide proofs for propositions in Sec. 2 and Sec. 3. Due to the similarity between the two methods, all propositions are restated for SiamDiff. DiffPreT can be seen as a special case that two siamese trajectories collapse into one.
For notations, we use the bold symbol to denote the representation of an object and use  $P_{1}^{0:T}$  and  $P_{2}^{0:T}$  to denote the corresponding random variables of representations of the siamese trajectories  $\mathcal{P}_{1}^{0:T}$  and  $\mathcal{P}_{2}^{0:T}$ .
Proposition 1 With some approximations, the mutual information between representations of two siamese trajectories is lower bounded by:
$$
I \left(\boldsymbol {P} _ {1} ^ {0: T}; \boldsymbol {P} _ {2} ^ {0: T}\right) \geq - \frac {1}{2} \left(\mathcal {L} ^ {(2 \rightarrow 1)} + \mathcal {L} ^ {(1 \rightarrow 2)}\right) + C,
$$
where  $C$  is a constant independent of our encoder and the term from trajectory  $\mathcal{P}_b^{0:T}$  to  $\mathcal{P}_a^{0:T}$  is defined as
$$
\mathcal {L} ^ {(b \rightarrow a)} := \mathbb {E} _ {\mathcal {P} _ {a} ^ {0: T}, \mathcal {P} _ {b} ^ {0: T}} \left[ \sum_ {t = 1} ^ {T} D _ {K L} \left(q (\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {a} ^ {0}) | | p (\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T})\right)\right],
$$
with  $b\to a$  being either  $2\to 1$  or  $1\rightarrow 2$
Proof. First, the mutual information between representations of two trajectories is defined as:
$$
I \left(\boldsymbol {P} _ {1} ^ {0: T}; \boldsymbol {P} _ {2} ^ {0: T}\right) = \mathbb {E} _ {\boldsymbol {\mathcal {P}} _ {1} ^ {0: T}, \boldsymbol {\mathcal {P}} _ {2} ^ {0: T} \sim p \left(\boldsymbol {P} _ {1} ^ {0: T}, \boldsymbol {P} _ {2} ^ {0: T}\right)} \left[ \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} , \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right)}{p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T}\right) p \left(\boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right)} \right], \tag {19}
$$
where the joint distribution is defined as  $p(\pmb{P}_1^{0:T}, \pmb{P}_2^{0:T}) = p(\pmb{P}_1^0, \pmb{P}_2^0) q(\pmb{P}_1^{1:T} | \pmb{P}_1^0) q(\pmb{P}_2^{1:T} | \pmb{P}_2^0)$ . Next, we can derive a lower bound with this definition:
$$
\begin{array}{l} I (\pmb {P} _ {1} ^ {0: T}; \pmb {P} _ {2} ^ {0: T}) = \mathbb {E} \left[ \log \frac {p (\pmb {\mathscr {P}} _ {1} ^ {0 : T} , \pmb {\mathscr {P}} _ {2} ^ {0 : T})}{p (\pmb {\mathscr {P}} _ {1} ^ {0}) q (\pmb {\mathscr {P}} _ {1} ^ {1 : T} | \pmb {\mathscr {P}} _ {1} ^ {0}) p (\pmb {\mathscr {P}} _ {2} ^ {0}) q (\pmb {\mathscr {P}} _ {2} ^ {1 : T} | \pmb {\mathscr {P}} _ {2} ^ {0})} \right] \\ \geq \mathbb {E} \left[ \log \frac {p (\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} , \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T})}{\sqrt {p (\boldsymbol {\mathcal {P}} _ {1} ^ {0}) p (\boldsymbol {\mathcal {P}} _ {2} ^ {0})} q (\boldsymbol {\mathcal {P}} _ {1} ^ {1 : T} | \boldsymbol {\mathcal {P}} _ {1} ^ {0}) q (\boldsymbol {\mathcal {P}} _ {2} ^ {1 : T} | \boldsymbol {\mathcal {P}} _ {2} ^ {0})} \right] \\ = \frac {1}{2} \mathbb {E} \left[ \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} , \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right) ^ {2}}{p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0}\right) p \left(\boldsymbol {\mathcal {P}} _ {2} ^ {0}\right) q \left(\boldsymbol {\mathcal {P}} _ {1} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {1} ^ {0}\right) ^ {2} q \left(\boldsymbol {\mathcal {P}} _ {2} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {2} ^ {0}\right) ^ {2}} \right] \\ = \frac {1}{2} \mathbb {E} \left[ \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} , \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right)}{p \left(\boldsymbol {\mathcal {P}} _ {2} ^ {0}\right) q \left(\boldsymbol {\mathcal {P}} _ {1} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {1} ^ {0}\right) q \left(\boldsymbol {\mathcal {P}} _ {2} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {2} ^ {0}\right)} + \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} , \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right)}{p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0}\right) q \left(\boldsymbol {\mathcal {P}} _ {1} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {1} ^ {0}\right) q \left(\boldsymbol {\mathcal {P}} _ {2} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {2} ^ {0}\right)} \right] \\ = \frac {1}{2} \mathbb {E} \left[ \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {1} ^ {0 : T} \mid \boldsymbol {\mathcal {P}} _ {2} ^ {0 : T}\right)}{q \left(\boldsymbol {\mathcal {P}} _ {1} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {1} ^ {0}\right)} + \log \frac {p \left(\boldsymbol {\mathcal {P}} _ {2} ^ {0 : T} \mid \boldsymbol {\mathcal {P}} _ {1} ^ {0 : T}\right)}{q \left(\boldsymbol {\mathcal {P}} _ {2} ^ {1 : T} \mid \boldsymbol {\mathcal {P}} _ {2} ^ {0}\right)} \right]. \\ \end{array}
$$
However, since the distribution of representations are intractable to sample for optimization, we instead sample the trajectories  $\mathcal{P}_1^{0:T}$  and  $\mathcal{P}_2^{0:T}$  from our defined diffusion process, i.e.,  $p(\mathcal{P}_1^{0:T},\mathcal{P}_2^{0:T}) = p(\mathcal{P}_1^0,\mathcal{P}_2^0)q(\mathcal{P}_1^{1:T}|\mathcal{P}_1^0)q(\mathcal{P}_2^{1:T}|\mathcal{P}_2^0)$ . Besides, instead of predicting representations, we use the representations from one trajectory to recover the other trajectory, which reflects more information than its representation. With these approximations, the lower bound above can be further written as:
$$
\frac {1}{2} \mathbb {E} \left[ \log \frac {p (\mathcal {P} _ {1} ^ {0 : T} | \mathcal {P} _ {2} ^ {0 : T})}{q (\mathcal {P} _ {1} ^ {1 : T} | \mathcal {P} _ {1} ^ {0})} + \log \frac {p (\mathcal {P} _ {2} ^ {0 : T} | \mathcal {P} _ {1} ^ {0 : T})}{q (\mathcal {P} _ {2} ^ {1 : T} | \mathcal {P} _ {2} ^ {0})} \right] \approx \frac {1}{2} \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ \log \frac {p (\mathcal {P} _ {1} ^ {0 : T} | \mathcal {P} _ {2} ^ {0 : T})}{q (\mathcal {P} _ {1} ^ {1 : T} | \mathcal {P} _ {1} ^ {0})} + \log \frac {p (\mathcal {P} _ {2} ^ {0 : T} \mid \mathcal {P} _ {1} ^ {0 : T})}{q (\mathcal {P} _ {2} ^ {1 : T} | \mathcal {P} _ {2} ^ {0})} \right]
$$
We now show the first term on the right hand side can be written as the loss defined in Proposition 1. The derivation is very similar with the proof of Proposition 3 in Xu et al. [82]. We include it here for
completeness:
$$
\begin{array}{l} \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ \log \frac {p \left(\mathcal {P} _ {1} ^ {0 : T} \mid \mathcal {P} _ {2} ^ {0 : T}\right)}{q \left(\mathcal {P} _ {1} ^ {1 : T} \mid \mathcal {P} _ {1} ^ {0}\right)} \right] \\ = \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ \sum_ {t = 1} ^ {T} \log \frac {p \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t} , \mathcal {P} _ {2} ^ {0 : T}\right)}{q \left(\mathcal {P} _ {1} ^ {t} \mid \mathcal {P} _ {1} ^ {t - 1}\right)} \right] \\ = \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ \right. \log \frac {\left(\mathcal {P} _ {1} ^ {0} \mid \mathcal {P} _ {1} ^ {1} , \mathcal {P} _ {2} ^ {0 : T}\right)}{q \left(\mathcal {P} _ {1} ^ {1} \mid \mathcal {P} _ {1} ^ {0}\right)} + \sum_ {t = 2} ^ {T} \log \left(\frac {p \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t} , \mathcal {P} _ {2} ^ {0 : T}\right)}{q \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t} , \mathcal {P} _ {1} ^ {0}\right)} \cdot \frac {q \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {0}\right)}{q \left(\mathcal {P} _ {1} ^ {t} \mid \mathcal {P} _ {1} ^ {0}\right)}\right) \Bigg ] \\ = \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ - \log q \left(\mathcal {P} _ {1} ^ {T} \mid \mathcal {P} _ {1} ^ {0}\right) + \log p \left(\mathcal {P} _ {1} ^ {0} \mid \mathcal {P} _ {1} ^ {1}, \mathcal {P} _ {2} ^ {0: T}\right) + \sum_ {t = 2} ^ {T} \log \frac {p \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t} , \mathcal {P} _ {2} ^ {0 : T}\right)}{q \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t} , \mathcal {P} _ {1} ^ {0}\right)} \right] \\ = - \mathbb {E} _ {\mathcal {P} _ {1} ^ {0: T}, \mathcal {P} _ {2} ^ {0: T}} \left[ \sum_ {t = 1} ^ {T} D _ {\mathrm {K L}} \left(q \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {1} ^ {0}\right) \mid \mid p \left(\mathcal {P} _ {1} ^ {t - 1} \mid \mathcal {P} _ {1} ^ {t}, \mathcal {P} _ {2} ^ {0: T}\right)\right)\right] + C ^ {(2 \rightarrow 1)} \\ = - \mathcal {L} ^ {(2 \rightarrow 1)} + C ^ {(2 \rightarrow 1)}, \\ \end{array}
$$
where we merge the term  $p(\mathcal{P}_1^0|\mathcal{P}_1^1,\mathcal{P}_2^{0:T})$  into the sum of KL divergences for brevity and use  $C^{(2\to 1)}$  to denote the constant independent of our encoder. Note that the counterpart can be derived in the same way. Adding these two terms together finishes the proof of Proposition 1.
C.2 Proof of Pre-Training Loss Decomposition
We restate the proposition of pre-training loss decomposition rigorously as below.
Proposition 2 Given the assumptions 1) the separation of the diffusion process on protein structures and sequences
$$
q \left(\mathcal {P} _ {a} ^ {t} \mid \mathcal {P} _ {a} ^ {t - 1}\right) = q \left(\mathcal {R} _ {a} ^ {t} \mid \mathcal {R} _ {a} ^ {t - 1}\right) \cdot q \left(\mathcal {S} _ {a} ^ {t} \mid \mathcal {S} _ {a} ^ {t - 1}\right), \tag {20}
$$
and 2) the conditional independence of the generation process
$$
p _ {\theta} \left(\mathcal {P} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \boldsymbol {\mathcal {P}} _ {b} ^ {t}\right) = p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \boldsymbol {\mathcal {P}} _ {b} ^ {t}\right) \cdot p _ {\theta} \left(\mathcal {S} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \boldsymbol {\mathcal {P}} _ {b} ^ {t}\right), \tag {21}
$$
it can be proved that
$$
\mathcal {L} ^ {(b \rightarrow a)} = \mathcal {L} ^ {(\mathcal {R}, b \rightarrow a)} + \mathcal {L} ^ {(\mathcal {S}, b \rightarrow a)}, \tag {22}
$$
where the three loss terms are defined as
$$
\mathcal {L} ^ {(b \rightarrow a)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {K L} \left(q (\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {a} ^ {0}) | | p _ {\theta} (\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t})\right)\right],
$$
$$
\mathcal {L} ^ {(\mathcal {R}, b \rightarrow a)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {K L} \left(q \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {R} _ {a} ^ {t}, \mathcal {R} _ {a} ^ {0}\right) | | p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}\right)\right)\right],
$$
$$
\mathcal {L} ^ {(\mathcal {S}, b \rightarrow a)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {K L} \left(q (\mathcal {S} _ {a} ^ {t - 1} | \mathcal {S} _ {a} ^ {t}, \mathcal {S} _ {a} ^ {0}) | | p _ {\theta} (\mathcal {S} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t})\right)\right],
$$
with  $b\to a$  referring to the term from trajectory  $\mathcal{P}_b^{0:T}$  to  $\mathcal{P}_a^{0:T}$ .
Proof. Let  $\mathcal{L}_t^{(\cdot)}$  to denote the t-th KL divergence term in  $\mathcal{L}^{(\cdot)}$ . Then, we have
$$
\begin{array}{l} \mathcal {L} _ {t} ^ {(b \rightarrow a)} = D _ {\mathrm {K L}} \left(q \left(\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {a} ^ {0}\right) | | p _ {\theta} \left(\mathcal {P} _ {a} ^ {t - 1} | \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right)\right) \\ = D _ {\mathrm {K L}} \left(\left[ q \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {R} _ {a} ^ {t}, \mathcal {R} _ {a} ^ {0}\right) q \left(\mathcal {S} _ {a} ^ {t - 1} \mid \mathcal {S} _ {a} ^ {t}, \mathcal {S} _ {a} ^ {0}\right) \right] | | \left[ p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right) p _ {\theta} \left(\mathcal {S} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right) \right]\right) \\ = D _ {\mathrm {K L}} \left(q \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {R} _ {a} ^ {t}, \mathcal {R} _ {a} ^ {0}\right) \| p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right)\right) + D _ {\mathrm {K L}} \left(q \left(\mathcal {S} _ {a} ^ {t - 1} \mid \mathcal {S} _ {a} ^ {t}, \mathcal {S} _ {a} ^ {0}\right) \| p _ {\theta} \left(\mathcal {S} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right)\right) \\ = \mathcal {L} _ {t} ^ {(\mathcal {R}, b \rightarrow a)} + \mathcal {L} _ {t} ^ {(\mathcal {S}, b \rightarrow a)}, \\ \end{array}
$$
where we use the assumptions (20) and (21) in the second equality. The third equality is due to the additive property of the KL divergence for independent distributions. Adding  $T$  KL divergence terms together will lead to (22).  $\square$
C.3 Proof of Simplified Structure Loss
For completeness, we show how to derive the simplified structure loss. The proof is directly adapted from [82].
Proposition 3 Given the definition of the forward process
$$
q \left(\mathcal {R} _ {a} ^ {t} \mid \mathcal {R} _ {a} ^ {t - 1}\right) = \mathcal {N} \left(\mathcal {R} _ {a} ^ {t}; \sqrt {1 - \beta_ {t}} \mathcal {R} _ {a} ^ {t - 1}, \beta_ {t} I\right), \tag {23}
$$
and the reverse process
$$
p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}\right) = \mathcal {N} \left(\mathcal {R} _ {a} ^ {t - 1}; \mu_ {\theta} \left(\mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}, t\right), \sigma_ {t} ^ {2} I\right), \tag {24}
$$
$$
\mu_ {\theta} \left(\mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}, t\right) = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathcal {R} _ {a} ^ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \epsilon_ {\theta} \left(\mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}, t\right)\right), \tag {25}
$$
the structure loss function
$$
\mathcal {L} ^ {\left(\mathcal {R}, b \rightarrow a\right)} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} D _ {K L} \left(q \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {R} _ {a} ^ {t}, \mathcal {R} _ {a} ^ {0}\right) \mid \mid p _ {\theta} \left(\mathcal {R} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}\right)\right)\right], \tag {26}
$$
can be simplified to
$$
\mathcal {L} ^ {(\mathcal {R}, b \rightarrow a)} = \sum_ {t = 1} ^ {T} \gamma_ {t} \mathbb {E} _ {\epsilon \sim \mathcal {N} (0, I)} \left[ \| \epsilon - \epsilon_ {\theta} (\mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {t}, t) \| _ {2} ^ {2} \right], \tag {27}
$$
where  $\gamma_{t} = \frac{\beta_{t}}{2\alpha_{t}(1 - \bar{\alpha}_{t - 1})}$  with  $\alpha_{t} = 1 - \beta_{t}$ ,  $\bar{\alpha}_{t} = \prod_{s = 1}^{t}\alpha_{s}$  and  $b\to a$  is either  $2\rightarrow 1$  or  $1\rightarrow 2$ .