SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our method begins by generating different conformers of a given protein. However, direct sampling requires an accurate characterization of the energy landscape of protein conformations, which can be difficult and time-consuming. To address this issue, we adopt a commonly used scheme for sampling randomly simulated conformers by adding torsional perturbations to the side chains [36].
Specifically, given the original protein  $\mathcal{P} = (\mathcal{S},\mathcal{R})$  we consider it as the native state  $\mathcal{P}_1 = (\mathcal{S}_1,\mathcal{R}_1)$  and generate a correlated conformer  $\mathcal{P}_2 = (\mathcal{S}_2,\mathcal{R}_2)$  by
randomly perturbing the protein structure. That is, we set  $S_{2}$  to be the same as  $S_{1}$ , and  $\mathcal{R}_2$  is obtained by applying a perturbation function  $\text{perturb}(\mathcal{R}_1, \epsilon)$  to the original residue structure, where  $\epsilon \in [0, 2\pi)^{n_r \times 4}$  is a noise vector drawn from a wrapped normal distribution [15]. The  $\text{perturb}(\cdot, \cdot)$  function rotates the side-chain of each residue according to the sampled torsional noises. To avoid atom clashes, we adjust the variance of the added noise and regenerate any conformers that violate physical constraints. It should be noted that the scheme can be adapted flexibly when considering different granularities of structures. For example, instead of considering side chain rotation on a fixed backbone, we can also consider a flexible backbone by rotating backbone angles, thereby generating approximate conformers.
Although the current torsional perturbation scheme is effective, there is potential for improvement. Future research can explore enhancements such as incorporating available rotamer libraries [62] and introducing backbone flexibility using empirical force fields [72].
Figure 3: High-level illustration of SiamDiff. Mutual denoising of diffusion trajectories is performed across two correlated conformers  $\mathcal{P}_1$  and  $\mathcal{P}_2$ .
3.2 Siamese Diffusion Trajectory Generation
To maintain the information-rich joint diffusion trajectories from DiffPreT, we generate trajectories for pairs of conformers, a.k.a., siamese trajectories. We first sample the diffusion trajectories  $\mathcal{P}_1^{0:T}$  and  $\mathcal{P}_2^{0:T}$  for conformers  $\mathcal{P}_1$  and  $\mathcal{P}_2$ , respectively, using the joint diffusion process outlined in Sec. 2. For example, starting from  $\mathcal{P}_1^0 = \mathcal{P}_1 = (\mathcal{S}_1,\mathcal{R}_1)$ , we use the joint diffusion on structures and sequences to define the diffusion process  $q(\mathcal{P}_1^{1:T}|\mathcal{P}_1^0) = q(\mathcal{R}_1^{1:T}|\mathcal{R}_1^0)q(\mathcal{S}_1^{1:T}|\mathcal{S}_1^0)$ . We derive trajectories on structures  $\mathcal{R}_1^{1:T}$  using the Gaussian noise in (3) and derive the sequence diffusion process  $\mathcal{S}^{1:T}$  using the random masking in (6). In this way, we define the trajectory  $\mathcal{P}_1^{0:T} = \{(S_1^t,\mathcal{R}_1^t)\}_{t=0}^T$  for  $\mathcal{P}_1$  and can derive the siamese trajectory  $\mathcal{P}_2^{0:T} = \{(S_2^t,\mathcal{R}_2^t)\}_{t=0}^T$  similarly.
3.3 Mutual Information Maximization between Representations of Siamese Trajectories
We aim to maximize the mutual information (MI) between representations of siamese trajectories constructed in a way that reflects the correlation between different conformers of the same protein. Direct optimization of MI is intractable, so we instead maximize a lower bound. In App. C.1, we show that this problem can be transformed into the minimization of a loss function for mutual denoising between two trajectories:  $\mathcal{L} = \frac{1}{2} (\mathcal{L}^{(2\to 1)} + \mathcal{L}^{(1\to 2)})$ , where
$$
\mathcal {L} ^ {(b \rightarrow a)} := \mathbb {E} _ {\mathcal {P} _ {a} ^ {0: T}, \mathcal {P} _ {b} ^ {0: T}} \left[ \sum_ {t = 1} ^ {T} D _ {\mathrm {K L}} \left(q \left(\mathcal {P} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {a} ^ {0}\right) \mid p \left(\mathcal {P} _ {a} ^ {t - 1} \mid \mathcal {P} _ {a} ^ {t}, \mathcal {P} _ {b} ^ {0: T}\right)\right)\right], \tag {13}
$$
with  $b\to a$  being either  $2\rightarrow 1$  or  $1\rightarrow 2$  and  $\mathcal{P}_b^{0:T}$  being representations of the trajectory  $\mathcal{P}_b^{0:T}$ .
The two terms share the similar formula as the ELBO loss in (8). Take  $\mathcal{L}^{(2\rightarrow 1)}$  for example. Here  $q(\mathcal{P}_1^{t - 1}|\mathcal{P}_1^t,\mathcal{P}_1^0)$  is a posteriorially tractable with our definition of each diffusion step  $q(\mathcal{P}_1^t |\mathcal{P}_1^{t - 1})$  in (3) and (6). The reverse process is learnt to generate a less noisy state  $\mathcal{P}_1^{t - 1}$  given the current state  $\mathcal{P}_1^t$  and representations of the siamese trajectory  $\mathcal{P}_2^{0:T}$ , which are extracted by the protein encoder to be pre-trained. The parameterization of the reverse process is similar as in Sec. 2.2, with the representations replaced by those of  $\mathcal{P}_2^{0:T}$  (see App. B for details).
Our approach involves mutual prediction between two siamese trajectories, which is similar to the idea of mutual representation reconstruction in [28, 14]. However, since  $\mathcal{P}_1$  and  $\mathcal{P}_2$  share information about the same protein, the whole trajectory of  $\mathcal{P}_2$  could provide too many clues for denoising towards  $\mathcal{P}_1^{t - 1}$ , making the pre-training task trivial. To address this issue, we parameterize  $p(\mathcal{P}_1^{t - 1}|\mathcal{P}_1^t,\mathcal{P}_2^{0:T})$  with  $p_{\theta}(\mathcal{P}_1^{t - 1}|\mathcal{P}_1^t,\mathcal{P}_2^t)$ . For diffusion on sequences, we further guarantee that the same set of residues are masked in  $S_1^t$  and  $S_2^t$  to avoid leakage of ground-truth residue types across correlated trajectories.
Final pre-training objective. Given the similarity between the one-side objective and the ELBO loss in (8), we can use a similar way to decompose the objective into structure and sequence losses and then derive simplified loss functions for each side. To summarize, the ultimate training objective for our method is
$$
\mathcal {L} _ {\text {s i m p l e}} = \frac {1}{2} \left(\mathcal {L} _ {\text {s i m p l e}} ^ {(\mathcal {R}, 2 \rightarrow 1)} + \mathcal {L} _ {\text {s i m p l e}} ^ {(\mathcal {S}, 2 \rightarrow 1)} + \mathcal {L} _ {\text {s i m p l e}} ^ {(\mathcal {R}, 1 \rightarrow 2)} + \mathcal {L} _ {\text {s i m p l e}} ^ {(\mathcal {S}, 1 \rightarrow 2)}\right), \tag {14}
$$
where  $\mathcal{L}_{\mathrm{simple}}^{(\cdot ,b\to a)}$  is the loss term defined by predicting  $\mathcal{P}_a^{0:T}$  from  $\mathcal{P}_b^{0:T}$  (see App. B for derivation).
3.4 Residue-level model
Residue-level protein graphs are simplified atom graphs that enable efficient message passing between nodes and edges. As in Zhang et al. [87], we only keep the alpha carbon atom of each residue and add sequential, radius and K-nearest neighbor edges as different types of edges. For SiamDiff, the residue-level model cannot discriminate conformers generated by rotating side chains, since we only keep CA atoms. To solve this problem, we directly add Gaussian noises to the coordinates instead to generate approximate conformers. Specifically, the correlated conformer  $\mathcal{P}_2 = (\mathcal{S}_2,\mathcal{R}_2)$  is defined by  $S_{2} = S_{1},\mathcal{R}_{2} = \mathcal{R}_{1} + \epsilon$  , where  $\epsilon \in \mathbb{R}^{n_a\times 3}$  is the noise drawn from a Gaussian distribution.
Pre-training Methods on Proteins. Self-supervised pre-training methods have been widely used to acquire co-evolutionary information from large-scale protein sequence corpus, inducing performant protein language models (PLMs) [20, 53, 61, 50]. Typical sequence pre-training methods include masked protein modeling [20, 61, 50] and contrastive learning [53]. The pre-trained PLMs have achieved impressive performance on a variety of downstream tasks for structure and function prediction [59, 81]. Recent works have also studied pre-training on unlabeled protein structures for generalizable representations, covering contrastive learning [87, 33], self-prediction of geometric quantities [87, 10] and denoising score matching [29, 76]. Compared with existing works, our methods model the joint distribution of sequences and structures via diffusion models, which captures both co-evolutionary information and detailed structural characteristics.
Diffusion Probabilistic Models (DPMs). DPM was first proposed in Sohl-Dickstein et al. [63] and has been recently rekindled for its strong performance on image and waveform generation [37, 11]. While DPMs are commonly used for modeling continuous data, there has also been research exploring discrete DPMs that have achieved remarkable results on generating texts [4, 49], graphs [71] and images [38]. Inspired by these progresses, DPMs have been adopted to solve problems in chemistry and biology domain, including molecule generation [82, 39, 78, 43], molecular representation learning [52], protein structure prediction [77], protein-ligand binding [15], protein design [2, 54, 40, 74] and motif-scaffolding [69]. In alignment with recent research efforts focused on diffusion-based image representation learning [1], this work presents a novel investigation into how DPMs can contribute to protein representation learning.
Now we discuss the relationship between our method and previous works.
Advantages of joint denoising. Compared with previous diffusion models focusing on either protein sequences [83] or structures [29] or cross-modal contrastive learning [12, 84], in this work, we
Table 1: Atom-level results on Atom3D tasks.
Method	PIP	MSP	RES	PSR	Mean Rank
AUROC	AUROC	Accuracy	Global ρ	Mean ρ
w/ pre-training	GearNet-Edge	0.868±0.002	0.633±0.067	0.441±0.001	0.782±0.021	0.488 ± 0.012	7.6
Denoising Score Matching	0.877±0.002	0.629±0.040	0.448±0.001	0.813±0.003	0.518±0.020	5.2
Residue Type Prediction	0.879±0.004	0.620±0.027	0.449±0.001	0.826±0.020	0.518±0.018	4.4
Distance Prediction	0.872±0.001	0.677±0.020	0.422±0.001	0.840±0.020	0.522±0.004	4.0
Angle Prediction	0.878±0.001	0.642±0.013	0.419±0.001	0.813±0.007	0.503±0.012	6.2
Dihedral Prediction	0.878±0.004	0.591±0.008	0.414±0.001	0.821±0.002	0.497±0.004	6.8
Multiview Contrast	0.871±0.003	0.646±0.006	0.368±0.001	0.805±0.005	0.502±0.009	7.2
DiffPreT	0.880±0.005	0.680±0.018	0.452±0.001	0.821±0.007	0.533±0.006	2.4
SiamDiff	0.884±0.003	0.698±0.020	0.460±0.001	0.829±0.012	0.546±0.018	1.2
Table 2: Residue-level results on EC and Atom3D tasks.
Method	EC	MSP	PSR	Mean Rank
AUPR	Fmax	AUROC	Global ρ	Mean ρ
w/ pre-training	GearNet-Edge	0.837±0.002	0.811±0.001	0.644±0.023	0.763±0.012	0.373±0.021	7.8
Denoising Score Matching	0.859±0.003	0.840±0.001	0.645±0.028	0.795±0.027	0.429±0.017	5.0
Residue Type Prediction	0.851±0.002	0.826±0.005	0.636±0.003	0.828±0.005	0.480±0.031	5.4
Distance Prediction	0.858±0.003	0.836±0.001	0.623±0.007	0.796±0.017	0.416±0.021	6.4
Angle Prediction	0.873±0.003	0.849±0.001	0.631±0.041	0.802±0.015	0.446±0.009	4.2
Dihedral Prediction	0.858±0.001	0.840±0.001	0.568±0.022	0.732±0.021	0.398±0.022	7.2
Multiview Contrast	0.875±0.003	0.857±0.003	0.713±0.036	0.752±0.012	0.388±0.015	4.0
DiffPreT	0.864±0.002	0.844±0.001	0.673±0.042	0.815±0.008	0.505±0.007	3.2
SiamDiff	0.878±0.003	0.857±0.003	0.700±0.043	0.856±0.007	0.521±0.016	1.2
perform joint diffusion on both modalities. Note that given a sequence  $\mathcal{S}$  and a structure  $\mathcal{R}$  that exist in the nature with high probability, the sequence-structure tuple  $\mathcal{P} = (\mathcal{S},\mathcal{R})$  may not be a valid state of this protein. Consequently, instead of modeling the marginal or conditional distribution, we model the joint distribution of protein sequences and structures.
Connection with diffusion models. Diffusion models excel in image and text generation [17, 49] and have been applied to unsupervised representation learning [1]. Previous works explored denoising objectives [23, 9] but lacked explicit supervision for different conformers, while our method incorporates mutual prediction between siamese diffusion trajectories to capture conformer correlation and regularize protein structure manifold.
Difference with denoising distance matching. While previous works rely on perturbing distance matrices [52, 29], which can violate the triangular inequality and produce negative values, our approach directly adds noise at atom coordinates, as demonstrated in Xu et al. [82]. This distinction allows us to address the limitations associated with denoising distance matching algorithms used in molecule and protein generation and pre-training.
Comparison with other deep generative models. Self-supervised learning essentially learns an Energy-Based Model (EBM) for modeling data distribution [48], making VAE [45], GAN [26], and normalizing flow [60] applicable for pre-training. However, these models limit flexibility or fail to acquire high sampling quality and diversity compared to diffusion models [54]. Therefore, we focus on using diffusion models for pre-training and leave other generative models for future work.
Pre-training datasets. Following Zhang et al. [87], we pre-train our models with the AlphaFold protein structure database v1 [44, 70], including 365K proteome-wide predicted structures.
Downstream benchmark tasks. In our evaluation, we assess EC prediction task [24] for catalysis behavior of proteins and four ATOM3D tasks [68]. The EC task involves 538 binary classification problems for Enzyme Commission (EC) numbers. We use dataset splits from Gligorijevic et al. [24] with a  $95\%$  sequence identity cutoff. The ATOM3D tasks include Protein Interface Prediction (PIP),
Mutation Stability Prediction (MSP), Residue Identity (RES), and Protein Structure Ranking (PSR) with different dataset splits based on sequence identity or competition year. Details are in App. D.
Baseline methods. In our evaluation, we utilize GearNet-Edge as the underlying model for both atom- and residue-level structures. GearNet-Edge incorporates various types of edges and edge-type-specific convolutions, along with message passing between edges, to model protein structures effectively. We compare our proposed methods with several previous protein structural pre-training algorithms, including multiview contrastive learning [87], denoising score matching [29], and four self-prediction methods (residue type, distance, angle, and dihedral prediction) [87]. For residue-level tasks, we include EC, MSP, and PSR in our evaluation, while PIP and RES tasks are specifically designed for atom-level models. Besides, we exclude EC from the atom-level evaluation due to the limited presence of side-chain atoms in the downloaded PDB dataset.
Training and evaluation. We pre-train our model for 50 epochs on the AlphaFold protein structure database following Zhang et al. [87] and fine-tune it for 50 epochs on EC, MSP, and PSR. However, due to time constraints, we only fine-tune the models for 10 epochs on the RES and PIP datasets. Results are reported as mean and standard deviation across three seeds (0, 1, and 2). Evaluation metrics include  $\mathrm{F}_{\mathrm{max}}$  and AUPR for EC, AUROC for PIP and MSP, Spearman's  $\rho$  for PSR, and micro-averaged accuracy for RES. More details about experimental setup can be found in App. D.
Tables 1 and 2 provide a comprehensive overview of the results obtained by GearNet-Edge on both atom- and residue-level benchmark tasks. The tables clearly demonstrate that both DiffPreT and SiamDiff exhibit significant improvements over GearNet-Edge without pre-training on both levels, underscoring the effectiveness of our pre-training methods.
An interesting observation from the tables is that previous pre-training methods tend to excel in specific tasks while showing limitations in others. For instance, Multiview Contrast, designed for capturing similar functional motifs [87], struggles with structural intricacies and local atomic interactions, resulting in lower performance on tasks like Protein Interface Prediction (PIP), Protein Structure Ranking (PSR), and Residue Identity (RES). Self-prediction methods excel at capturing structural details or residue type dependencies but show limitations in function prediction tasks, such as Enzyme Commission (EC) number prediction and Mutation Stability Prediction (MSP), and do not consistently improve performance on both atom and residue levels.
In contrast, our DiffPreT approach achieves top-3 performance in nearly all considered tasks, showcasing its versatility and effectiveness across different evaluation criteria. Moreover, SiamDiff surpasses all other pre-training methods, achieving the best results in 6 out of 7 tasks, establishing it as the state-of-the-art pre-training approach. These results provide compelling evidence that our joint diffusion pre-training strategy successfully captures the intricate interactions between different proteins (PIP), captures local structural details (RES) and global structural characteristics (PSR), and extracts informative features crucial for accurate function prediction (EC) across various tasks.
To analyze the effect of different components of SiamDiff, we perform ablation study on atom-level tasks and present results in Table 3. We first examine two degenerate settings of joint diffusion, i.e., "w/o sequence diffusion" and "w/o structure diffusion". These settings lead to a deterioration in performance across all benchmark
Table 3: Ablation study on atom-level Atom3D tasks.
Method	PIP	MSP	RES	PSR
AUROC	AUROC	Accuracy	Global ρ
GearNet-Edge	0.868±0.002	0.633±0.067	0.441±0.001	0.782±0.021
SiamDiff	0.884±0.003	0.698±0.020	0.460±0.001	0.829±0.008
w/o seq. diff.	0.873±0.004	0.695±0.002	0.443±0.001	0.803±0.010
w/o struct. diff.	0.878±0.003	0.652±0.021	0.456±0.001	0.805±0.005
w/o MI max.	0.880±0.005	0.680±0.018	0.452±0.001	0.821±0.007
w/small noise	0.875±0.002	0.646±0.031	0.444±0.001	0.828±0.005
w/large noise	0.867±0.003	0.683±0.020	0.443±0.001	0.819±0.011
tasks, highlighting the importance of both sequence diffusion for residue type identification in RES and structure diffusion for capturing the structural stability of mutation effects in MSP. Next, we compare SiamDiff with DiffPreT, which lacks mutual information maximization between correlated conformers. The consistent improvements observed across all tasks indicate the robustness and effectiveness of our proposed mutual information maximization scheme.
Besides, we compare our method to baselines with fixed small  $(T = 1)$  and large  $(T = 100)$  noise levels to demonstrate the benefits of multi-scale denoising in diffusion pre-training. Interestingly, we observe that denoising with large noise enhances performance on MSP by capturing significant structural changes that lead to structural instability, while denoising with small noise improves performance in PSR by capturing fine-grained details for protein structure assessment. By incorporating multi-scale noise, we eliminate the need for manual tuning of the noise level as a hyperparameter and leverage the advantages of both large- and small-scale noise, as evidenced in the table.
5.4 Combine with Protein Language Models
Protein language models (PLMs) have recently become a standard method for extracting representations from protein sequences, such as ESM [50]. However, these methods are unable to directly handle structure-related tasks
Table 4: ESM2-650M-GearNet on residue-level tasks.
Method	EC	MSP	PSR
AUPR	Fmax	AUROC	Global ρ	Mean ρ
GearNet-Edge	0.837±0.002	0.811±0.001	0.664±0.023	0.764±0.012	0.373±0.021
w/ SiamDiff	0.878±0.003	0.857±0.003	0.700±0.043	0.856±0.007	0.521±0.016
ESM-GearNet	0.904±0.002	0.890±0.002	0.685±0.027	0.829±0.013	0.595±0.010
w/ SiamDiff	0.907±0.001	0.897±0.001	0.692±0.010	0.863±0.009	0.656±0.011
in Atom3D without using protein structures as input. A recent solution addresses this by feeding residue representations outputted by ESM into the protein structure encoder GearNet [86]. To showcase the potential of SiamDiff on PLM-based encoders, we pre-trained the ESM-GearNet encoder using SiamDiff and evaluated its performance on residue-level tasks. Considering the model capacity and computational budget, we selected ESM-2-650M as the base PLM. The results in Table 4 demonstrate the performance improvements obtained by introducing the PLM component in ESM-GearNet. Furthermore, after pre-training with SiamDiff, ESM-GearNet achieves even better performance on all tasks, especially on PSR where ESM-only representations are not indicative for structure ranking. This highlights the benefits of our method for PLM-based encoders.
In addition, we provide experimental results about pre-training datasets in App. E, multimodal baselines in App. F, different diffusion strategies in App. G, and different backbone models in App. H.
In this work, we propose the DiffPreT approach to pre-train a protein encoder by sequence-structure joint diffusion modeling, which captures the inter-atomic interactions within structure and the residue type dependencies along sequence. We further propose the SiamDiff method to enhance DiffPreT by additionally modeling the correlation between different conformers of one protein. Extensive experiments on diverse types of tasks and on both atom- and residue-level structures verify the competitive performance of DiffPreT and the superior performance of SiamDiff.
The authors would like to thank Zhaocheng Zhu, Chence Shi, Jiarui Lu, Huiyu Cai, Xinyu Yuan and Bozitao Zhong for their helpful discussions and comments.
This project is supported by AIHN IBM-MILA partnership program, the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino-Bird Gift Fund, a NRC Collaborative R&D Project (AI4D-CORE-06) as well as the IVADO Fundamental Research Project grant PRF-2019-3583139727.
[83] Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. *bioRxiv*, 2022.
[84] Yuning You and Yang Shen. Cross-modality and self-supervised protein embedding for compound-protein affinity and contact prediction. Bioinformatics, 38(Supplement_2):ii68-ii74, 2022.
[85] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=tYIMtogyee.
[86] Zuobai Zhang, Minghao Xu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, and Jian Tang. Enhancing protein language model with structure-based encoder and pre-training. In ICLR 2023-Machine Learning for Drug Discovery workshop.
[87] Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. arXiv preprint arXiv:2203.06125, 2022.
[88] Zhaocheng Zhu, Chence Shi, Zuobai Zhang, Shengchao Liu, Minghao Xu, Xinyu Yuan, Yangtian Zhang, Junkun Chen, Huiyu Cai, Jiarui Lu, et al. Torchdrug: A powerful and flexible machine learning platform for drug discovery. arXiv preprint arXiv:2202.08320, 2022.
Protein Structure Encoder. The community witnessed a surge of research interests in learning informative protein structure representations using structure-based encoders. The encoders are designed to capture protein structural information on different granularity, including residue-level structures [24, 87, 80], atom-level structures [34, 41, 73] and protein surfaces [22, 66, 64]. In this work, we focus on pre-training a typical residue-level structure encoder, i.e., GearNet-Edge [87], and a typical atom-level structure encoder, i.e., GVP [41].