Title: Pre-Training Protein Encoder via Siamese Sequence-Structure Diffusion Trajectory Prediction
Venue: NeurIPS 2023
Accepted: yes
Keywords: Protein representation learning

diffusion models

self-supervised learning
Abstract: Self-supervised pre-training methods on proteins have recently gained attention, with most approaches focusing on either protein sequences or structures, neglecting the exploration of their joint distribution, which is crucial for a comprehensive understanding of protein functions by integrating co-evolutionary information and structural characteristics. In this work, inspired by the success of denoising diffusion models in generative tasks, we propose the DiffPreT approach to pre-train a protein encoder by sequence-structure joint diffusion modeling. DiffPreT guides the encoder to recover the native protein sequences and structures from the perturbed ones along the joint diffusion trajectory, which acquires the joint distribution of sequences and structures. Considering the essential protein conformational variations, we enhance DiffPreT by a method called Siamese Diffusion Trajectory Prediction (SiamDiff) to capture the correlation between different conformers of a protein. SiamDiff attains this goal by maximizing the mutual information between representations of diffusion trajectories of structurally-correlated conformers. We study the effectiveness of DiffPreT and SiamDiff on both atom- and residue-level structure-based protein understanding tasks. Experimental results show that the performance of DiffPreT is consistently competitive on all tasks, and SiamDiff achieves new state-of-the-art performance, considering the mean ranks on all tasks. Code will be released upon acceptance.

Pre-Training Protein Encoder via Siamese Sequence-Structure Diffusion Trajectory Prediction

Abstract

Self-supervised pre-training methods on proteins have recently gained attention, with most approaches focusing on either protein sequences or structures, neglecting the exploration of their joint distribution, which is crucial for a comprehensive understanding of protein functions by integrating co-evolutionary information and structural characteristics. In this work, inspired by the success of denoising diffusion models in generative tasks, we propose the DiffPreT approach to pre-train a protein encoder by sequence-structure joint diffusion modeling. DiffPreT guides the encoder to recover the native protein sequences and structures from the perturbed ones along the joint diffusion trajectory, which acquires the joint distribution of sequences and structures. Considering the essential protein conformational variations, we enhance DiffPreT by a method called Siamese Diffusion Trajectory Prediction ( SiamDiff ) to capture the correlation between different conformers of a protein. SiamDiff attains this goal by maximizing the mutual information between representations of diffusion trajectories of structurally-correlated conformers. We study the effectiveness of DiffPreT and SiamDiff on both atom- and residue-level structure-based protein understanding tasks. Experimental results show that the performance of DiffPreT is consistently competitive on all tasks, and SiamDiff achieves new state-of-the-art performance, considering the mean ranks on all tasks. Code will be released upon acceptance.

1 Introduction

Machine learning-based methods have made remarkable strides in predicting protein structures [44, 5, 50] and functionality [55, 24]. Among them, self-supervised (unsupervised) pre-training approaches [20, 61, 87] have been successful in learning effective protein representations from available protein sequences or from their experimental/predicted structures. These pre-training approaches are based on the rationale that modeling the input distribution of proteins provides favorable initialization of model parameters and serves as effective regularization for downstream tasks [25]. Previous methods have primarily focused on modeling the marginal distribution of either protein sequences to acquire co-evolutionary information [20, 61], or protein structures to capture essential characteristics for tasks such as function prediction and fold classification [87, 33]. Nevertheless, both these forms of information hold significance in revealing the underlying functions of proteins and offer complementary perspectives that are still not extensively explored. To address this gap, a more promising approach for pre-training could involve modeling the joint distribution of protein sequences and structures, surpassing the limitations of unimodal pre-training methods.

To model this joint distribution, denoising diffusion models [37, 65] have recently emerged as one of the most effective methods due to their simple training objective and high sampling quality and diversity [2, 54, 40, 74]. However, the application of diffusion models has predominantly been explored in the context of generative tasks, rather than within pre-training and fine-tuning frameworks that aim to learn effective representations for downstream tasks. In this work, we present a novel approach that adapts denoising diffusion models to pre-train structure-informed protein encoders 1 . Our proposed approach, called DiffPreT , gradually adds noise to both protein sequence and structure to transform them towards random distribution, and then denoises the corrupted protein structure and sequence using a noise prediction network parameterized with the output of the protein encoder. This approach enables the encoder to learn informative representations that capture the inter-atomic interactions within the protein structure, the residue type dependencies along the protein sequence, and the joint effect of sequence and structure variations.

In spite of these advantages, DiffPreT ignores the fact that any protein structure exists as a population of interconverting conformers, and elucidating this conformational heterogeneity is essential for predicting protein function and ligand binding [27]. In both DiffPreT and previous studies, no explicit constraints are added to acquire the structural correlation between different conformers of a specific protein or between structural homologs, which prohibits capturing the conformational energy landscape of a protein [56]. Therefore, to consider the physics underlying the conformational change, we propose Siamese Diffusion Trajectory Prediction ( SiamDiff ) to augment the DiffPreT by maximizing the mutual information between representations of diffusion trajectories of structurallycorrelated conformers ( i.e. , siamese diffusion trajectories). We first adopt a torsional perturbation scheme on the side chain to generate randomly simulated conformers [36]. Then, for each protein, we generate diffusion trajectories for a pair of its conformers. We theoretically prove that the problem can be transformed to the mutual prediction of the trajectories using representations from their counterparts. In this way, the model can keep the advantages of DiffPreT and inject conformer-related information into representations as well.

Both DiffPreT and SiamDiff can be flexibly applied to atom-level and residue-level structures to pre-train protein representations. To thoroughly assess their capabilities, we conduct extensive evaluations of the pre-trained models on a wide range of downstream protein understanding tasks. These tasks encompass protein function annotation, protein-protein interaction prediction, mutational effect prediction, residue structural contributions, and protein structure ranking. In comparison to existing pre-training methods that typically excel in only a subset of the considered tasks, DiffPreT consistently delivers competitive performance across all tasks and at both the atomic and residuelevel resolutions. Moreover, SiamDiff further enhances model performance, surpassing previous state-of-the-art results in terms of mean ranks across all evaluated tasks.

2 DiffPreT: Diffusion Models for Pre-Training

Recently, there have been promising progress on applying denoising diffusion models for protein structure-sequence co-design [54, 2]. The effectiveness of the joint diffusion model on modeling the distribution of proteins suggests that the process may reflect physical and chemical principles underlying protein formation [3, 18], which could be beneficial for learning informative representations. Based on this intuition, in this section, we explore the application of joint diffusion models on pre-training protein encoders in a pre-training and fine-tuning framework.

2.1 Preliminary

Notation. A protein with n r residues (amino acids) and n a atoms can be represented as a sequencestructure tuple P = ( S , R ) . We use S = [ s 1 , s 2 , · · · , s n r ] to denote its sequence with s i as the type of the i -th residue, while R = [ r 1 , r 2 ..., r n a ] ∈ R n a × 3 denotes its structure with r i as the Cartesian coordinates of the i -th atom. To model the structure, we construct a graph for each protein with edges connecting atoms whose Euclidean distance below a certain threshold.

Equivariance. Equivariance is ubiquitous in machine learning for modeling the symmetry in physical systems [67, 75] and is shown to be critical for successful design and better generalization of 3D

!

!

!

"#$

!

"

!

%

· · ·

· · ·

…YLVCGERGFF…

…**********…

"(!

"

|!

"#$

)

&

&

(!

"#$

|!

"

)

…YLVCGERGFF…

…GALGKRP…

Atom3D

PIP

𝝓(𝓢, 𝓡)

𝝓(𝓢, 𝓡)

Pre-Training

Fine-Tuning

MSP

RES

PSR

networks [46]. Formally, a function F : X → Y is equivariant w.r.t. a group G if F ◦ ρ X ( x ) = ρ Y ◦ F ( x ) , where ρ X and ρ Y are transformations corresponding to an element g ∈ G acting on the space X and Y , respectively. The function is invariant w.r.t G if the transformations ρ Y is identity. In this paper, we consider SE(3) group, i.e. , rotations and translations in 3D space.

Problem Definition. Given a set of unlabeled proteins D = {P 1 , P 2 , ... } , our goal is to train a protein encoder ϕ ( S , R ) to extract informative d -dimensional residue representations h ∈ R n r × d and atom representations a ∈ R n a × d that are SE(3)-invariant w.r.t. protein structures R .

2.2 Diffusion Models on Proteins

Diffusion models are a class of deep generative models with latent variables encoded by a forward diffusion process and decoded by a reverse generative process [63]. We use P 0 to denote the groundtruth protein and P t for t = 1 , · · · , T to be the latent variables over T diffusion steps. Modeling the protein as an evolving thermodynamic system, the forward process gradually injects small noise to the data P 0 until reaching a random noise distribution at time T . The reverse process learns to denoise the latent variable towards the data distribution. Both processes are defined as Markov chains:

where q ( P t |P t -1 ) defines the forward process at step t and p θ ( P t -1 |P t ) with learnable parameters θ defines the reverse process at step t . We decompose the forward process into diffusion on protein structures and sequences, respectively. The decomposition is represented as follows:

where the reverse processes use representations a t and h t learned by the protein encoder ϕ θ ( S t , R t ) .

Forward diffusion process q ( P t |P t -1 ) . For diffusion on protein structures, we introduce random Gaussian noises to the 3D coordinates of the structure. For diffusion on sequences, we utilize a Markov chain approach with an absorbing state [MASK], where each residue either remains the same or transitions to [MASK] with a certain probability at each time step [4]. Specifically, we have:

Here, β 1 , ..., β T and ρ 1 , ..., ρ T are a series of fixed variances and masking ratios, respectively. random_mask ( S t -1 , ρ t ) denotes the random masking operation, where each residue in S t -1 is masked with a probability of ρ t at time step t .

Reverse process on structures p θ ( R t -1 |P t ) . The reverse process on structures is parameterized as a Gaussian with a learnable mean µ θ ( P t , t ) and user-defined variance σ t . Given the availability of

R t as an input, we reparameterize the mean µ θ ( P t , t ) following Ho et al. [37]:

where α t = 1 -β t , ¯ α t = ∏ t s =1 α s and the network ϵ θ ( · ) learns to decorrupt the data and should be translation-invariant and rotation-equivariant w.r.t. the structure R t .

To define our noise prediction network, we utilize the atom representations a t (which is guaranteed to be SE(3)-invariant w.r.t. R t by the encoder) and atom coordinates r t (which is SE(3)-equivariant w.r.t. R t ). We build an equivariant output based on normalized directional vectors between adjacent atom pairs. Each edge ( i, j ) is encoded by its length ∥ r t i -r t j ∥ 2 and the representations of two end nodes a t i , a t j , and the encoded score m i,j will be used for aggregating directional vectors. Specifically,

where N t ( i ) denotes the neighbors of the atom i in the corresponding graph of P t . Note that ϵ θ ( P t , t ) achieves the equivariance requirement, as m i,j is SE(3)-invariant w.r.t. R t while r t i -r t j is translation-invariant and rotation-equivariant w.r.t. R t .

Reverse process on sequences p θ ( S t -1 |P t ) . For the reverse process p θ ( S t -1 |P t ) , we adopt the parameterization proposed in [4]. The diffusion trajectory is characterized by the probability q ( S t -1 |S t , ˜ S 0 ) , and we employ a network ˜ pθ to predict the probability of S 0 :

We define the predictor ˜ p θ with residue representations h t . For each masked residue i in S t , we feed its representation h t i to an MLP and predict the type of the corresponding residue type s 0 i in S 0 :

where the softmax function is applied over all residue types.

2.3 Pre-Training Objective

Now we derive the pre-training objective of DiffPreT by optimizing the diffusion model above with the ELBO loss [37]:

Under the assumptions in (2), it can be shown that the objective can be decomposed into a structure loss L ( R ) and a sequence loss L ( S ) (see proof in App. C.2):

Both loss functions can be simplified as follows.

Structure loss L ( R ) . It has been shown in Ho et al. [37] that the loss function can be simplified under our parameterization by calculating KL divergence between Gaussians as weighted L2 distances between means ϵ θ and ϵ (see details in App. C.3):

where the coefficients γ t are determined by the variances β 1 , ..., β t . In practice, we follow Ho et al. [37] to set all weights γ t = 1 for the simplified loss L ( R ) simple .

Since ϵ θ is designed to be rotation-equivariant w.r.t. R t , to make the loss function invariant w.r.t. R t , the supervision ϵ is also supposed to achieve such equivariance. Therefore, we adopt the chain-rule approach proposed in Xu et al. [82], which decomposes the noise on pairwise distances to obtain the modified noise vector ˆ ϵ as supervision. We refer readers to Xu et al. [82] for more details.

Sequence loss L ( S ) . Since we parameterize p θ ( S t -1 |P t ) with ˜ p θ ( ˜ S 0 |P t ) and q ( S t -1 |S t , ˜ S 0 ) as in (6), it can be proven that the t -th KL divergence term in L ( S ) reaches zero when ˜ p θ ( ˜ S 0 |P t ) assigns

all mass on the ground truth S 0 (see proof in App. C.4). Therefore, for pre-training, we can simplify the KL divergence to the cross-entropy between the correct residue type s 0 i and the prediction:

where CE ( · , · ) denotes the cross-entropy loss.

The ultimate training objective is the sum of simplified structure and sequence diffusion losses:

2.4 Two-Stage Noise Scheduling

Previous studies on scheduled denoising autoencoders on images [23] have shown that large noise levels encourage the learning of coarse-grained features, while small noise levels require the model to learn fine-grained features. We observe a similar phenomenon in structure diffusion, as depicted in Fig. 2. The diffusion loss with large noise at a fixed scale (orange) is smaller than that with small noise at a fixed scale (green), indicating the higher difficulty of structure diffusion with small noises. Interestingly, the opposite is observed in sequence diffusion, where the denoising accuracy with small noise (green) is higher than that with large noise (orange). This can be attributed to the joint diffusion effects on protein sequences and structures. The addition of large noise during joint diffusion significantly disrupts protein structures, making it more challenging to infer the correct protein sequences. For validation, we com-

0.0

0.2

0.4

0.6

0.8

1-st stage - large noise

2-nd stage - small noise

Joint diffusion (two-stage)

Sequence diffusion (two-stage)

Structure diffusion (two-stage)

Joint diffusion (large noise)

Joint diffusion (small noise)

0

10

20

30

40

50

Epoch

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Sequence Denoising Accuracy

pare the loss of sequence-only diffusion pre-training (blue) with joint diffusion pre-training (red) in Fig. 2. Sequence diffusion achieves higher denoising accuracy than joint diffusion when using uncorrupted structures, supporting our hypothesis.

Unlike recent denoising pre-training methods that rely on a fixed noise scale as a hyperparameter [85], DiffPreT incorporates a denoising objective at various noise levels to capture both coarse- and finegrained features and consider the joint diffusion effect explained earlier. Both granularities of features are crucial for downstream tasks, as they capture both small modifications for assessing protein structure quality and large changes leading to structural instability. In our implementation, we perform diffusion pre-training using T = 100 noise levels (time steps). Following the intuition of learning coarse-grained features before fine-grained ones in curriculum learning [8], we divide the pre-training process into two stages: the first stage focuses on large noise levels ( t = 10 , ..., 100 ), while the second stage targets small noise levels ( t = 1 , ..., 9 ). As observed in Fig. 2, structure diffusion becomes more challenging and sequence diffusion becomes easier during the second stage, as we discussed earlier. However, even with this two-stage diffusion strategy, there remains a significant gap between the accuracy of sequence diffusion and joint diffusion. We hypothesize that employing protein encoders with larger capacities could help narrow this gap, which is left as our future works.

3 SiamDiff: Siamese Diffusion Trajectory Prediction

Through diffusion models on both protein sequences and structures, the pre-training approach proposed in Sec. 2 tries to make representations capture (1) the atom- and residue-level spatial interactions and (2) the statistical dependencies of residue types within a single protein. Nevertheless, no constraints or supervision have been added for modeling relations between different protein structures, especially different conformers of the same protein. Generated under different environmental factors, these conformers typically share the same protein sequence but different structures due to side chain rotation driving conformational variations. These conformers' properties are highly correlated [58], and their representations should reflect this correlation.

In this section, we introduce Siamese Diffusion Trajectory Prediction ( SiamDiff ), which incorporates conformer-related information into DiffPreT by maximizing mutual information (MI) between diffusion trajectory representations of correlated conformers. We propose a scheme to generate simulated conformers (Sec. 3.1), generate joint diffusion trajectories (Sec. 3.2), and transform MI maximization into mutual denoising between trajectories (Sec. 3.3), sharing similar loss with DiffPreT.

3.1 Conformer Simulation Scheme

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 P = ( S , R ) , we consider it as the native state P 1 = ( S 1 , R 1 ) and generate a correlated conformer P 2 = ( S 2 , R 2 ) by

𝓟

𝟏

𝟎

𝓟

𝟏

𝒕$𝟏

𝓟

𝟏

𝒕

𝓟

𝟏

𝑻

· · ·

· · ·

𝓟

𝟐

𝟎

𝓟

𝟐

𝒕$𝟏

𝓟

𝟐

𝒕

𝓟

𝟐

𝑻

· · ·

· · ·

𝓟

𝟐

Torsional

Perturbation

𝓟 = (𝓢, 𝓡)

Identity

𝓟

𝟏

…YLVCGERGFF…

…**********…

…YLVCGERGFF…

…**********…

randomly perturbing the protein structure. That is, we set S 2 to be the same as S 1 , and R 2 is obtained by applying a perturbation function perturb ( R 1 , ϵ ) to the original residue structure, where ϵ ∈ [0 , 2 π ) n r × 4 is a noise vector drawn from a wrapped normal distribution [15]. The perturb ( · , · ) 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].

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 P 0: T 1 and P 0: T 2 for conformers P 1 and P 2 , respectively, using the joint diffusion process outlined in Sec. 2. For example, starting from P 0 1 = P 1 = ( S 1 , R 1 ) , we use the joint diffusion on structures and sequences to define the diffusion process q ( P 1: T 1 |P 0 1 ) = q ( R 1: T 1 |R 0 1 ) q ( S 1: T 1 |S 0 1 ) . We derive trajectories on structures R 1: T 1 using the Gaussian noise in (3) and derive the sequence diffusion process S 1: T using the random masking in (6). In this way, we define the trajectory P 0: T 1 = { ( S t 1 , R t 1 ) } T t =0 for P 1 and can derive the siamese trajectory P 0: T 2 = { ( S t 2 , R t 2 ) } T t =0 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: L = 1 2 ( L (2 → 1) + L (1 → 2) ) , where

with b → a being either 2 → 1 or 1 → 2 and P 0: T b being representations of the trajectory P 0: T b .

The two terms share the similar formula as the ELBO loss in (8). Take L (2 → 1) for example. Here q ( P t -1 1 |P t 1 , P 0 1 ) is a posterior analytically tractable with our definition of each diffusion step q ( P t 1 |P t -1 1 ) in (3) and (6). The reverse process is learnt to generate a less noisy state P t -1 1 given the current state P t 1 and representations of the siamese trajectory P 0: T 2 , 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 P 0: T 2 (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 P 1 and P 2 share information about the same protein, the whole trajectory of P 2 could provide too many clues for denoising towards P t -1 1 , making the pre-training task trivial. To address this issue, we parameterize p ( P t -1 1 |P t 1 , P 0: T 2 ) with p θ ( P t -1 1 |P t 1 , P t 2 ) . For diffusion on sequences, we further guarantee that the same set of residues are masked in S t 1 and S t 2 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

where L ( · ,b → a ) simple is the loss term defined by predicting P 0: T a from P 0: T b (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 P 2 = ( S 2 , R 2 ) is defined by S 2 = S 1 , R 2 = R 1 + ϵ , where ϵ ∈ R n a × 3 is the noise drawn from a Gaussian distribution.

4 Related Work

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

perform joint diffusion on both modalities. Note that given a sequence S and a structure R that exist in the nature with high probability, the sequence-structure tuple P = ( S , 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.

5 Experiments

5.1 Experimental Setups

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 Gligorijevi´ c 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 atomand 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 Fmax and AUPR for EC, AUROC for PIP and MSP, Spearman's ρ for PSR, and micro-averaged accuracy for RES. More details about experimental setup can be found in App. D.

5.2 Experimental Results

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.

5.3 Ablation Study

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

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

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.

6 Conclusions

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.

Acknowledgments

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.

References

Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Diffusion-based representation learning. arXiv preprint arXiv:2105.14257 , 2021.

Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019 , 2022.

Christian B Anfinsen. The formation and stabilization of protein structure. Biochemical Journal , 128(4):737, 1972.

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Minkyung Baek, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, R Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using a three-track neural network. Science , 373(6557):871-876, 2021.

Inigo Barrio-Hernandez, Jingi Yeo, Jürgen Jänes, Tanita Wein, Mihaly Varadi, Sameer Velankar, Pedro Beltrao, and Martin Steinegger. Clustering predicted structures at the scale of the known protein universe. bioRxiv , pages 2023-03, 2023.

Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International conference on machine learning , pages 531-540. PMLR, 2018.

Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In International Conference on Machine Learning , 2009.

B Chandra and Rajesh Kumar Sharma. Adaptive noise schedule for denoising autoencoder. In International conference on neural information processing , pages 535-542. Springer, 2014.

Can Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein selfsupervised learning. arXiv preprint arXiv:2204.04213 , 2022.

Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. Wavegrad: Estimating gradients for waveform generation. arXiv preprint arXiv:2009.00713 , 2020.

Tianlong Chen, Chengyue Gong, Daniel Jesus Diaz, Xuxi Chen, Jordan Tyler Wells, qiang liu, Zhangyang Wang, Andrew Ellington, Alex Dimakis, and Adam Klivans. Hotprotein: A novel framework for protein thermostability prediction and editing. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum? id=YDJRFWBMNby .

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15750-15758, 2021.

Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. International Conference on Learning Representations (ICLR) , 2023.

Christian Dallago, Jody Mou, Kadina E Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks in fitness landscape inference for proteins. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2) , 2021.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Ken A Dill and Justin L MacCallum. The protein-folding problem, 50 years on. science , 338 (6110):1042-1046, 2012.

Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications on Pure and Applied Mathematics , 36 (2):183-212, 1983.

Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Wang Yu, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. Prottrans: Towards cracking the language of lifes code through self-supervised deep learning and high performance computing. IEEE Transactions on Pattern Analysis and Machine Intelligence , pages 1-1, 2021. doi: 10.1109/TPAMI.2021.3095381.

Bent Fuglede and Flemming Topsoe. Jensen-shannon divergence and hilbert space embedding. In International Symposium onInformation Theory, 2004. ISIT 2004. Proceedings. , page 31. IEEE, 2004.

Pablo Gainza, Freyr Sverrisson, Frederico Monti, Emanuele Rodola, D Boscaini, MM Bronstein, and BE Correia. Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods , 17(2):184-192, 2020.

Krzysztof J Geras and Charles Sutton. Scheduled denoising autoencoders. arXiv preprint arXiv:1406.3269 , 2014.

Vladimir Gligorijevi´ c, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, et al. Structure-based protein function prediction using graph convolutional networks. Nature communications , 12(1):1-14, 2021.

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning . MIT press, 2016.

Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Proceedings of the 27th International Conference on Neural Information Processing Systems-Volume 2 , pages 2672-2680, 2014.

Barry J Grant, Alemayehu A Gorfe, and J Andrew McCammon. Large conformational changes in proteins: signaling and other functions. Current opinion in structural biology , 20(2):142-147, 2010.

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271-21284, 2020.

Yuzhi Guo, Jiaxiang Wu, Hehuan Ma, and Junzhou Huang. Self-supervised pre-training for protein embeddings using tertiary structures. In AAAI , 2022.

Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 297-304. JMLR Workshop and Conference Proceedings, 2010.

Michael U Gutmann and Aapo Hyvärinen. Noise-contrastive estimation of unnormalized statistical models, with applications to natural image statistics. Journal of machine learning research , 13(2), 2012.

Kaveh Hassani and Amir Hosein Khasahmadi. Contrastive multi-view representation learning on graphs. In International Conference on Machine Learning , pages 4116-4126. PMLR, 2020.

Pedro Hermosilla and Timo Ropinski. Contrastive representation learning for 3d protein structures. arXiv preprint arXiv:2205.15675 , 2022.

Pedro Hermosilla, Marco Schäfer, Matˇ ej Lang, Gloria Fackelmann, Pere Pau Vázquez, Barbora Kozlíková, Michael Krone, Tobias Ritschel, and Timo Ropinski. Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. International Conference on Learning Representations , 2021.

R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint arXiv:1808.06670 , 2018.

Bosco K Ho and David A Agard. Probing the flexibility of large conformational changes in protein structures through local perturbations. PLoS computational biology , 5(4):e1000343, 2009.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning , pages 8867-8887, 2022.

John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent Xue, Fritz Obermeyer, Andrew Beam, et al. Illuminating protein space with a programmable generative model. bioRxiv , 2022.

Bowen Jing, Stephan Eismann, Pratham N. Soni, and Ron O. Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=1YLJDvSx6J4 .

Bowen Jing, Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106.03843 , 2021.

Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S. Jaakkola. Torsional diffusion for molecular conformer generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=w6fj2r62r_H .

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning , pages 5361-5370. PMLR, 2020.

Andriy Kryshtafovych, Torsten Schwede, Maya Topf, Krzysztof Fidelis, and John Moult. Critical assessment of methods of protein structure prediction (casp)-round xiii. Proteins: Structure, Function, and Bioinformatics , 87(12):1011-1020, 2019.

Yann LeCun and Ishan Misra. Self-supervised learning: The dark matter of intelligence, 2021. URL https://ai.facebook.com/blog/ self-supervised-learning-the-dark-matter-of-intelligence/ .

Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. arXiv preprint arXiv:2205.14217 , 2022.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv , 2022.

Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. In International Conference on Learning Representations , 2021.

Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)invariant denoising distance matching. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=CjTHVo1dvR .

Amy X Lu, Haoran Zhang, Marzyeh Ghassemi, and Alan M Moses. Self-supervised contrastive learning of protein representations by mutual information maximization. BioRxiv , 2020.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. 2022. URL https://openreview.net/forum?id=jSorGn2Tjg .

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. bioRxiv , 2021. doi: 10.1101/2021.07.09.450648. URL https://www.biorxiv.org/content/10. 1101/2021.07.09.450648v1 .

G Ulrich Nienhaus, Joachim D Müller, Ben H McMahon, and Hans Frauenfelder. Exploring the conformational energy landscape of proteins. Physica D: Nonlinear Phenomena , 107(2-4): 297-311, 1997.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Clare M O'Connor, Jill U Adams, and Jennifer Fairman. Essentials of cell biology. Cambridge, MA: NPG Education , 1:54, 2010.

Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with tape. Advances in neural information processing systems , 32, 2019.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538. PMLR, 2015.

Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences , 118(15), 2021.

Maxim V Shapovalov and Roland L Dunbrack. A smoothed backbone-dependent rotamer library for proteins derived from adaptive kernel density estimates and regressions. Structure , 19(6):844-858, 2011.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265, 2015.

Vignesh Ram Somnath, Charlotte Bunne, and Andreas Krause. Multi-scale representation learning on proteins. Advances in Neural Information Processing Systems , 34, 2021.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Freyr Sverrisson, Jean Feydy, Bruno E Correia, and Michael M Bronstein. Fast end-to-end learning on protein surfaces. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15272-15281, 2021.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Raphael JL Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon Anderson, Stephan Eismann, et al. Atom3d: Tasks on molecules in three dimensions. arXiv preprint arXiv:2012.04035 , 2020.

Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motifscaffolding problem. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=6TxBxqNME1Y .

Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, et al. Alphafold protein structure database: massively expanding the structural coverage of protein-sequence space with high-accuracy models. Nucleic acids research , 2021.

Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.14734 , 2022.

Junmei Wang, Romain M Wolf, James W Caldwell, Peter A Kollman, and David A Case. Development and testing of a general amber force field. Journal of computational chemistry , 25 (9):1157-1174, 2004.

Limei Wang, Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning protein representations via complete 3d graph networks. arXiv preprint arXiv:2207.12600 , 2022.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv , 2022.

Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. Advances in Neural Information Processing Systems , 31, 2018.

Fang Wu, Qiang Zhang, Dragomir Radev, Yuyang Wang, Xurui Jin, Yinghui Jiang, Zhangming Niu, and Stan Z Li. Pre-training of deep protein models with molecular dynamics simulations for drug binding. arXiv preprint arXiv:2204.08663 , 2022.

Kevin E Wu, Kevin K Yang, Rianne van den Berg, James Y Zou, Alex X Lu, and Ava P Amini. Protein structure generation via folding diffusion. arXiv preprint arXiv:2209.15611 , 2022.

Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. arXiv preprint arXiv:2209.00865 , 2022.

Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graphlevel representation learning with local and global structure. In International Conference on Machine Learning , pages 11548-11558. PMLR, 2021.

Minghao Xu, Yuanfan Guo, Yi Xu, Jian Tang, Xinlei Chen, and Yuandong Tian. Eurnet: Efficient multi-range relational modeling of spatial multi-relational data. arXiv preprint arXiv:2211.12941 , 2022.

Minghao Xu, Zuobai Zhang, Jiarui Lu, Zhaocheng Zhu, Yangtian Zhang, Chang Ma, Runcheng Liu, and Jian Tang. Peer: A comprehensive and multi-task benchmark for protein sequence understanding. arXiv preprint arXiv:2206.02096 , 2022.

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id= PzcvxEMzvQC .

Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. bioRxiv , 2022.

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.

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 .

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 .

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.

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.

A More Related Works

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

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 θ ( P t -1 1 |P t 1 , P t 2 ) 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 P 0: T 1 and P 0: T 2 to denote the representation of the siamese trajectories P 0: T 1 and P 0: T 2 , respectively. Different from the generation process in traditional diffusion models, the parameterization of p θ ( P t -1 1 |P t 1 , P t 2 ) should inject information from P t 2 . Therefore, we use the extracted residue and atom representations (denoted as a t 2 and h t 2 ) of P t 2 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.

Generation process on protein structures. As in (4), modeling the generation process of protein structures is to model the noise on R t 1 and gradually decorrupt the noisy structure. This can be parameterized with a noise prediction network ϵ θ ( P t 1 , P t 2 , t ) that is translation-invariant and rotationequivariant w.r.t. R t 1 . Besides, the noise applied on R t 1 should not change with transformations on R t 2 , so ϵ θ should be SE(3)-invariant w.r.t. R t 2 .

To achieve these goals, we build our noise prediction network with atom representations a t 2 (which is SE(3)-invariant w.r.t. R t 2 ) and atom coordinates r t 1 (which is SE(3)-equivariant w.r.t. R t 1 ). We define an equivariant output similarly as in DiffPreT. Specifically, we have

where N t 1 ( i ) denotes the neighbors of the atom i in the corresponding graph of P t 1 . Note that ϵ θ ( P t 1 , P t 2 , t ) achieves the equivariance requirement, as m i,j is SE(3)-invariant w.r.t. R t 1 and R t 2 while r t 1 i -r t 1 j is translation-invariant and rotation-equivariant w.r.t. R t 1 .

Generation process on protein sequences. As in (7), the generation process on sequences aims to predict masked residue types in S 0 1 with a predictor ˜ p θ . In our setting of mutual prediction, we define the predictor based on representations of the same residues in S t 2 , which are also masked. Hence, for each masked residue i in S t 2 , we feed its representation h t 2 i to an MLP and predict the type of the corresponding residue type s 0 1 i in S 0 1 :

where the softmax function is applied over all residue types.

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 L (2 → 1) for example and its counterpart can be derived in the same way. The objective can be decomposed into a structure loss L ( R , 2 → 1) and a sequence loss L ( S , 2 → 1) :

Based on the derivation in Sec. 2.3, the structure loss L ( R , 2 → 1) can be simplified as

and the sequence loss L ( S , 2 → 1) can be simplified as

Then, the final objective in Sec. 3.3 can be easily derived.

C Proofs

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.

C.1 Proof of Proposition 1

For notations, we use the bold symbol to denote the representation of an object and use P 0: T 1 and P 0: T 2 to denote the corresponding random variables of representations of the siamese trajectories P 0: T 1 and P 0: T 2 .

Proposition 1 With some approximations, the mutual information between representations of two siamese trajectories is lower bounded by:

where C is a constant independent of our encoder and the term from trajectory P 0: T b to P 0: T a is defined as

with b → a being either 2 → 1 or 1 → 2 .

Proof. First, the mutual information between representations of two trajectories is defined as:

where the joint distribution is defined as p ( P 0: T 1 , P 0: T 2 ) = p ( P 0 1 , P 0 2 ) q ( P 1: T 1 | P 0 1 ) q ( P 1: T 2 | P 0 2 ) . Next, we can derive a lower bound with this definition:

However, since the distribution of representations are intractable to sample for optimization, we instead sample the trajectories P 0: T 1 and P 0: T 2 from our defined diffusion process, i.e. , p ( P 0: T 1 , P 0: T 2 ) = p ( P 0 1 , P 0 2 ) q ( P 1: T 1 |P 0 1 ) q ( P 1: T 2 |P 0 2 ) . 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:

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:

where we merge the term p ( P 0 1 |P 1 1 , P 0: T 2 ) into the sum of KL divergences for brevity and use C (2 → 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

and 2) the conditional independence of the generation process

it can be proved that

where the three loss terms are defined as

with b → a referring to the term from trajectory P 0: T b to P 0: T a .

Proof. Let L ( · ) t to denote the t-th KL divergence term in L ( · ) . Then, we have

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). □

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

and the reverse process

the structure loss function

can be simplified to

where γ t = β t 2 α t (1 -¯ α t -1 ) with α t = 1 -β t , ¯ α t = ∏ t s =1 α s and b → a is either 2 → 1 or 1 → 2 .

Proof. First, we prove q ( R t a |R 0 a ) = N ( R t a ; √ ¯ α t R 0 a , (1 -¯ α t ) I ) . Let ϵ i be the standard Gaussian random variable at time step i . Then, we have

which suggests that the mean of R a is ¯ α t R a and the variance matrix is ( α t α t -1 ...α 2 β 1 + · · · + α t -1 β t -1 + β t ) I = (1 -¯ α ) I .

Next, we derive the posterior distribution as:

Let ˜ β t = 1 -¯ α t -1 1 -¯ α t β t , then the t -th KL divergence term can be written as:

C.4 Proof of Simplified Sequence Loss

Now we show the equivalence of optimizing sequence loss L ( S ,b → a ) and the masked residue type prediction problem on S 0 a .

Proposition 4 Given the definition of reverse process on protein sequences

the sequence loss L ( S ,b → a ) reaches zero when ˜ p θ ( ˜ S 0 a |P t a , P t b ) puts all mass on the ground truth S 0 a .

Proof. The loss function can be written as:

where Z is the normalization constant. Hence, when ˜ p θ ( ˜ S 0 a |P t a , P t b ) puts all mass on the ground truth S 0 a , the distribution p θ ( S t -1 a |P t a , P t b ) will be identical with q ( S t -1 a |S t a , S 0 a ) , which makes the KL divergence become zero. □

D Experimental Details

In this section, we introduce the details of our experiments. All these methods are developed based on PyTorch and TorchDrug [88].

Downstream benchmark tasks. For downstream evaluation, we adopt the EC prediction task [24] and four ATOM3D tasks [68].

Enzyme Commission (EC) number prediction task aims to predict EC numbers of proteins which describe their catalysis behavior in biochemical reactions. This task is formalized as 538 binary classification problems. We adopt the dataset splits from Gligorijevi´ c et al. [24] and use the test split with 95% sequence identity cutoff following Zhang et al. [87].

Protein Interface Prediction (PIP) requires the model to predict whether two amino acids from two proteins come into contact when the proteins bind (binary classification). The protein complexes of this benchmark are split with 30% sequence identity cutoff.

Mutation Stability Prediction (MSP) task seeks to predict whether a mutation will increase the stability of a protein complex or not (binary classification). The benchmark dataset is split upon a 30% sequence identity cutoff among different splits.

Residue Identity (RES) task studies the structural role of an amino acid under its local environment. A model predicts the type of the center amino acid based on its surrounding atomic structure. The environments in different splits are with different protein topology classes.

Protein Structure Ranking (PSR) predicts global distance test scores of structure predictions submitted to the Critical Assessment of Structure Prediction (CASP) [47] competition. This dataset is split according to the competition year.

Graph construction. For atom graphs, we connect atoms with Euclidean distance lower than a distance threshold. For PSR and MSP tasks, we remove all hydrogen atoms following Jing et al. [42]. For residue graphs, we discard all non-alpha-carbon atoms and add three different types of directed edges: sequential edges, radius edges and K-nearest neighbor edges. For sequential edges, two atoms are connected if their sequential distance is below a threshold and these edges are divided into different types according to these distances. For two kinds of spatial edges, we connect atoms according to Euclidean distance and k-nearest neighbors. We further apply a long range interaction filter that removes edges with low sequential distances. We refer readers to Zhang et al. [87] for more details.

Atom-level backbone models. To adapt GearNet-Edge to atom-level structures with moderate computational cost, we construct the atom graph by using only the spatial edge with the radius

d radius = 4 . 5 Å. We concatenate one-hot features of atom types and residue types as node features and concatenate (1) one-hot features of residue types of end nodes, (2) one-hot features of edge types, (3) one-hot features of sequential distance, (4) spatial distance as edge features. The whole model is composed of 6 message passing layers with 128 hidden dimensions and ReLU activation function. For edge message passing, we employ the discretized angles to determine the edge types on the line graph. The final prediction is performed upon the hidden representation concatenated across all layers.

Residue-level backbone models. We directly borrow the best hyperparameters reported in the original paper of GearNet-Edge [87]. We adopt the same configuration of relational graph construction, i.e. , the sequential distance threshold d seq = 3 , the radius d radius = 10 . 0 Å, the number of neighbors k = 10 and the long range interaction cutoff d long = 5 . We use one-hot features of residue types as node features and concatenate (1) features of end nodes, (2) one-hot features of edge types, (3) one-hot features of sequential distance, (4) spatial distance as edge features. Then we use 6 message passing layers with 512 hidden dimensions and ReLU as the activation function. For edge message passing, the edge types on the line graph are determined by the discretized angles. The hidden representations in each layer of GearNet will be concatenated for the final prediction.

Baseline pre-training methods. Here we briefly introduce the considered baselines. Multiview Contrast aims to maximize the mutual information between correlated views, which are extracted by randomly chosen augmentation functions to capture protein sub-structures. Residue type, distance, angle and dihedral prediction masks single residues, single edges, edge pairs and edge triplets, respectively, and then predict the corresponding properties. Denoising score matching performs denoising on noised pairwise distance matrices based on the learnt representations.

For all baselines in [87], we adopt the original configurations. For Multiview Contrast, we use subsequence cropping that randomly extracts protein subsequences with no more than 50 residues and space cropping that takes all residues within a 15 Å Euclidean ball with a random center residue. Then, either an identity function or a random edge masking function with mask rate equal to 0.15 is applied for constructing views. The temperature τ in the InfoNCE loss function is set as 0.07. We set the number of sampled items in each protein as 256 for Distance Prediction and as 512 for Angle and Dihedral Prediction. The mask rate for Residue Type Prediction is set as 0.15. When masking a residue on atom graphs, we discard all non-backbone atoms and set the residue features as zero. Since the backbone models and tasks in our paper are quite different with those in Guo et al. [29], we re-implement the method on our codebase. We consider 50 different noise levels log-linearly ranging from 0.01 to 10.0.

In DiffPreT, for structure diffusion, we use a sigmoid schedule for variances β t with the lowest variance β 1 = 1 e -4 and the highest variance β T = 0 . 1 . For sequence diffusion, we simply set the cumulative transition probability to [MASK] over time steps as a linear interpolation between minimum mask rate 0.15 and maximum mask rate 1.0. The number of diffusion steps is set as 100. In SiamDiff, we adopt the same hyperparameters for multimodal diffusion models. We set the variance of torsional perturbation noises as 0 . 1 π on the atom level and that of Gaussian perturbation noises as 0 . 3 on the residue level when constructing the correlated conformer.

All other optimization configurations for these pre-training methods are reported in Table 5. All methods are pre-trained on 4 Tesla A100 GPUs and Table 5 reports the batch sizes on each GPU.

Fine-tuning on downstream tasks. For all models on all downstream tasks, we apply a three-layer MLP head for prediction, the hidden dimension of which is set to the dimension of model outputs. The number of used gpus and batch sizes for each model are chosen according the memory limit. All residue-level tasks are run on 4 V100 GPUs while all atom-level tasks are run on A100 GPUs.

Evaluation metrics. We clarify the definitions of Fmax (used in EC), global Spearman's ρ (used in PSR) and mean Spearman's ρ (used in PSR) as below:

Fmax denotes the protein-centric maximum F-score. It first computes the precision and recall for each protein at a decision threshold t ∈ [0 , 1] :

where f denotes a functional term in the ontology, T i is the set of experimentally determined functions for protein i , P i ( t ) is the set of predicted functions for protein i whose scores are greater or equal to t , and ✶ [ · ] represents the indicator function. After that, the precision and recall are averaged over all proteins:

where N denotes the total number of proteins, and M ( t ) denotes the number of proteins which contain at least one prediction above the threshold t , i.e. , | P i ( t ) | > 0 .

Based on these two metrics, the Fmax score is defined as the maximum value of F-measure over all thresholds:

Global Spearman's ρ for PSR measures the correlation between the predicted global distance test (GDT_TS) score and the ground truth. It computes the Spearman's ρ between the prediction and the ground truth over all test proteins without considering the different biopolymers that these proteins lie in.

Mean Spearman's ρ for PSR also measures the correlation between GDT_TS predictions and the ground truth. However, it first splits all test proteins into multiple groups based on their corresponding biopolymers, then computes the Spearman's ρ within each group, and finally reports the mean Spearman's ρ over all groups.

E Results of Pre-Training on Different Sizes of Datasets

In the main paper, we followed the setting in Zhang et al. [87] and used AlphaFold Database v1 as our pre-training dataset for fair comparison. Here, we investigate the impact of pre-training on different dataset sizes. Since previous work by Zhang et al. [87] showed minimal differences between using experimental or predicted structures, we conduct experiments on the AlphaFold Database in this section and do not use PDB as our pre-training dataset. We utilize the preprocessed clustered

AlphaFold Database provided in [6], which includes 2.2M non-singleton clusters with an average of 13.2 proteins per cluster and an average pLDDT of 71.59. For each cluster, we use the representative structure from [6]. To explore the effects of dataset size, we pre-train our encoder using 10K, 50K, 100K, 500K, and 2.2M clusters in the database. The results, shown in Table 7, reveal a general trend of increased performance with larger datasets. However, for certain tasks like MSP, the performance does not consistently improve, possibly due to the limited size and variability of the downstream datasets. Overall, scaling the model to the entire AlphaFold Database holds promise for achieving performance gains.

F Results of Multimodal Pre-Training Baselines

In Sec. 5, we include all pre-training baselines from existing works, which solely focus on unimodal pre-training objectives and overlook the joint distribution of sequences and structures. In this section, we introduce three additional pre-training baselines that leverage both sequence and structure information. These baselines combine residue type prediction with distance/angle/dihedral prediction objectives. Following established settings, we pre-train the encoder and present the results in Table 8. While the Residue Type & Angle Prediction method achieves higher performance on PSR, there are still substantial gaps compared to our DiffPreT and SiamDiff across other tasks. Notably, the introduction of geometric property prediction tasks leads to a drop in performance on RES, indicating that a simple combination of pre-training objectives diminishes the benefits of each individual objective. Once again, these findings underscore the effectiveness of our methods in modeling the joint distribution of protein sequences and structures.

G Results of Different Diffusion Models for Pre-Training

In Sec. 2, we consider joint diffusion models on protein sequences and structures for pre-training. In this section, we explore the performance of different diffusion models when applied for pre-training. The results are shown in Table 9.

First, we simply run diffusion models on protein sequences and structures for pre-training, both of which achieve improvement compared with the baseline GearNet-Edge. By combining two diffusion models, DiffPreT can achieve better performance on PIP and PSR. This advantage will be further

increased after using siamese diffusion trajectory prediction as shown in Table 3. It can be observed that sequence diffusion achieves better performance than DiffPreT due to the consistency of objectives between pre-training and RES tasks.

Besides, we also consider diffusion models on torsional angles of protein side chains for pre-training. This method has shown its potential in the protein-ligand docking task [15]. For each residue, we randomly select one side-chain torsional angle and add some noises drawn from wrapped Gaussian distribution during the diffusion process. Then, we predict the added noises with the extracted atom representations corresponding to the torsional angle. In Table 9, we can see clear improvements on PIP and PSR tasks compared with GearNet-Edge. This suggests that it would be a promising direction to explore more different diffusion models for pre-training, e.g. , diffusion models on backbone dihedral angles.

H Results of Pre-Training GVP

To study the effect of our proposed pre-training methods on different backbone models, we show the pre-training results on GVP [41, 42] in this section.

Setup. GVP constructs atom graphs and adds vector channels for modeling equivariant features. The original design only includes atom types as node features, which makes pre-training tasks with residue type prediction very difficult to learn. To address this issue, we slightly modify its architecture to add the embedding of atom and corresponding residue types as atom features. Then, the default configurations in Jing et al. [42] are adopted. We construct an atom graph for each protein by drawing edges between atoms closer than 4 . 5 Å. Each edge is featured with a 16-dimensional Gaussian RBF encoding of its Euclidean distance. We use five GVP layers and hidden representations with 16 vector and 100 scalar channels and use ReLU and identity for scalar and vector activation functions, respectively. The dropout rate is set as 0.1. The final atom representations are followed by two mean pooling layers to obtain residue and protein representations, respectively. All other hyperparameters for pre-training and downstream tasks are the same as those in App. D.

Experimental results. The results are shown in Table 10. Among all pre-training methods, SiamDiff, angle prediction and dihedral prediction are the top three. This is different from what we observe in Table 1, where residue type and distance prediction are more competitive baselines. We hypothesize that this is because GearNet-Edge includes angle information in the encoder while GVP does not. Therefore, GVP will benefit more from pre-training methods with supervision on angles. Nevertheless, we find that SiamDiff is the only method that performs well on different backbone models.

I Results on Protein Engineering Task

To further prove the effectiveness of our proposed pre-training methods, we add experiments on the GB1 dataset from FLIP [16]. As this is a protein engineering task with mutated sequences, we assume that the backbone structure remains unchanged after mutation, to save costs in generating mutant structures. We only keep CA atoms in the wild type protein structure as the input to the encoder. We benchmark residue-level methods in Table 11 in the attached file, alongside CNN and ESM-1b baselines from the FLIP paper.

According to Table 11, we observe that modeling structural information is beneficial compared with using only sequential information, even under the assumption that all mutants share the same backbone structure. Among all pre-training methods, SiamDiff demonstrates the most significant improvements over the baseline, once again validating the effectiveness of our method.

J Effects of Approximate Conformer Generation

In SiamDiff, we introduce approixmate conformer generation mechanism by randomly rotating torsional angles. We hypothesize that highly realistic conformers are not vital for better representations. To confirm, an extra rebuttal experiment is performed in Table 12. Instead of random perturbation, we sample from a rotamer library [62] based on residue types and backbone angles. Table 12 shows random torsional perturbation still outperforms sampling from a rotamer library in most tasks, confirming our hypothesis. This can be attributed to the fact that the objective of pre-training is to learn common information between diverse views through mutual prediction, as SimCLR and SimSiam. Considering this perspective, introducing random torsional noise allows us to generate more diverse conformers compared to solely relying on realistic conformer distributions.

(A) GearNet-Edge representations with random initialization

AF-Q6P3I7-F1-model_v1.pdb

AF-Q6P3I8-F1-model_v1.pdb

AF-Q6P3J0-F1-model_v1.pdb

AF-Q6P3J2-F1-model_v1.pdb

Noise Level

(B) GearNet-Edge representations pre-trained with first-stage SiamDiff

AF-Q6P3I7-F1-model_v1.pdb

AF-Q6P3I8-F1-model_v1.pdb

AF-Q6P3J0-F1-model_v1.pdb

AF-Q6P3J2-F1-model_v1.pdb

Noise Level

(C) GearNet-Edge representations pre-trained with two-stage SiamDiff

AF-Q6P3I7-F1-model_v1.pdb

AF-Q6P3I8-F1-model_v1.pdb

AF-Q6P3J0-F1-model_v1.pdb

AF-Q6P3J2-F1-model_v1.pdb

Noise Level

K Visualization of SiamDiff Embeddings upon Noises

To explore pre-training insights, we visualize UMAP representations of 4 random AlphaFold DB proteins in Fig. 4. Several interesting phenomena can be observed:

Randomly initiated representations in Fig. 4A form a clear, continuously color-changing trajectory (blue to red). This confirms that the forward diffusion process gradually adds noise to proteins, leading to smooth changes in their representations, as expected for diffusion models.

After pre-training with large noise scales , the encoder maintains the color smoothness of the trajectory, which is desired for effective denoising during the backward diffusion process. Intriguingly, pre-training narrows the trajectory compared to the broader trajectory without pre-training, particularly at the two ends. This suggests that first-stage pre-training clusters proteins with similar levels of added noise, even for large and diverse noises. This clustering property proves useful for detecting large perturbations in downstream tasks, such as mutation stability prediction in MSP, as opposed to the diverse representation distributions in Fig. 4A.

Continuing with small noise scale pre-training , the trajectory becomes much narrower in the middle and even breaks for some proteins. This indicates that by focusing on only slightly perturbed samples during pre-training, our model becomes capable of discerning proteins with small and large noises, making it more effective for fine-grained downstream tasks like PSR and

PIP. However, the red end of the trajectory is thicker than that in Fig. 4B, which may imply some forgetting behavior in the second-stage pre-training.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to use the diffusion on both the protein sequence and the protein structure for pretraining. Additionally, they also take the correlation between different conformers of the same protein into consideration and maximize the mutual information between their trajectories via mutual denoising. Specifically, different conformers of the same protein are generated by perturbations on the side-chain torsional angles. Experiments on Enzyme Commission classification and four Atom3D tasks demonstrate the benefit of diffusion-based pretraining and mutual denoising.\n\nSTRENGTHS:\n1. The idea of maximizing the mutual information between diffusion trajectories of different conformers is interesting and inspiring. The authors also show the objective can be transformed into a loss function for mutual denoising between two relevant trajectories.\n2. The evaluation is comprehensive. The authors compare with different pretraining strategy and conduct many ablation studies. They also explore the effect of the size of the pretraining data and different diffusion models.\n\nWEAKNESSES:\n1. Random perturbations on the side-chain torsional angles may produce conformations with high energy, which does not reflect the actual \"physics underlying the conformational change\" as claimed in line 51-52. This may also downgrade the effect of the mutual denosing scheme. As shown in the ablation studies, discarding the mutual information maximization scheme only has very minor effect on the performance.\n2. The organization of the paper can be improved. The main insight lies in how to maximize the mutual information between representations of Siamese Trajectories. However, the major illustration of this contribution is put into the Appendix while the authors spend a great deal of time introducing the preliminaries on diffusion.\n\nQUESTIONS:\n1. Do you think the mutual denoising scheme can also work on homologs as you have mentioned the problem in line 50? There might not be a 1-to-1 mapping on the node representations when it comes to homologs since they may have different number of residues.\n2. What do you think is the reason that diffusion-based pretraining produces better results than denoising score matching pretraining in your experiments? In Section 4.4, you have mentioned that perturbing distance matrices may produce negative values, however, it is also natural to do denoising score matching pretraining directly on the absolute coordinates [1][2]. Could the denosing score mathching baseline be underestimated because it is implemented on the distance matrix instead of the absolute coordinates?\n3. Can you theoretically show how much deviation will be introduced by the approximation in line 743 (Appendix C.1)? For example, can you derive the bound of the gap between the lowerbound of mutual information and the loss you derived from the RHS of the approximation?\n4. What is the value of $t$ in the diffusion model in the fine-tuning phase?\n\n[1] Pre-training via Denoising for Molecular Property Prediction\n\n[2] Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs\n\nLIMITATIONS:\nRandom perturbations on the side-chain torsional angles may not conform to the actual distribution of the conformational change.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel pre-training method for proteins by jointly modeling sequences and structures with a diffusion model (DiffPreT). The encoder, which is the noise prediction network, learns representations for the sequence and the structure, respectively. This representation can then be used for downstream tasks. Additionally, the authors sugest an extension to their method to capture the correlation between conformer structures, based on the maximization of mutual information between diffusion trajectories for different simulated conformers, obtained by perturbing side chain rotations.\n\nSTRENGTHS:\nThe authors bring forward some very interesting ideas and the results are promising, possibly signifying a step forwards within the field of protein representations. Moreover, it is clear that an extensive amount of work has gone into running experiments, including ablation studies and additional results in the appendix.\n\nWEAKNESSES:\nThe paper is confusing and unclear in some places, both in terms of text and in terms of mathematical notation. There are some essential parts of the method that are not properly explained and/or motivated, and not all results seem to indicate a significant improvement. Finally, I think the results could be presented in a more diverse way, rather than only showing big tables.\n\nFor comments and suggestions on how to imrpove the manuscript, I refer to the \"Questions\" part.\n\nQUESTIONS:\n- Main questions / comments:\n  1. In the DiffPreT setting (and to some extent also in the SiamDiff setting), it is unclear to me what noise level is used to get the representation for downstream tasks. It is mentioned that mutiple levels of noise are used, but it remains vague wether this means concatenating respresentations at multiple noise levels, all noise levels, or picking a noise level at random when getting the representation for the downstream task. \n  2. Maybe even more generally, could you give some intuition on why representations on noisy structures could serve better for downstream tasks?\n  3. Could you comment on the size of the latent representations that you get? I might have overlooked it, but I missed what the size of $d$ is. This seems like important information, especially when the SiamDiff setting is used where all these representations seem to be concatenated.\n  4. For the siamDiff setting, it is not intuitive to me what kind of distribution is really being modelled. Randomly rotating the side chains without perturbing the backbone, and simply throwing away structures where there are clashes seems a bit messy. Is the resulting something \"biologically relevant\"?\n  5. When considering the standard deviation of some of the results, the improvements are not always statistically significant. Perhaps it would be good to acknowledge this.\n  6. The structuring of the paper is a bit off in some places, e.g. 3.4 is showing results even before all methods are explained, and the discussion in 4.4 is more a \"related works\" section, which has mostly been discussed in the introduction already, and which somewhat disrupts the flow of the paper.\n  7. From the main text, it took me a long time to discover that the encoder $\\phi$ is the noise predicting network, and even longer to see that it is GearNet-Edge (and GVP in the appendix). This could be more clear earlier on in the manuscript.\n  8. In the appendix it is stated that GearnNet-Edge is a structural encoder. However, from the main paper it seems like you get a representation for both the sequence and the structure. How does this work? Are these separate networks?\n  9. All results are presented as big tables, it would help to include some intuitive graphical results. Perhaps visualize the protein structure / sequence reconstructions of the diffusion model, or somehow visualize the representations (for example see if clusters appear when PCA is done on the representations or something similar).\n\n\n- Other questions / comments:\n  10. The term \"encoder\" can be confusing, as the forward process of a diffusion model can also be seen as a series of unparameterized encoders. It would help if this distinction is clear from the beginning.\n  11. Perhaps a missing reference for protein structure-sequence co-design: Lisanza et al. [2023]\n  12. The task explanation in the main paper is very minimal. I appreciate that there is a more extensive description in the appendix, but it could still be much more elaborate.\n  13. Sometimes abbreviations are introduced before their meaning (e.g. EC, ESM).\n  14. The \"Mean Rank\" seems like a strange metric to me that does not add much to all other results. \n  15. The results using GVP are interesting and would maybe be worth including in the main paper, including a discussion on why some scores are much higher than for GearNet-Edge (except for PSR).\n  16. Some discussion on computational cost is missing, either in the main paper or in the appendix.\n\nLIMITATIONS:\nThe authors discuss limitations in the appendix. It would be nice if some of this discussion is transferred to the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

comments: SUMMARY:
In this work, the authors perform a thorough investigation of different pre-training strategies on joint sequence-structure diffusion models for representation learning, rather than generative modeling. They evaluate on an EC prediction task and four tasks from the Atom3D benchmark. They find that joint sequence-structure pre-training with the SiamDiff method for maximizing mutual information between diffusion trajectories of protein conformers consistently outperforms other methods on downstream task evaluation.

STRENGTHS:
Great observation that prior pre-training strategies excel for particular tasks and have shortcomings in others; and while SiamDiff does not yield huge improvements over other strategies, it does consistently outperform them across tasks. 

This work takes a thorough and critical look at the benefits and tradeoffs of considering sequence- or structure-only methods, and systematically evaluates different pre-training approaches.

WEAKNESSES:
Lack of evaluation against other downstream predictors and backbone models for downstream prediction.

QUESTIONS:
The sizable gap in sequence denoising accuracy for sequence vs joint diffusion is very interesting, especially considering that in the ablation, the model without structure diffusion is competitive with the others. Can the authors comment on the additional complexity of joint diffusion compared to sequence-only diffusion, and why this accuracy gap is so pronounced? The authors mention model capacity, but clearly capacity is sufficient for downstream tasks.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to use joint protein sequence and structure diffusion as a pretraining task, which they call DiffPreT. In order to account for the fact that proteins can exist as ensembles of conformers, the paper further proposes to generate pairs of conformers, use the diffusion forward process to corrupt them, and then train the reverse process across diffusion trajectories. They call this pretrained model SiamDiff. The paper then shows that DiffPreT and SiamDiff are effective on a panel of protein function prediction tasks.\n\nSTRENGTHS:\nThe paper addresses the important problem of protein function prediction using novel and effective methods. Using diffusion as a pretraining task is novel, interesting, and apparently quite effective. The SiamDiff objective is creative and also empirically effective. The experiments are well-done and convincing. The writing is clear, and the paper is easy to follow. I appreciate the uncertainties on metrics and the attempt to benchmark across a wide variety of tasks. Using two-stage noise-scheduling is also an intuitive way to adapt a model framework originally meant for generation to pretraining.\n\nWEAKNESSES:\n### Major\n\nIn general, this is a very good paper that clearly describes a promising method for protein pretraining. However, it only evaluates on Atom3D and EC number. The paper would be much more significant if they evaluated GearNet on zero-shot fitness tasks such as those in ProteinGym, protein engineering tasks such as those in FLIP, and more general structure prediction tasks such as those used in CAFA. \n\nIn addition, the results would be more general if the paper considered different diffusion schemes. For structure generation, diffusion on [frames](https://arxiv.org/abs/2301.12485) or [angles](https://arxiv.org/abs/2209.15611) seems to outperform diffusion directly on atoms. In sequence diffusion, D3PM with an absorbing state isn't as computationally efficient as [autoregressive diffusion](https://arxiv.org/abs/2110.02037) while not allowing iterative refinement like D3PM with a uniform prior over amino acids. The paper would be slightly stronger if it discussed these choices, and much stronger if it evaluated the effect of these choices. However, I realize that this is probably not feasible during the discussion period!\n\nIn general, I think the exposition of the diffusion model could be shortened, as that is not the main contribution of the paper, in order to make more space for contributions that would increase the paper's significance. \n\n### Minor\n\n- The paper should include ablations on EC\n- It's not clear to me why ESM-2-650M-GearNet doesn't require structure as input.\n\nQUESTIONS:\n- How is EC a residue-level task? Don't the labels apply to the entire protein?\n- Is there a way to use this model on proteins without a structure?\n- How does the choice of diffusion process effect performance?\n- How well does this work on function prediction, zero-shot fitness prediction, or protein engineering tasks?\n\nLIMITATIONS:\n- SiamDiff requires structures at inference time. \n- The paper only considers one diffusion process out of many possible choices.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}