Pre-Training Protein Encoder via Siamese Sequence-Structure Diffusion Trajectory Prediction
Authors: Zuobai Zhang
Minghao Xu
Aurelie Lozano
Vijil Chenthamarakshan
Payel Das
Jian Tang
Venue: NeurIPS 2023 spotlight (2023)
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
Zuobai Zhang 12*, Minghao Xu 12+ Aurélie Lozano Vijil Chenthamarakshan', Payel Das Jian Tang" |4.5
equal contribution ccorresponding author Mila Québec Al Institute, Université de Montréal, 'IBM Research, HEC Montréal, SCIFAR AI Chair
{zuobai zhang, minghao xu>@mila @mila..la.......m.mmmmm...al...a quebec, fecvijil, secvijil,aclozang aclozano, daspa Huss bmm com, j jian. tang@hec.co ca
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 diffu- sion models generative tasks, we propose the DiffPreT approach pre-train a joint diffusion modeling. DiffPreTT guides
protein encoder by sequence- the encoder recover the native protein sequences and structures from the per- turbed ones along the joint diffusion trajectory, which acquires the joint distribution 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 a protein. SifmDiff attains this goal by maximizing the mutual information between representations of diffusion trajectories of wwwetturally-oorelattte conformers. We study the effectiveness of DifiPreT and SiamDiff both atom- and residue-level structure-bare based protein understanding tasks. Experimental results show that the performance of DiffPreT is consistently competitive all tasks, and SifmDiff achieves new state-of-the-an performance, considering the mean ranks on all tasks. Code will be released upon acceptance.
Introduction
learning saaringg-basdd methods have made remarkable strides in predicting protein struc- 50) and functionality (55] 24]. Among them, self-supervised (unsupervised) pre-training have been successful in learning effective protein representations from available protein sequences from their experimental/prediceed structures. These pre-training ap- proaches 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 acquire co-evolutionary information 220 61). or protein structures to capture essential characteristics for tasks such as function prediction and fold classification [87 (3). 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, 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.
Machine
tures 44 [5]
approaches 20 61
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 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 novel approach that adapts denoising diffusion models to pre-train structure-informed protein encoders 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.
diversity 22 54 40
74..
In spite of these advantages, DiffPreTT ignores the fact that any protein structure exists 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 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 structurally- correlated conformers (i.e., siamese diffusion trajectories). We first adopt 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 wide range of downstream protein understanding tasks. These tasks encompass protein function annotation, protein 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, DiffPreTT consistently delivers competitive performance across all tasks and both the atomic and residue- level resolutions. Moreover, SifmDiff further enhances model performance, surpassing previous state-of-the-arr results in terms of mean ranks across all evaluated tasks.
2 DiffPreT: Diffusion Models for Pre- Pre-Training
Recently, there have been promising progress on applying denoising diffusion models for protein structure-sequence co-design [54]] 2 The efffcctiveness of the joint diffusion model modeling the distribution of proteins suggests that the process may reflect physical and chemical principles underlying protein formation B 18] which could be beneficial for learning informative representa- tions. Based on this intuition, in this section, we explore the application of joint diffusion models on pre-training protein encoders in pre-training and fine-tuning framework.
2.1 Preliminary
Notation. A protein with n, residues (amino acids) and nh atoms can be represented as sequence- structure tuple P = (S, R). We use S [$,, $2 ... sn,]] to denote its sequence with 8; as the type of the i-th residue, while R [r1, r2..., To,, R x3 denotes its structure with r. 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 certain threshold.
Equivariance. Equivariance is ubiquitous in machine learning for modeling the symmetry in physical and is shown to be critical for successful design and better generalization of 3D
systems 67 75
It is important to note that protein structure encoders in this context take both sequences and structures as input, distinguishing them from protein sequence encoders established in previous works.

---

Pre- Pre-Training
Fine-Tuning
..LLCGGEGGE....
GALCKRP....
@CP
(b R)
Pat so
(b(3)))
FF..
..YYGGEGGG.......
Atom3D
MSP
RES
PSR
A pre-training and fine-tuning framework for DiffPreT. During pre-training, diffusion processes are applied both protein structures and sequences, with * indicating masked residues. Noise prediction networks, parameterized with an encoder o(S, R), are employed to restore the original states. The learned encoder e(S, R.) is subsequently fine-tuned on downstream tasks.
Figure 1:
(x = Py o F(z), where px and py are transformations corresponding to an element g G acting on the space X and Y, respectively. The function is invariant w.zt G' if the transformations Py is identity. In this paper, we consider SB(3) group, i.e., rotations and translations in 3D space.
networks 46 Formally, function F X + ) is equivariant w.E.. group G J
Problem Definition. Given set of unlabeled proteins = PP,.P........ our goal is to train a protein encoder (SS R) to extract informative d-dimensional residue representations Ra,xd and atom representations R that are SE(3)-invarias w.r.f. protein structures R.
2.2 Diffusion Models Proteins
Diffusion models are a class deep generative models with latent variables encoded by a forward diffussio process and decoded by a reverse generative process (63). We use po to denote the ground- truth protein and Pe for = 1.... T to be the latent variables over T diffusion steps. Modeling the protein an evolving thermodynamic system, the forward process gradually injects small noise to the data po until reaching 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:
q(Pl:T II" q(P (PP'(P (PO:TT-11 II (Pt-11P)))
where qPP'ppt-1) defines the forward process at step and pe (pt with learnable parameters 0 defines the reverse process at step t. We decompose the forward process into diffusion on protein structures and sequences, respectively. The decomposition represented as follows:
q qRR R ('" pa (p
IP*) (R IP') ( pe(S"-1p') (2)
where the reverse processes use representations a' and h' learned by the protein encoder Φo(S", Rt).
Forward diffusion process (PP')-) For diffusion on protein structures, we introduce random Gaussian noises to the 3D coordinates of the structure. For diffusion on sequences, we utilize 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:
(R!) = -N(': VI B,R'-1, BTT), ((S"(S"=)) = random_ (aadlom_mss((s" ,t),,
(3)
Here, B1,..., By and PI, PT random mask(S' P:) denotes the random masking operation, where each residue in s is masked with probability of Pt at time step
series of fixed variances and masking ratios, respectively.
Reverse process on structures po (R'-1(P'). The reverse process on structures is parameterized as Gaussian with learnable mean μr (Pt, ,) and user-defined variance o,, Given the availability of

---

R2 as an input, we reparameterize the mean PP ,t) following Ho et al.
[37
po Pe(R"~11p1) - NNRRR (o o)) po(P"
voa
~
R
रíー.. co (P(,,
(4)
where = =1 a; II: a, and the network es(-) learns to decorrupt the data and should be cunslation-invariann and rotation naauiin-equivariant w.r.f. the structure R'.
To define our noise prediction network, we utilize the atom representations a' (which guaranteed to be SE(3)-invariant w.E. R$ by the encoder) and atom coordinates r' (which is SE(3)-equivarian w.e.. R²). We build an equivariant output based on normalized directional vectors between adjacent atom pairs. Each edge (i, is encoded by its length ||r rill2 and the representations of end nodes at al. and the encoded score ma, will be used for aggregating directional vectors. Specifically,
(5) where (i) denotes the neighbors of the atom in the corresponding graph of Pt. Note that co(P", achieves the equivariance requirement, as ma.j is SE(3)-invariant w.r.t. R while r rj is translation invarrinn and cotation-equivaren w.r.f. R'.
[ea(P' (aP,,t)) t)] = =EjEN( maj
[---))11 with ma.j = MLP(a), a). MLP([Ir - rrjll)),
Reverse process sequences po pp For the reverse process (P we adopt the parameterization proposed in [4. The diffusion trajectory is characterized by the probability (SS ))), and we employ a network pi to predict the probability of 5o:
Po(S x Es
11st 55) (S))))
(6)
We define the predictor pe with residue representations ht, For each masked residue i in St, we feed its representation h to an MLP and predict the type of the corresponding residue type sf in So:
Pa S = II, pe (tt) = II, Softmax(3 (MMmmmxx)(MMM(())))) MLP(h))),
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 337
C= E
T Ei=1 DEL (g(Pt-1 (p' (P)
(8)
Under the assumptions in 2 it can be shown that the objective can be decomposed into structure loss CRR and a sequence loss C(S) (see proof in App. C.2
((R) =E D t-1 DKL (e(R'~1(R' ( R (Ro
((S) = D DKL (5.5
Both loss functions can be simplified as follows.
Structure loss ((R), It has been shown Ho et 37 that the loss function can be simplified under our parameterization by calculating KL divergence between Gaussians weighted L2 distances between means eg and (see details App. C.3):
((R)
Σr B.E.witon) [ol'-co(P(.)))
(10)
where the coefficients ? are determined by the variances $1, Br. In practice, we follow Ho et al. 37 to set all weights "t = 1 for the simplified loss (R)
simple.
Since eo is designed to be rotation- retttiin--quvariiant w.E.t. R', to make the loss function invariant w.r.t. Rt. the supervision is also supposed to achieve such equivariance. Therefore, adopt the chain-rule approach proposed in Xu et al. [22 which decomposes the noise on pairwise distances to obtain the modified noise vector è supervision. We refer readers Xu al.
82 for more details.
Sequence loss ( Since we parameterize Po(S /P') with po 5555pp")) and (St as in (6). it can be proven that the t-th KL divergence term in (S) reaches zero when Pe (p)) assigns

---

all mass on the ground truth S (see proof in App. C.4 Therefore, for pre-training, we can simplify the KL divergence the cross-entropy between the correct residue type of and the prediction:
((S)
Ei= E, CE C5,, (o)))(()))P)))
(1))
simple
where CE(.,) denotes the cross-entropy loss.
The ultimate training objective is the sum of simplified structure and sequence diffusion losses:
(R) simple
(S)
simple
simple
(12)
2.4 Two-Stage Noise Scheduling
Previous studies on scheduled denoising autoen-
Stuutur Diffuzion Las
stge arge
small
coders images [3 have shown that large noise levels encourage the learning of coarse-grained fea-
JIRt dissin the saaee
Sequence dfftaten iave-dage)
tures, while small noise levels require the model
24
JanLd Rualon slagend
juird flusion roral nowe
learn fine-grained features. We observe similar phe-
nomenon structure diffusion, as depicted in Fig. 2]
The diffusion loss with large noise at a fixed scale
Scquenoe Acurrsy
(orange) 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)
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- pare the loss of sequencc-only diffusion pre-training in Fig. 2 Sequence diffusion achieves higher denoising accuracy than joint diffusion when using uncorrupted structures, supporting hypothesis.
Figure 2: Structure diffusion loss and se- quence denoising accuracy of pre-training with different noise scheduling strategies. (bluc) with joint diffusion pre-training (red)
Unlike recent denoising pre-training methods that rely on a fixed noise scale as a hyperparameter 85], DifiPreT incorporates denoising objective at various noise levels to capture both coarse- and fine- grained features and consider the joint diffusion effect explained earlier. Both granularities features crucial for downstream tasks, they capture both small modifications for assessing protein structure quality and large changes leading to structural instability. In our implementation, we perform diffusion pre- pre-training using T = 100 noise levels (time steps). Following the intuition of learning coarse-grained features before fine -re-grained ones in curriculum learning [8] we divide the pre-training process into stages: the first stage focuses
large noise levels (t = 10,..., 100), while the second stage targets small noise levels (t = 1, 1..., 9). As observed in Fig. 2 structure diffusion becomes more challenging and sequence diffusion becomes easier during the second stage, as 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.
SiamDiff: Siamese Diffusion Trajectory Prediction
Through diffusion models on both protein sequences and structures, the pre-training approach pro- posed in Sec. frries 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 con- straints 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.
5

---

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. B. : generate joint diffusion trajectories (See. (.2). and transform MI maximization into mutual denoising between trajectories (Sec. 3.3 sharing similar loss with DifiPreT.
3.1 Conformer Simulation Scheme
GENCGENNGS......
..
Our method begins by generating differ- ent conformers of a given protein. How- ever, direct sampling requires 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 perturba- tions to the side chains (36..
liselity
(5.38
Pentetalve
..YNNGGEGGP...
Specifically, given the original protein (8, R), we consider as the na- tive state Pi = (S1, R₂) and generate correlated conformer P3= (S2, R2) by randomly perturbing the protein structure.
Figure 3: High-level illustration of SiamDiff. Mutual denoising of diffusion trajectories is performed across two correlated conformers P1 and P2.
That is, we set S₂ to be the same as S₁, and R₂ is obtained by applying a perturbation function perturb(R,, e) to the original residue structure, where E [0, 2m="* is noise vector drawn from wrapped normal distribution [15]. The perturb(...) function rotates the side-chain of each residue according to the sampled torsional noises. To avoid atom clashes, 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 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 incorporating available rotamer libraries 622 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 pu:7 and Po:T for conformers 'P1 and P2, respectively, using the joint diffusion process outlined in Sec. For example, starting from P P1 (S. R1). we use the joint diffusion on structures and sequences
to define the diffusion process q :T (P!) (RR 1:T (R)) (T))) We derive trajectories on structures R.P.T using the Gaussian noise 3 and derive the sequence diffusion process $1:7' using
the random masking in (6). In this way, we define the trajectory pu:T ((S), R() for P, and
can derive the siamese trajectory PD. (S)) R()
T similarly.
3.3 Mutual Information Maximization between Representations of Siamese Trajectories
We aim to maximize the mutual information (MI) between representations of siamese trajectories constructed way that reflects the correlation between different conformers of the same protein. Direct optimization of MI is intractable, so we instead maximize lower bound. In App. C.I we show that this problem can be transformed into the minimization of loss function for mutual denoising between two trajectories: L (C(2-11) + C(1-2))), where
(b-ra)) = Expar .pa.T
ΣTT DKI. (q(Pi-1(Pt Pe Pa)l (P)))
)}
(13)
with b + a being either 2 + 1 or 1 2 and PoT 0:T being representations of the trajectory PO:T.

---

The two terms share the similar formula as the ELBO loss in 8 Take [(2-1) for example. Here (PP [pt_ Pi)
is a posterior analytically tractable with our definition of each diffusion step q(Pi)P in ( and (6) The reverse process is learnt to generate a less noisy state Pi-1 given the current state Pf and representations of the siamese trajectory P22 which are extracted by the protein encoder to be pre-trained. The parameterization the reverse process is similar as in Sec. 2.2
with the representations replaced by those of P.-T (see App. B for details).
Our approach involves mutual prediction between two siamese trajectories, which is similar to the idea of mutual representation reconstruction in 14]. However, since Pi and P2 share information about the same protein, the whole trajectory of P2 could provide too many clues for denoising towards PI- making the pre-training task trivial. To address this issue, we parameterize p(P)-1IP(: (P:.PP with po (P P2). For diffusion on sequences, we further guarantee that the same set of residues masked in Sf and S 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 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
Camplle
1 (C((ppte))) (Siiple
( ((122-11) + simple
simple
+ -.
simple
(14)
where clingre is the loss term defined by predicting po:T from PO-T (see App.
B for derivation).
3.4 Residue-level model
Residue-level protein graphs are simplified atom graphs that enable efficient message passing between nodes and edges. As in Zhang et al. [87), we only keep the alpha carbon atom of each residue and add sequential, radius and K-nearest neighbor edges as different types of edges. For Sifm 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 P2 (S2, R2) is defined
by S2 = Si, R2 Ri + e, where e E R x3 is the noise drawn from a Gaussian distribution.
4 Related Work
Pre-training Methods 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))
61 50 Typical sequence pre-training methods include 50 and contrastive learning (53] The pre-trained PLMs have achieved impressive performance on variety of downstream tasks for structure and function pre- diction [59 81). Recent works have also studied pre-training on unlabeled protein structures for generalizable representations, covering contrastive learning [87 (3), self- celf-prediction of geometric
20
53
masked protein modeling 20 61
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 has been recently rekindled for its strong performance on image and waveform generation 37
et al. 65 and
While DPMs 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 77 and
images [38]. Inspired by these progresses, DPMs have been adopted 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 [5] protein design 2 54 40)
74] and motif-scaffoldinn
[69] In alignment with recent research efforts focused on diffusion-basee image representation learning [I] this work presents novel investigation into how DPMs can contribute to protein representation learning.
Now 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 [[2] 84] in this work, we
7

---

Table I: Atom-level results Atom3D tasks.
PIP
MSP
PSR
Mcan
Method
Rank
AUROC
AUROC
Accuracy
Global
Mean
GearNet-Edge
0.868=00002
0.633-00007
0.441=0.001
0.782400021
0.488 =00012
7.6
Denoising Score Matching
0.877=00002
0.629440000 040
0.448=0.00
0.81340.003
0.518+0.000
5.2
w/pre-traning
Residue Type Prediction
0.879=0.004
0.620=0.027
0,449=0.001
0.8261:0.020
0.518:00018
4.4
Distance Prediction
0.872=0.001
0.677+40000 020
0.422=0.00
0.840:=0.020
0.52240
4,0
Angle Prediction
0.878=0.001
0.642
0,419=0.001
0.813:40.007
0.503+20.012
6.2
Dibedral Prediction
0.878:10.004
0.591 1000088
0.41410.001
0.821
0.497.10.004
6.8
Multiview Contrast
0.871+00.003
0.646-+0.000
-3688-0.001
0.805 0-805--0.005
0.502 0.002--0.009
7.2
DitPPreT
0.880=00005
0.680:200018
0.452=0.000
0.821=0.007
0.533
2.4
SiamDift'
0.884=0.003
0.69810.020
0.460:1.000
0.82910.012
0.546:100018
Table 2: Residue-level results on EC and Atom3D tasks.
EC
MSP
PSR
Mean
Method
Rank
AUPR
Faus
AUROC
Global
Mcan
GearNet Edge
0.837=#0.002
00.01
0.644-+0.023
0.763++0.12
0.373--0.021
7.8
Denoising Score Matching
0.859=0.003
0,840+0.00
0.645:10.023
0.795+0.027
0.429::00017
5.0
w/ pre-training
Residue Type Prediction
0.851=0.002
0.826:10.00
0.636:10.003
0.828=0.005
0.4801100031
5,4
Distance Prediction
0.858=00.033
0.836:10.001
0.623:1:0.007
0.796110017
0.41611.0021
6.4
Angle Prediction
0.873=0.003
0,849:0 001
0.631 0.1
0.802=0.015
0.44610.009
4.2
Dihedral Prediction
0.858=000001
0.840:10.001
0,568:10.022
0.732:10.02
0.398 1.0022
7.2
Multiview Contrast
0.875=00.003
0.713 0.036
0.752++0001
0.388
DitPPreT SiamDiff
0.864=0.002
0.844440.001
0.673::0.04
0.815+0.0006
0.505=00.077
3.2
0.878
0.85710.003
0.700:10.043
0.856:10.007
0.521
1.2
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 aequence-structure tuple P (S,R) R) may not be 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 [7 49] and have been applied to unsupervised representation learning [I]. 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 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 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 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.
Experiments
5.1 Experimental Setups
Following Zhang et al. 87 we pre-train our models with the AlphaFold protein structure database v1 44 77), including 365K proteome-wide predicted structures.
Pre-training datasets.
Downstream benchmark tasks. 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ć et [24 with 95% sequence identity cutoff. The ATOMBD 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-Fdge as the underlying model for both atom- and residue-level structures. GearNet-Fdge incorporates various types of edges and edge-type-specing convolutions, along with message passing between edges, to model protcin structures effectively, We compare our proposed methods with several previous protein structural pre-training algorithms, including multiview contrastive learning [87]. denoising score matching (2 and four self- self-prediction methods (residue type, distance, angle, and dihedral prediction) 87. For residue- 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 the AlphaFold protein structure database following Zhang et al. [87) and fine-tune it for 50 epochs on EC, MSP, and PSR. However, due 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 Fmx and AUPR for EC, AUROC for PIP and MSP, Spearman's p for PSR, and micro-averaged accuracy for RES. More details about experimental setup can be found in App. D
5.2 Experimental Results
Tables D 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- vee-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 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 both atom and residue levels.
In contrast, our DiffPreT approach achieves top-3 performance in nearly all considered tasks, show- casing its versatility and effectiveness across different evaluation criteria. Moreover, SifmDiff surpasses all other pre-training methods, achieving the best results in 6 out of 7 tasks, establishing it as the state-of-the 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 com- ponents of SifmDiff, we perform ab-
Table 3: Ablation study on atom- Atom3D tasks.
Method
PIP
MSP
RES
PSR
lation study on atom-level tasks and
AUROC
AUROC
Accuracy
Global
present results in Table 3 We first examine two degenerate settings of joint diffusion, i.e., "w/o sequence dif- fusion" and "w/o structure diffusion" " These settings lead to a deterioration in performance across all benchmark
GeuuNet-Edge
0.868:20.002
0.633:000077
0.441:20001
0.782:00021
SinmDiff
0.884:0.003
0.698
0.460
10.00
who seg, slilf
0.873=00.04
0.62
0.443:0 001
0.803-200010
strict, diff.
0.878.10.003
0.65210.021
0.45620.00
0.805
wio MI max.
0.880:0.000 005
0.68010.0188
0.452:0.000
001
0.001
w/ small noise
0.875.10.002
0.646:100031
0.444:0.001
0.828
iarge
0.867+00.003
0.683++0.020
0.443+00001
0.819+0.011
tasks, highlighting the importance of both sequence diffusion for residue type identification in RES and structure diffusion for capturing the structural stability of mutation effects in MSP. Next, we compare SifmDiff 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 perfor- mance in PSR by capturing iee-iia 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, evidenced in the table.
5.4 Combine with Protein Language Models
Protein language models (PLMs) have recently become standard method for extracting representations from
Table 4: ESM2-650M Gear-ee residue-level tasks.
Method
Fem
Global
Mcan y
protein sequences, such as ESM [50] However, these methods are unable to directly handle structure-related tasks
GearNat-Fllge
SiamDii
0.85
0.708=1U043
0185-00007
0521
ESM-CrarNee
8.50000000
0.329-0:113
SiamDDir
0.00
0.602=00100
0.863-0.099
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 show- case the potential of SiamDiff on PLM-based encoders, we pre- pre-traine trairrined the ESM-GearNet encoder using SiamDiff and evaluated its performance residue-level tasks. Considering the model capacity and computational budget, we selected ESM--26500M as the base PLM. The results in Table demon- strate the performance improvements obtained by introducing the PLM component in ESM-GearNet. Furthermore, after pre-training with SiamDiff, ESM-GearNe achieves even better performance on all tasks, especially PSR where ESM-only representations 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. multimodal baselines in App. different diffusion strategies App. G and different backbone models in App.
Conclusions
In this work, propose the DiffPreT approach to pre-train 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 DifiPreTT by additionally modeling the correlation between different conformers of one protein. Extensive experiments on diverse types of tasks and both atom- and residue-level structures verify the competitive performance of DiffPre'T 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-3583119722
10

---

References
[1] Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Diffusion-basee representation learning. arXiv preprint 1.Kiv:2205..422. 2021.
[2] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019, 2022.
[3] Christian B Anfinsen. The formation and stabilization of protein structure. Biochemical Journal, 128(4):737, 1972.
[4] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 14:17981-17993 2021.
[5] Minkyung Baek, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang. Qian Cong, Lisa N Kinch, R Dustin Schaeffer, al. Accurate prediction of protein structures and interactions using a three-track neural network. Science, 373(655)):871-876 2021.
[6] Inigo Barrio-Hernande 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 1023-03. 2023.
[7] 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.
[8] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In International Conference on Machine Learning, 2009.
[9] B Chandra and Rajesh Kumar Sharma. Adaptive noise schedule for denoising autoencoder. In International conference on neural information processing, pages 535-542. Springer, 2014.
[10] Can Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self- supervised learning. arXiv preprint wrXiv:2204.04213 2022.
[11] 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.
[12] 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 httpa//////eeeeee.eee...e. id-YDJRFWEMNby
[13] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607.. PMIR, 2020.
[14] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Patlern Recognition, pages 15750-15758, 2021.
[15] 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.
[16] Christian Dallago, Jody Mou, Kadina Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks in fitness landscape inference for proteins. In Thirty-fithh Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round2 2), 2021.

---

[17] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances Neural Information Processing Systems, 14:8780-8794, 2021.
[18] Ken A Dill and Justin L MacCallum. The pootein-folding problem, 50 years on. science, 338 (6110):1042-1046 2012.
[19] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications Pure and Applied Mathematics, 36 (2):183-212, 1983.
[20] 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. 1190/TPAA..2222.3355555
[21] 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.
[22] 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):84-192, 2020.
[23] Krzysztof J Geras and Charles Sutton. Scheduled denoising autoencoders. arXiv preprint 1406.33400.2200 2014.
[24] Vladimir Gligorijević, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, et al. cruucture-bssd protein function prediction using graph convolutional networks. Nature communications, 12(1):1-14, 2021.
[25] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
[26] 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.
[27] 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.
[28] lean-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 2020.
[29] Yuzhi Guo, Jiaxiang Wu, Hehuan Ma, and Junzhou Huang. Self-supervised pre-training for protein embeddings using tertiary structures. AAAI, 2022.
[30] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirleenth international con- ference artificial intelligence and statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[31] Michael U Gutmann and Aapo Hyvärinen. Noise-contrastive estimation of unnormalized statistical models, with applications natural image statistics. Journal of machine learning research, 13(2), 2012.
[32] Kaveh Hassani and Amir Hosein Khasahmadi. Contrastive multi-view representation learning on graphs. In International Conference on Machine Learning, pages 4116-4126. PMLR, 2020.
12

---

[33] Pedro Hermosilla and Timo Ropinski. Contrastive representation learning for 3d protein structures. arXiv preprint arXiv:2205.156755 2022.
[34] Pedro Hermosilla, Marco Schäfer, Matëj 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.
[35] 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.
[36] Bosco K Ho and David A Agard. Probing the flexibility of large conformational changes in protein structures through local perturbations. PLoS computational biology, 100)ee00 2009.
[37] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 6886840-6851, 2020.
[38] 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: 34:12454-12465, 2021.
[39] Emiel Hoogeboom, Victor 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.
[40] 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 programmable generative model. bioRxiv, 2022.
[41] Bowen Jing, Stephan Eismann, Pratham N. Soni, and Ron O. Dror. Learning from protein structure with geometric vector perceptrons. International Conference on Learning Repre- sentations, 2021. URL URLpsttts:://ppeereviiee
[42] Bowen Jing, Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106.03844 038433 2021.
[43] 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: htttps://openreview wwt.:ntoorrmdddweeeeeeem j2r62r
[44] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-599 2021.
[45] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 131266 14, 2013.
[46] 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.
[47] 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, ((12)):011-1020) 2019.
[48] Yann LeCun and Ishan Misra. matter of intelligence, 2021. self suue-vsupev -earnddd-ernnnnnn the- ofatter-of-ineellieeee
Self-supervised learning:
The dark
URL https://ai. facebook.co. ctt.blog.///////////oo//////oooc
[49] Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- Im improves controllable text generation. arXiv preprint 2rXiv:2205.14211 2022.
13

---

[50] 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. bioRviv, 2022.
[51] 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.
[52] 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 btttps://operreeee wwweenrvviee.eet/frrrreeeee htt.eee..ee..ootu.raddddooooooo id=CjTHVoldvR
[53] Amy X Lu, Haoran Zhang, Marzych Ghassemi, and Alan M Moses. Self-supervised contrastive learning of protein representations by mutual information maximization. BioRxiv, 2020.
[54] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng. and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion-based generative models for protein structures. 2022.
http:/rrreieeeeeeeerrr/aaaaaaaee net/forumPicc:Sor6n22]
[55] Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Lan- guage models enable zero-shot prediction of the effects of mutations on protein function. bioRxiv, 2021. doi: 10. 1a.11002222..00 w1t..09.4555 4556488. 1101/2021 07 09. 450648v1
URL htttps://www..... wworxi..wwi...oooo..o...wwwwwww. org/ content/10
[56] G UIrich 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.
[57] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint trXiv:1807.03744 2018
[58] Clare M O'Connor, Jill U Adams, and Jennifer Fairman. Essentials of cell biology. Cambridge, MA: NPG Education, 1:54, 2010.
[59] 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.
[60] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[61] 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. 2021.
[62] Maxim V Shapovalov and Roland I. Dunbrack. A smoothed backbone-dependent rotamer library for proteins derived from adaptive kernel density estimates and regressions. Structure, 1944)844-8588 2011.
[63] Jascha vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265, 2015.
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper-
[64] Vignesh Ram Somnath, Charlotte Bunne, and Andreas Krause. Multi-scale representation learning on proteins. Advances Neural Information Processing Systems, 34, 2021.
[65] 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 2rXiv:2011.13456 2020.
[66] Freyr Sverrisson, Jean Feydy, Bruno E Correia, and Michael M Bronstein. Fast end-to-end learning on protein surfaces. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pages 15272-15281, 2021.
14

---

[67] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-ant translation-equiie -qquivarinn neural networks for 3d point clouds. arXiv preprint arXiv: 18Xiv:8802.000 002188022228121 19, 2018.
[68] 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 preprini arXiv:2012.04000 2020.
[69] Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S. Jaakkola. Diffusion probabilistic modeling of protein backbones 3d for the motif- scaffolding problem. In The Eleventh International Conference on Learning Representations, 2023.
net www.eeeww.e/eerree wtt/:/oornnndddddd forrf
[70] 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-sequent space with high-accuracy models. Nucleic acids research, 2021.
[71] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprini arXiv:2209. 14734, 2022.
[72] Junmei Wang, Romain M Wolf, James W Caldwell, Peter A Kollman, and David Case. Development and testing of general amber force field. Journal of computational chemistry, 25 157:1157-1174, 2004.
[73] Limei Wang. Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning protein representa- tions via complete 3d graph networks, arXiv preprint arXiv:2207.12600 2022.
[74] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian Trippe, Jason Yim, Helen 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.
[75] 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.
[76] 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 2244:022004.0663, 2022.
[77] Kevin E Wu, Kevin K Yang, Rianne van den Berg, James Zou, Alex X Lu, and Ava Amini. Protein structure generation via folding diffusion. urXiv preprint arXiv:2209, 156/1, 2022.
[78] Lemeng Wu, Chengyue Gong, Xingehao Liu, Mao Ye, and Qiang Liu. Diffusson- molecule generation with informative prior bridges. arXiv preprint arXiv:2209.0m 220.00050000 2022.
[79] Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graph- level representation learning with local and global structure. In International Conference on Machine Learning, pages 1 11548-1 11558. PMLR, 2021.
[80] 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.
[81] 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.02090 220622299. 2022.
[82] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Con- ference on Learning Representations, 2022. URI. https:///pperreiie net/ formmmddd
P2cvxEMZvQC
15

---

[83] Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. bioRxiv, 2022.
[84] Yuning You and Yang Shen. Cross-modality and self-supervised protein embedding for compound-protein affinity and contact prediction. Bioinformatics, 1Ssuupplement22::i868---1
2022.
[85] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzale 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
[86] Zuobai Zhang, Minghao Xu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, and Jian Tang. Enhancing protein language model with structure- encoder and pre-training. In ICLR 2023-Machine Learning for Drug Discovery workshop.
[87] Zuobai Zhang. Minghao Xu, Arian Jamash, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining, arXiv preprint arXiv:2203.06112, 2022.
[88] Zhaocheng Zhu, Chence Shi. Zuobai Zhang, Shengehao Liu, Minghao Xu, Xinyu Yuan, Yang- tian Zhang, Junkun Chen, Huiyu Cai, Jiarui Lu, al. Torchdrug: A powerful and flexible machine learning platform for drug discovery. arXiv preprint arXiv:2202.08320, 2022.
16

---

A More Related Works
Protein Structure Encoder. The community witnessed a surge of research interests learning informative protein structure representations using structure-basee based encoders. The encoders are designed to capture protein structural information on different granularity, including residue-level structures 124 87 80], atom-level structures [34] 4173 and protein surfaces [22] 66, 64] In this
work, we focus on pre-training a typical residue-level structure encoder, i.e., GearNet-Edge typical atom-level structure encoder, i.e., GVP
and
41)
Mutual Information (MI) Estimation and Maximization. MI can measure both the linear and non- linear dependency between random variables. Some previous works 35] try to use neural networks Vonsker-Varadhan
estimate the lower bound of MI, including Donsker- Varadhan representation [9] Jensen-Shanne
divergence [1 and
Estimation (NCE) (30 [1. The optimization with InfoNCE loss [57] maximizes a lower bound of MI and broadly shown to be a superior representation learning strategy [3 32 79 51 87]. In this work, we adopt the MI lower bound proposed by Liu 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 to enhance protein representations by utilizing joint pre-training using 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. 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, discuss details of our SiamDiff method. We first describe the parameterization of the generation process Po (PPP Pb) in Sec. B.1 derive the pre-training objective in Sec. B.2 and discuss some modifications when applied on residue-level models Sec. 3.4
B.1 Parameterization of Generation Process
Remember that we use PO:T and P to denote the representation of the siamese trajectories 'DO-T' and puu respectively. Different from the generation process in traditional diffusion models, the parameterization of (P Pb) should inject information from P!. Therefore, we use extracted residue and atom representations (denoted as al and hτ) of P for this denoising step. Given the conditional independence in (2), this generation process can be decomposed into that on protein structures and sequences similarly Sec. 2.3
17

---

Algorithm SiamDiff Pre- Tre-Training
Input: training dataset D, learning rate a
Output: trained encoder db
1: for (S, R2) in D do
2: (Si, R1) < (S, R);
3:
(S2, R2) (S, perturbbRR c)):
4:
sample noise scale E (1...));
5:
Ri ~ Nivō, (1 ā. 1)
6:
R2 ~ ~N
7:
Si, raandm_mmak((5)
8:
h1 di(
9:
ln &- Ri))
10:
((R) 1(R1, [2R,, hi):
11:
( [(Si, h2) LSS, h1);
12:
- 8 aVo(L + 3
13: end for
Generation process on protein structures. As in ( modeling the generation process of protein structures is to model the noise on R2 and gradually decorrupt the noisy structure. This can be parameterized with a noise prediction network eg (Pi P I) that translation-invariant and rotation- equivariant w.r.t Ri. Besides, the noise applied on Ri should not change with transformations on R2,, so eg should be SE(3)-invariant w.r.f. R2--
To achieve these goals, we build our noise prediction network with atom representations až (which is SE(3)-invariant w.r.t R2) and atom coordinates r (which is SE3))eqqivvarian w.r.t. R2). We define an equivariant output similarly as in DiftPreT. Specifically, we have
(P) P), EjENÇ() Maj"
| with m,j MLP(a abj. MIP([Iri, rfyllz)),
where (i) denotes the neighbors of the atom in the corresponding graph of Pt. Note that (P P achieves the equivariance requirement, m,, is SEE()-invariant w.r.t R4 and Ri while rf. rf, is translation-invarian and rotation-equivariant w.r.f. Ri-
Generation process protein sequences. As in (7). the generation process sequences aims to predict masked residue types in S with predictor Pg. our setting of mutual prediction, we define the predictor based on representations of the same residues in St. which are also masked. Hence, for each masked residue i in S2, we feed its representation hig to an MLP and predict the type of the corresponding residue type s), in SP:
Po (o Pb) = II, (o (P. Pb)) = II,
where the softmax function applied over all residue types.
B.2 Pre- 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 ((2-1)) for example and its counterpart can be derived in the same way. The objective can be decomposed into structure loss ((R.2-1) and sequence loss (1)2-+1);
(RR,2+11) =E DI DKL ((8 =E I DEL (1(St-'1S,,SS (1o (P|, Pb))]
RR) Pb))]
(15)
t-1
(16)
Based on the derivation in Sec. 2.3 the structure loss ( can be simplified as
(R,2-+1) simple
Σi= BenN(onn) [[le ea (PI, P/ 1)1)
(17)
and the sequence loss 2 can be simplified as
((5,2=+1) ΣI Σ, CE (e Pt, Pb)). simple
(18)
Then, the final objective in Sec. 3.3 can be easily derived.
18

---

C
Proofs
In this section, we provide proofs for propositions Sec. 2 and Sec. Due to the similarity between the two methods, all propositions restated for SiamDiff. DiffPree can be seen 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 PT and por denote the corresponding random variables of representations of the siamese trajectories P and P..
Proposition 1 With approximations, the mutuaal information between representations of two siamese trajectories is lower bounded by:
I(PTT ;P -((0011 + c(1+2))
+ +C,
where is a constant independent of our encoder and the term from trajectory P to Po: defined as
((b-ba)
-Epar .pg.T ET= DKL ( (q(P:-u1Pa PP)lp(( p!.
with + a being either 2 + or 2
Proof. First, the mutual information between representations of two trajectories defined as:
T(PO:T; :.:T) Epprr P (PpP T p.l [og
p(p4: pp.
(19)
where the joint distribution is defined as p( Ph:T" PO-T") = p(Pp: Po ) PIIT Pi) (22
Next, we can derive lower bound with this definition:
(PATT PDT
I(PPIT: P
E log
1:T ((PD)((P)
P
Pi)p(PP) 1:f
P!!)
p
PPOTT
O.T .O.TT
log
(((P)) 2) 1:T
PP))(PP 1:T $1P!)
PPO:T PO
log
2
P(P (PP))((P)) (P
!P)) (P) 1:T PD))2
p PPI:T P...T
p P P::T PD:T
log
+ llog P(P 1:7
((P2)q(P 1:7
PP (P)
1PPT (PD)
|P (1gg 1: (p!)
P(P
p(P) U:T (!! log
log
2
(P 1:7
P
(P) 1:7 PD
However, since the distribution of representations are intractable to sample for optimization, we in- stead sample the trajectories ppp:T and po: from our defined diffusion process, i.e., p(PP:T, pDoTT) p(PP, P.) (P(]) (PP):T Besides, instead of predicting representations, we use the repre- sentations 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:
log P(PP:T (PTT) (P!))P))
+ log p(P (Pl.T) (PI
log (PP-T1P9.T)) + log ((P)'"(p))
P(PS I(P.TTPPPT)) (PPT) q(p) 7 (PD)
Ep Pa-T
We now show the first term on the right hand side can be written as the loss defined in Proposition The derivation is very similar with the proof of Proposition 3 in Xu et al. (82). We include it here for
19

---

completeness:
(PP.T(PP.T))
E Epp pg-
log
(PPI.T(PP)
p(P IP,,
T T
log
(PPiPPT-1)
log (Pi11) P?
p(P) |P| Po:7)
q(P) (P9)
=Bpp7 .2g T
log
(PH(P)))
((Pi -1 IPI (P, PT)
(P((P))
pPP 'PI p!:T)
log (P) (Pt) + log P.iT) + Σ log
-BppT,pp-
((P ((P(-1|P4, PP))
Epp. 1p8-T
Σ t=
DKL (q(P) 1
PI. P!-T))
))
+ ((2-1)
((2++1) + C(2-11)
where we merge the term p P] P
into the sum of KL divergences for brevity and use (22-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] [
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
((P:)Pt-t)> = (RR ((S
(20)
and 2) the conditional independence of the generation process
PP Pt)) po - ( P() po P(),
(21)
it can be proved that
((b-sa) C(Rbb-aa) + ((S
(22)
where the three loss terms are defined as
(bb-a) B Dtm DxL (9(P!- Puy pppp:ppppppi (Pt-1 P())]
(RRbboa) =E [ DRL. (1(R!- RD. (R P()
I
((S,b-a) =E EET DKL (9(St- (e P())]
with + a referring to the term from trajectory Ph: to P.:T.
Proof. Let c,)) to denote the t-th KL divergence term in L Then, have
c(0-ta) DxI. (q(Pt=1 P9) (P: 1pt ppt P!:r)
DKL (Iq(R)-1 R1, R. ( - S2] po (R (R--LIPt P.:T }e (S - pb:T) ))
Ipl DKL. (q(Rt- RDD
(R - p. + DKL (q(S) 12) -ppt po:T)
(RR.b+o)) C (S,b+a)
where we use the assumptions 20 and 21 in the second equality. The third equality is due to the additive property of the KI. divergence for independent distributions, Adding T KI. divergence terms
together will lead to (22
20

---

C.3 Proof of Simplified Structure Loss
For completeness, we show how to derive the simplified structure loss. The proof is directly adapted from [2).
Proposition 3 Given the definition of the forward process
q(Rg)Raa = =N(RRa) /1
B.R ,11))
(23)
and the reverse process
pa (R (R ( = =N(R (o Pi 1), o)
(24)
Br (P', ai P/ 1).
no Pb.)) Vot
R₂
(25)
vI - ā
the structure loss function can be simplified to
(( Et=r DxL (9(R)-11R) (Po Pb))]
(26)
((Rbb+a)) = = (elerN((O.)) [[le- Eg Pi, t) t)]
(27)
Bc, ai = II 0g and b - is either 2 or - where 7t 20,(1-0, 1) with at = -3
Proof.
First, we prove g (
R2) = NR(R Vor (1 T). Let & be the standard Gaussian
random variable at time step i. Then, we have
R Va VBe
Var-lar + aa BBEe
=aaR + ii +.... Var-1Br-1t-11 + BeEe,
which suggests that the mean of R2 is va,Ra and the variance matrix is 10.00-1-..00. + +.. 1 + 3,) = a.
Next, we derive the posterior distribution as:
(R', = (R,, R (R RR.)
a(Rg, RD)
N(R Vat 13,1) N N 1 Vā (1 - d1-1)1)
N(R v vo, Ro (1 (22 04
1 Rd Varl Ō 1)
1- āi-1 3,T).
=N(R 1
Rt.
1
Let 3,
1-ū; Br, then the 1-th KI. divergence term can be written as:
I-a,
DKL (g(R) - RD) 11pe ( 1 P())
Vā Rd.
+ Vas(1 - -Z=1)
R¹
Br
co (Pi a: Pb,())
l
23
1- σt
Vat
VI - āi
Vo4-191 13
R. VI dic
Var x(-1)
( Rk.
Be
Eg
+
Rt.
co Pb:
23,
1 - āi
Vāz
1 ā
Vas
v1 -1-a;
88
Er ||c co (P t)
23,
an(1 ōit)
=7,
[||e- - ea (Pi, P/, 1)]22
which completes the proof.
0
21

---

C.4 Proof of Simplified Sequence Loss
Now we show the equivalence of optimizing sequence loss (( and the masked residue type prediction problem on Sg
Proposition 4 Given the definition of reverse process on protein sequences
Po P()) x Ese q(S) Saa St) pa(S PP:, Pi).
(28)
the sequence loss ((5.b-a) reaches zero when pa (S Pt)) puts all mass on the ground truth St.
Proof.
The loss function can written as: ((S
E
DKL (1(5) Sh S. S Pb))
Esg 9(5 1
S Sa) (5 Pt. P22)
DKL. 1 )
E
Z
where is the normalization constant. Hence, when Po (5 Pa Pt) puts all mass on the ground truth Sg. the distribution pa ( Pt)) will be identical with g(S) - -Is S2), 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 538 binary classification problems. We adopt the dataset splits from Gligorijević et al. (4 and use the test split with 95% sequence identity cutoff following Zhang 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 protein complex not (binary classification). The benchmark dataset is split upon 30% sequence identity cutoff among different splits.
4. Residue Identity (RES) task studies the structural role of an amino acid under its local envi- ronment. 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.
5. 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, remove all hydrogen atoms following Jing 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 long range interaction filter that removes edges with low sequential distances. We refer readers Zhang al. 87 for more details.
Atom-level backbone models. To adapt GearNet-Fdge atom-level structures with moderate computational cost, construct the atom graph by using only the spatial edge with the radius
22

---

dradius = 4.5A. 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 ReLUU 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 orig- inal paper of GearNet-EGde Edge [87]. We adopt the same configuration of relational graph construction, i.e., the sequential distance threshold dseq = 3. the radius dradius = 10.0A, the number of neighbors k 10 and the long range interaction cutoff' choog 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, adopt the original configurations. For Multiview Contrast, we 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 random edge masking function with mask rate equal to 0.15 is applied for constructing views. The temperature r 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 15 When masking residue 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 DifiPreT, for structure diffusion, we use a sigmoid schedule for variances 3r with the lowest variance 31 = 1e and the highest variance 3T = 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 SifmDiff, 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 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 F Fmax (used in EC), global Spearman's p (used in PSR) and mean Spearman's 's p (used in PSR) as below:
F Fmax denotes the protein-centric maximum F-score. It first computes the precision and recall for each protein at a decision threshold E [0, 1]:
precision, (1) E, 1l e P.(t) nT]
recall: E, 1lf E P(( E, (11E ET)
(29)
E, 1f P((O)
23

---

Table 5: Optimization configurations for pre- pre-training methods. Here max length denotes the maximum
number of residues kept in each protein and Ir stands for learning rate.
Max length
Batch size
Method
Optimizer
Ir
residue
atom
residue
atom
Residue Type Prediction
100
100
96
64
Adam
le-3
Distance Prediction
100
100
128
64
Adam
le-3
Angle Prediction
100
100
96
64
Adam
le-3
Dihedral Prediction
100
100
96
64
Adam
le-3
Multiview Contrast
96
64
Adam
le-3
Denoising Score Matching
200
200
12
Adam
le-4
Dif
150
100
16
64
Adam
le-4
SiamDiff
150
100
32
Adam
le-4
Table 6: Optimization configurations for downstream evaluations. Here max length denotes the
maximum number of residues kept in each protein and stands for learning rate.
GPUS
Batch size
Task
Optimizer
Ir
residue
atom
residue
atom
EC
4
N/A
2
N/A
Adam
lc-4
PIP
N/A
N/A
8
Adam
le-4
MSP
4
8
Adam
le-4
RES
N/A
4
N/A
64
Adam
le-4
PSR
4
8
Adam
le-4
where f denotes functional term in the ontology, T, is the set of experimentally determined functions for protein i, (t) is the set of predicted functions for protein whose scores are greater equal to f, and [ represents the indicator function. After that, the precision and recall averaged over all proteins:
precision())
M(1 Σ precision, recall(()
N Σ recall:(1).
(30)
where N denotes the total number of proteins, and (t) denotes the number of proteins which contain at least one prediction above the threshold i.e., [(())) O.
Based on these two metrics, Fmax score defined as the maximum value of F-measure over all thresholds:
2 precision(t) recall(t)
Fmax
max
(31)
precision (t) recall(t)
Global Spearman's p for PSR measures the correlation between the predicted global distance test (GDT TS) score and the ground truth. It computes the Spearman's 's between the prediction and the ground truth over test proteins without considering the different biopolymers that these proteins lie in.
Mean Spearman's for PSR also measures the correlation between GDT 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 's within each group, and finally reports the mean Spearman's 's p over all groups.
E Results of Pre-Training 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 different dataset sizes. Since previous work by Zhang et al. [7 showed minimal differences between using experimental or predicted structures, we conduct experiments on the AlphaFold Database in this section and do not use PDB pre-training dataset. We utilize the preprocessed clustered
impact of pre-training on

---

Table Atom-level results of pre-training on different sizes of datasets on Atom3D tasks.
PSR
MSP
PIP
RES
Pre-training Dataset
Size
Global p
Mcan P
AUROC
AUROC
Acc.
AlphaFold DB vl
365K
0.829
0.546
0.698
0.884
0.460
Clustered AlphaFold DB
10K
0.816
0.501
0.586
0.880
0,444
Clustered AlphaFold DB
50K
0.797
0,498
0.648
0.879
0,450
Clustered AlphaFold DB
100K
0.805
0.540
0.685
0.882
0.454
Clustered AlphaFold DB
500K
0.848
0.537
0.599
0.880
0.459
Clustered AlphaFold DB
2.2M
0.840
0.560
0.700
0.882
0.460
AlphaFold Database provided in [6, which includes 2.2M non-singleton clusters with an average of 13.2 proteins per cluster and average PLDDT of 71.59. For each cluster, 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 reveal a general trend of increased performance with larger datasets. However, for certain tasks like MSP, the performance does consistently improve, possibly due the limited size and variability of the downstream datasets. Overall, scaling the model to the entire AlphaFold Database holds promise for achieving performance gains.
Results of Multimodal Pre-Training Baselines
Table 8: Atom-level results of multimodal pre-training baselines Atom3D tasks.
PSR
MSP
PIP
RES
Method
Global p
Mean p
AUROC
AUROC
Acc.
GearNet-Edge Edge
0.782
0.488
0,633
0.868
0,441
w/ pre-ttraining
Residue Type Prediction
0.826
0,518
0.620
0.879
0.449
Residue Type & Distance Prediction
0.817
0,518
0.665
0,873
0.402
Residue Type & Angle Prediction
0.837
0.524
0.642
0.878
0.415
Residue Type & Dihedral Prediction
0.823
0.494
0,597
0.871
0.414
DiffPreT
0.821
0,533
0,680
0.880
0.452
SlamDifr
0.829
0.546
0.698
0.884
0.460
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/dihedrr prediction objectives. Following established settings, we pre-train the encoder and present the results in Table g While the Residue Type & Angle Prediction method achieves higher performance PSR, there are still substantial gaps compared to our DiffPreT and Sifm across other tasks. Notably, the introduction of geometric property prediction tasks leads to drop in performance 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.
RES, indicating
G Results of Different Diffusion Models for Pre-Training
In Sec. we consider joint diffusion models on protein sequences and structures for pre-training. In this section, explore the performance of different diffusion models when applied for pre-training The results are shown Table
First, we simply run diffusion models on protein sequences and structures for pre-training, both of which achieve improvement compared with the baseline GearNet-Fdge. By combining two diffusion models, DiffPreT can achieve better performance PIP and PSR. This advantage will be further
25

---

Table 9: Results of different diffusion models for pre-training atom-level Atom3D tasks.
PIP
RES
PSR
Method
AUROC
Accuracy
Global p
Mean p
GearNet-Edge
0.868
0,441
0.782
0.488
w/ pre-training
DiffPreT
0.880
0,452
0.821
0.533
SiamDiff
0.884
0.460
0.829
0.546
Sequence Diffusion
0.879
0.456
0,802
0.508
Structure Diffusion
0,877
0,448
0,813
0.518
Torsional Diffusion
0.877
0.442
0.819
0.505
increased after using siamese diffusion trajectory prediction as shown in Table 3 It can be observed that sequence diffusion achieves better performance than DiffPreTT 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, 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 promising direction to explore more different diffusion models for pre-training, e.g., diffusion models on backbone dihedral angles.
H
Results of Pre-Traeae Fre-Training GVP
Table 10: Atom-level results of GVP on Atom3D tasks. Accuracy is abbreviated as Acc.
PSR
MSP
PIP
RES
Mean
Method
Rank
Glohal
Mean p
AUROC
AUROC
Acc.
GVP
0.809
0.486
0.610
0.846
0.481
8.6
Denoising Score Matching
0.849
0.535
0.625
0.851
0,529
5.4
pre-training
Residue Type Prediction
0.845
0,527
0.652
0.847
0,518
5.8
Distance Prediction
0.825
0.513
0.632
0.836
0.483
7.6
Angle Prediction
0.872
0.545
0.637
0.881
0.557
2.0
Dihedral Prediction
0.852
0,538
0.677
0.881
0.555
2.2
Multiview Contrast
0.848
0,518
0.656
0.833
0.490
6.4
DiffPreT
0.850
0.540
0.631
0.851
0.542
4.4
SiamDiff
0.854
0.548
0.673
0.863
0.554
2.2
To study the effect of our proposed pre-training methods on different backbone models, we show the pre-training results GVP [1] 42 in this section.
Setup. GVP constructs atom graphs and adds vector channels for modeling equivariant features. The original design only includes atom types 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 atom features. Then, the default configurations Jing 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 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 those in App.
26

---

Table 11: Spearman correlation GB1 with the 2-vs-rest split.
Method
Spearmanr
CNN
0.320
ESM-1b
0,550
GearNet-Edge
0.651
w/ Multiview Contrast
0.714
w/ Residue Type Prediction
0.697
w/ Distance Prediction
0.685
w/ Angle Prediction
0.686
w/ Dihedral Prediction
0.661
w/ SiamDiff
0.747
Table 12: Comparison between SifmDiff with random torsional perturbation and sampling from rotamer library.
PIP
MSP
RES
PSR
Method
AUROC
AUROC
Accuracy
Global p
Mean
SiamDiff
0.884 0.877
0.698
0.460 0.449
0.829 0.834
0.546
w/ rotamer library
0.631
0.521
Experimental results. The results are shown in Table I0 angle prediction and dihedral prediction are the top three. This is different from what we observe in Table I where residue type and distance prediction are more competitive baselines. We hypothesize that this is because GearNet-Fdge includes angle information in the encoder while GVP does not. Therefore, GVP will benefit more from pre-training methods with supervision on angles. Nevertheless, find that SifmDiff is the only method that performs well on different backbone models.
Among all pre- pre-training methods, Sifm
I Results on Protein Engineering Task
To further prove the effectiveness of proposed pre-training methods, add experiments the GB1 dataset from FLIP [6]. As this is protein engineering task with mutated sequences, we assume that the backbone structure remains unchanged after mutation, save costs 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 II the attached file, alongside CNN and ESM-Ib baselines from the FLIP paper.
According Table II with using only sequential information, even under the assumption that 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.
observe that modeling structural information is beneficial compared
J Effects of Approximate Conformer Generation
In SiamDiff, we introduce approixmate conformer generation mechanism by randomly rotating tor- sional angles. We hypothesize that highly realistic conformers are not vital for better representations. To confirm, an extra rebuttal experiment is performed in Table [2] Instead of random perturbation, we sample from rotamer library 622] based on residue types and backbone angles. Table T2 shows random torsional perturbation still outperforms sampling from a rotamer library in most tasks, con- firming our hypothesis. This can 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 realistic conformer distributions.
27

---

(A) GearNet- Edge representations with random initialization
(B) GearNet-Eol Edge representations pre-trained with sirst-stage SiamDiff
S
(C) GearNet- representations pre-trained with two-stage SiamDiff
3r-ge330-Eimaen v1.ach
AF-OGEBBEImas (E-00P0-Fi vlats
Figure 4: Visualization of GearNet-Edge representations on 4 randomly selected proteins from AlphaFold Database with sequence and structure noises at different levels. Three models are employed for extracting representations: (A) without pre-training,
(B) with first-stage SiamDiff, and (C) with two-stage SiamDiff pre-training. For each protein, we consider t = I.... 100 noise levels and randomly generate 32 noisy samples for each noise level using the scheme in SiamDiff. Samples with small noise levels are highlighted in blue, while those with large noise levels are marked in red.
K
Visualization of SiamDiff Embeddings upon Noises
To explore pre-training insights, visualize UMAP representations of 4 random AlphaFold DB proteins in Fig. Several interesting phenomena can be observed:
Randomly initiated representations in Fig. 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 seales, the encoder maintains the color smoothness of the trajectory, which desired for effective denoising during the backward diffusion process. Intrigu- ingly, pre-training narrows the trajectory compared to the broader trajectory without pre-training, particularly the 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 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 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 more effective for fine-grained downstream tasks like PSR and
28

---

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


Reviews:
[{"review_id": "SHRPPdfYg5", "reviewer": "Reviewer_MBi3", "summary": "The 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. ", "questions": "- 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?", "limitations": "- SiamDiff requires structures at inference time. \n- The paper only considers one diffusion process out of many possible choices. ", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "### 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your appreciation of our work! We respond to your questions and concerns below:\n\n>**Q1: The evaluation on more types of downstream tasks would make this paper more significant.**\n\nThanks for the suggestion. We believe this additional experiment in the global response showcases the potential of SiamDiff on protein engineering tasks. However, we acknowledge that more effort is required when considering more complicated settings, e.g., insertion and deletion. Also, task-specific designs need to be explored to apply structure-based models for zero-shot fitness prediction. Due to the limited time in the rebuttal period, we leave these explorations as future work.\n\n>**Q2: Different structure and sequence diffusion schemes are not thoroughly explored in the current draft.**\n\nThanks for pointing out this important aspect! \n\nFor structure diffusion, we resort to **diffusion on coordinates** for its **effectiveness on learning molecular representations that has been proven in recent small molecule pre-training works [a,b]**. We agree that **diffusion on amino acid frames and bond/torsion angles better fit the inductive biases of protein structure generation** as shown in FoldingDiff [c] and RFdiffusion [d]. It is intriguing to study the effectiveness of these structure diffusion schemes on protein representation learning. We will definitely investigate them in our future work.\n\nFor sequence diffusion, Autoregressive Diffusion Models (ARDMs) [e] are equivalent to the infinite time limit of Discrete Denoising Diffusion Probabilistic Models (D3PMs) [f] that are used in our current work. Therefore, **ARDMs are maximally expressive and can potentially enhance the effectiveness of our proposed DiffPreT and SiamDiff methods.** To explore this possibility, we conducted an initial experiment during the rebuttal by replacing D3PMs in SiamDiff with ARDMs. The results are presented in Table A.\n\nTable A: Comparison between SiamDiff with D3PMs and ARDMs.\n|#Method|PIP|MSP|RES|PSR||\n|:----:|:----:|:----:|:----:|:----:|:----:|\n||AUROC|AUROC|Acc.|Global $\\rho$|Mean $\\rho$|\n|**SiamDiff**|**0.884**|**0.698**|**0.460**|**0.829**|**0.546**|\n|w/ ARDM|0.883|0.640|0.450|0.828|0.533|\n\nAs shown in Table A, the results of ARDMs are quite close to those of SiamDiff on some tasks, but generally fall short in performance. Our hypothesis is that the advantages of ARDMs in the original paper stem from their random assigned order and causal masking in Transformers. The causal masking scheme poses challenges for GNN-based encoders. In our implementation, we simply remove edges that do not adhere to the assigned order. Further investigation is needed to explore how to adapt ARDMs to proteins, and this aspect is left for future work.\n\n[a] Zaidi, Sheheryar, et al. \"Pre-training via denoising for molecular property prediction.\" ICLR, 2023.\n\n[b] Liu, Shengchao, Hongyu Guo, and Jian Tang. \"Molecular geometry pretraining with se (3)-invariant denoising distance matching.\" ICLR, 2023.\n\n[c] Wu, Kevin Eric, et al. \"Protein structure generation via folding diffusion.\" arXiv, 2022.\n\n[d] Watson, Joseph L., et al. \"De novo design of protein structure and function with RFdiffusion.\" Nature, 2023. \n\n[e] Hoogeboom, Emiel, et al. \"Autoregressive diffusion models.\" ICLR, 2022.\n\n[f] Austin, Jacob, et al. \"Structured denoising diffusion models in discrete state-spaces.\" NeurIPS, 2021.\n\n>**Q3: The paper should include ablations on EC.**\n\nThanks for the suggestion. We agree that including the ablation on EC will provide a better understanding of components of SiamDiff on residue-level tasks. However, this is infeasible during the rebuttal period, due to the large amount of computational resources needed (5 pre-training baselines * 3 repeated experiments, each with 4 GPUs and 24 hours fine-tuning). We will run the experiments and add the ablation study in the final version.\n\n>**Q4: It’s not clear why ESM-2-650M-GearNet doesn’t require structure as input.**\n\nWe would like to clarify that what we state in the paper is ESM-2-650M (the sequence encoder part of ESM-2-650M-GearNet) can only extract representations from protein sequences and cannot take protein structures as input, making it unsuitable for structure-related tasks in Atom3D. **ESM-2-650M-GearNet can definitely take protein structures as input and extract structural representations by its GearNet component.** We will make this point clearer in the revision.\n\n>**Q5: How is EC a residue-level task? Don’t the labels apply to the entire protein?**\n\nWe clarify that the “residue-level” here means that **residue-level structures are used for task prediction, instead of predicting per-residue labels**. Therefore, the labels of the EC task are applied to entire proteins, and we use residue-level protein structures to predict such per-protein labels. We will state this point more clearly in the revised version.\n\n>**Q6: Is there a way to use this model on proteins without a structure? SiamDiff requires structures at inference time.** \n\nAs joint pre-training methods upon protein sequences and structures, DiffPreT and SiamDiff require protein structures as input for model prediction. Given the recent advancement in protein structure prediction, it can be easy to obtain accurate structures for most proteins, though with additional inference time. However, we argue that in most protein-related tasks, the importance of accuracy outweighs that of inference speed, including protein function and protein-protein interaction prediction. Besides, even for high-throughput tasks dominated by sequence-based methods like protein engineering, structure-based methods can be easily applied with little extra effort and large improvements, as shown in the GB1 experiment in the global response. Hence, we argue that the requirement of protein structures should not be a downside of our method.\n"}, {"review_id": "NpATMpWgPw", "reviewer": "Reviewer_s72W", "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.", "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", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThis 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.\n", "weaknesses": "Lack of evaluation against other downstream predictors and backbone models for downstream prediction.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful comments and great suggestions! We respond to your concerns  below:\n\n>**Q1: Lack of evaluation against other downstream predictors and backbone models for downstream prediction.**\n\n**We would like to argue that the focus of the paper is the development and evaluation of new pre-training algorithms, not achieving state-of-the-art performance on a specific task.** Therefore, we constantly use GearNet-Edge as the backbone model and 2-layer MLP as the prediction head so as to **have a fair comparison among different pre-training methods**. We also try GVP as the backbone model in Appendix H, to show that our methods can be applied on different encoders. Indeed, it is interesting to study the performance of different pre-training methods under other backbone models like ProNet [a] and CDConv [b]. Since the response period is short, we will try our best to finish these studies and add them to the future paper version. \n\n[a] Wang, Limei, et al. \"Learning hierarchical protein representations via complete 3d graph networks.\" ICLR, 2022.\n\n[b] Fan, Hehe, et al. \"Continuous-Discrete Convolution for Geometry-Sequence Modeling in Proteins.\" ICLR, 2022.\n\n>**Q2: In the ablation, the model without structure diffusion is competitive with the others. What’s the additional complexity of joint diffusion against sequence-only diffusion?**\n\nFirst, we want to argue that structure diffusion is an important component of SiamDiff. According to Table 3, the baseline without structure diffusion does not perform well on structure-informed tasks like PIP, MSP and PSR. We also perform ablation studies on residue-level tasks during rebuttal, the results of which are shown in Table A. Based on these results, we can conclude that structure diffusion brings consistent improvements across all considered tasks, which proves its necessity.\n \nTable A: Ablation study on residue-level tasks.\n|#Method|EC||MSP|PSR||\n|:----:|:----:|:----:|:----:|:----:|:----:|\n||AUPR|$F_{\\max}$|AUROC|Global $\\rho$|Mean $\\rho$|\n|**SiamDiff**|**0.878**|**0.857**|**0.700**|**0.856**|**0.521**|\n|w/o structure diffusion|0.868|0.850|0.671|0.826|0.502|\n\nIn terms of complexity, it should be noted that the bottleneck of computation still lies in the protein encoder. The encoder requires message passing between protein atoms and introduces linear (or quadratic) cost with respect to the number of edges in the protein graph. The process of perturbing and recovering protein structures only requires linear complexity with respect to the number of atoms/residues in the protein, which is similar to sequence diffusion and is cheap compared with the protein encoder. Therefore, it is good to include structure diffusion considering its benefits and little cost.\n\n\n>**Q3: Why is the gap of sequence denoising accuracy between joint diffusion and sequence-only diffusion so pronounced?**\n\nThanks for the question! It is an interesting observation that sequence diffusion only can achieve around 0.8 recovery accuracy, while using joint diffusion will decrease the accuracy to 0.4. As explained in the paper, the phenomenon can be attributed to the introduction of large structural noise. The large effect of structural noise implies that the sequence recovery may rely on the backbone conformation of predicted residues (note that we remove the corresponding side chain atoms to avoid information leakage). After perturbing the backbone structures, it will be much more difficult to infer the residue types. This hypothesis can be validated by the phenomenon observed on residue-level pre-training: when using residue-level SiamDiff, we only keep the alpha Carbon atoms instead of all three backbone atoms. The sequence diffusion can only achieve around 0.4 recovery, while the joint diffusion will decrease the accuracy to 0.3, which has a much smaller gap. Therefore, we conclude that the large gap is because atom-level sequence denoising is *too easy* with the correct backbone information, while introducing structural noise brings difficulty to the pre-training task."}, {"review_id": "nHag55nPh5", "reviewer": "Reviewer_Qtp8", "summary": "The 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.", "questions": "- 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.", "limitations": "The authors discuss limitations in the appendix. It would be nice if some of this discussion is transferred to the main paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your careful reading! While some questions may be due to misunderstanding, we find that your suggestions are very helpful for us to improve the quality and clarity of our paper! We respond to your concerns below:\n\n>**Q1: What noise level is used to get the representation for downstream tasks? Why noisy protein representations can help?**\n\nWe clarify that we only introduce noises during pre-training and will use the original protein during fine-tuning, as in denoising pre-training [a,b]. The aim of pre-training is to make the learned representations capture structural and sequential details. We don’t need to introduce additional noise for downstream tasks. We will clarify this point in the revised version.\n\n[a] Zaidi et al. \"Pre-training via denoising for molecular property prediction.\" ICLR, 2023.\n\n[b] Liu et al. \"Molecular geometry pretraining with se (3)-invariant denoising distance matching.\" ICLR, 2023.\n\n>**Q2: Please comment on the size of the latent representation that you use.**\n\nAs stated in the Appendix D, our approach employs 128 hidden dimensions for atom-level protein structure tasks and 512 for residue-level tasks. We use a smaller number of hidden dimensions in the atom-level cases so as to have moderate computational cost under the larger-scale atom-level graphs. \n\n>**Q3: Are the performance improvements on downstream tasks statistically significant?**\n\nThanks for the question. To validate whether the improvements are statistically significant, we conduct a one-tailed t-test between our methods and the second best pre-training baselines on tasks that our method is the best at. We repeat the experiment with five different seeds and show the results in Table A.\n\nTable A: Statistical significance test results.\n|#Task|#Method|p-value|t-statistics|\n|:----:|:----:|:----:|:----:|\n|PIP (atom)|SiamDiff v.s. Residue Type Prediction|0.001|6.78|\n|RES (atom)|SiamDiff v.s. Residue Type Prediction|$5.0 \\times 10^{-5}$|15.54|\n|PSR mean $\\rho$ (atom)|SiamDiff v.s. Distance Prediction|0.003|5.11|\n|MSP (atom)|SiamDiff v.s. Distance Prediction|0.04|2.33|\n|EC auprc (residue)|SiamDiff v.s. Multiview Contrast|>0.1|-|\n|PSR global $\\rho$ (residue)|SiamDiff v.s. Residue Type Prediction|0.01|3.68|\n|PSR mean $\\rho$ (residue)|SiamDiff v.s. Residue Type Prediction|0.01|3.62|\n\nIt can be observed that for all tasks except EC, the performance improvements are statistically significant under a p-value less than 0.1 (i.e., with t-statistics surpassing the critical value of the corresponding test). We will include the test results in the camera ready version and acknowledge that the improvements on EC are not statistically significant. \n\n>**Q4: Are the conformers derived by side chain perturbation “biologically relevant”?**\n\nPlease see the global response for details.\n\n>**Q5: Sec. 3.4 is showing results even before all methods are explained.**\n\nFor Sec. 3.4, we want to emphasize the contribution of the proposed two-stage noising scheme as a part of our method and **the results are necessary to support our claims in Sec. 3.4**. First, we provide intuitive insights into the challenges of denoising sequence and structure with varying noise levels, emphasizing the advantages of two-stage noise scheduling. To validate these points, we examine structure denoising loss and sequence denoising accuracy across diverse noise levels and schedules, effectively showcasing our two-stage noise scheduling approach. \n\n>**Q6: How can you get both sequence and structure representations using GearNet-Edge?**\n\nIn the paper, we have mentioned in the first footnote that the structure encoder refers to those that take both sequences and structures as input. Notably, GearNet-Edge employs sequential edges, linking consecutive residues in a protein sequence. In this way, message passing is performed according to the protein sequence, and thus GearNet-Edge can extract sequence representations. In another way, spatial and KNN edges are constructed to represent the protein structure, which enables GearNet-Edge to extract structure representations. By combining these different types of edges, GearNet-Edge extracts both sequence and structure representations. \n\n>**Q7: It would help to include some intuitive graphical results.**\n\nPlease see the global response for details.\n\n>**Q8: A related work of protein structure-sequence co-design [a] is not referred to.**\n\nThis work is closely related to ours in terms of performing sequence-structure joint diffusion. However, our work focuses on learning informative protein representations, while this work focuses on protein generation. We will discuss these connections and differences in the revision.   \n\n[a] Lisanza, et al. \"Joint generation of protein sequence and structure with RoseTTAFold sequence space diffusion.\" bioRxiv, 2023.\n\n>**Q9: The metric “Mean Rank” does not add much to all other results.**\n\nWe would like to argue that **“Mean Rank” is a commonly used metric for benchmarking methods on a diverse set of tasks**. As shown in Tables 1 and 2, some baseline methods perform well (even the best) on one or two tasks while cannot perform consistently well on all tasks. Since it is hard to measure such performance consistency using individual task-specific metrics, we introduce “Mean Rank'' as a metric for measuring performance consistency across tasks. Based on this metric, we can observe that DiffPreT and SiamDiff are the only two methods that perform consistently well on all tasks. \n\n>**Q10: The discussion on computational cost is missing.**\n\nRegarding computational cost, the protein encoder remains the main bottleneck, involving linear/quadratic cost due to message passing among atoms in the graph.  Perturbing/recovering protein sequences and structures carries linear complexity relative to atoms/residues, significantly cheaper than the protein encoder. We will add a paragraph in the method section to discuss the computational cost introduced in pre-training.\n"}, {"review_id": "YYv6VrGOP6", "reviewer": "Reviewer_E714", "summary": "The 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.", "questions": "1. 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", "limitations": "Random perturbations on the side-chain torsional angles may not conform to the actual distribution of the conformational change.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful comments and golden suggestions! We respond to your concerns as below:\n\n>**Q1: The random side-chain perturbation does not reflect the actual \"physics underlying the conformational change**\n\nPlease see the global response for details.\n\n>**Q2: Discarding the mutual information maximization scheme only has very minor effect on the performance**\n\nWe respectfully disagree with the reviewer on this point. It should be noted that the SiamDiff method without mutual information maximization scheme is exactly the DiffPreT method. As shown in Tables 1 & 2, the improvement of SiamDiff over DiffPreT is consistent and significant, especially on residue-level tasks.\n\n>**Q3: The organization of the paper can be improved. 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\nThanks for the suggestion! In the revised version, we will move a portion of preliminaries of diffusion models to the appendix, thus leaving more space to discuss and evaluate the effect of different diffusion design choices and conduct more downstream evaluation. We believe this can be done in the camera ready version with an additional content page limit.\n\n>**Q4: Do you think the mutual denoising scheme can also work on homologs as you have mentioned the problem in line 50?**\n\nThis is an interesting question! It is indeed possible to use structural homologs instead of randomly generated conformers for mutual denoising. Though there might not be one-to-one residue mapping, we can still use the alignment tools between protein structures [a,b] to get residue/node correspondence for mutual denoising. Nevertheless, extra efforts are required to study whether these realistic conformers are better for pre-training, as discussed in our response to Q1. We leave this question to future work.\n\n[a] Zhang, Yang, and Jeffrey Skolnick. \"TM-align: a protein structure alignment algorithm based on the TM-score.\" Nucleic acids research 33.7 (2005): 2302-2309.\n\n[b] Holm, Liisa. \"Using Dali for protein structure comparison.\" Structural Bioinformatics: Methods and Protocols (2020): 29-42.\n\n>**Q5: What do you think is the reason that diffusion-based pretraining produces better results than denoising score matching pretraining in your experiments? Could the denosing score mathching baseline be underestimated because it is implemented on the distance matrix instead of the absolute coordinates?**\n\nRegarding denoising score matching, we agree that it is more appropriate and intuitive to add noises directly to the absolute coordinates instead of distance matrices. In our ablation study, the baseline (SiamDiff w/o sequence diffusion) can be considered as an augmented version of denoising score matching. The superiority of SiamDiff over denoising score matching can be attributed to two main factors: sequence-structure joint diffusion and multi-level noise scheduling, both of which have been extensively discussed in our ablation study.\n\nFurthermore, as mentioned in Sec. 3.4, another drawback of previous denoising-based molecular pre-training methods is that they typically treat the noise level as a hyperparameter to be tuned [a]. This presents a significant challenge in selecting the optimal hyperparameter for pre-training, as it becomes difficult to capture both coarse- and fine-grained features effectively.\n\n[a] Zaidi, Sheheryar, et al. \"Pre-training via denoising for molecular property prediction.\" ICLR, 2023.\n\n>**Q6: 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 lower bound of mutual information and the loss you derived from the RHS of the approximation?**\n\nThanks for your careful reading and insightful question! The approximation made in line 743 is based on replacing representations with diffusion trajectories. The latter is easier to sample and provides more informative targets for recovery. Therefore, this decision is more of a practical choice. In theory, it would be possible to derive a lower bound for the error introduced in this approximation by introducing some assumptions about the information loss between representations and trajectories. However, since the primary focus of the paper lies in the application of the pre-training algorithm, the presented theorems serve as a justification for the reasonableness of our method and only consider rough approximations. Characterizing theoretically the deviation introduced by such approximations is a very interesting direction for future work.\n\n>**Q7: What is the value of $t$ in the diffusion model in the fine-tuning phase?**\n\nWe want to clarify that we only introduce noises in the pre-training stage and will use the original protein in the fine-tuning stage, which is common in denoising pre-training methods [a,b]. The aim of pre-training is to make the learned representations capable of capturing structural and sequential details. We don’t need to introduce additional noise for downstream tasks. We will clarify this point in the revised version.\n\n[a] Zaidi, Sheheryar, et al. \"Pre-training via denoising for molecular property prediction.\" ICLR, 2023.\n\n[b] Liu, Shengchao, Hongyu Guo, and Jian Tang. \"Molecular geometry pretraining with se (3)-invariant denoising distance matching.\" ICLR, 2023.\n"}]

Author rebuttal:
We extend our gratitude to all reviewers for valuable feedback. We’ve made significant improvements based on your suggestions. Here is a brief summary of important points:

>**New benchmark results on protein engineering task (Reviewer MBi3)** 

During the rebuttal period, we followed your suggestion to include the GB1 dataset from FLIP in experiments. 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.**

>**Biological relevance of random torsional perturbation scheme (Reviewer Qtp8, E714)** 

First, we reiterate that random side-chain perturbation is commonly used for simulating conformers [a,b]. Moreover, removing structures with clashes ensures that generated conformers are physically plausible. So the generated conformer distribution is biologically relevant.

Methodologically, while adding torsional noises without changing the backbone when generating conformers, Gaussian noises are introduced to the backbone during forward diffusion. This makes our encoder capture both backbones and side-chain noises effectively.

Besides, in pre-training, highly realistic conformers aren’t vital for better representations. To confirm, an extra rebuttal experiment (Table 12 in attached file) is performed. Instead of random perturbation, we sample from a rotamer library [c] 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.

In summary, while random torsional perturbation may not be as realistic as rotamer library-based or force field-based methods, it holds **biological relevance, is easy to implement, and proves to be a practical pre-training choice due to performance advantages**.

[a] Ho et al. "Probing the flexibility of large conformational changes in protein structures through local perturbations." PLoS computational biology, 2009.

[b] Ho et al. "Conserved tertiary couplings stabilize elements in the PDZ fold, leading to characteristic patterns of domain conformational flexibility." Protein Science, 2010.

[c] Shapovalov et al. "A smoothed backbone-dependent rotamer library for proteins derived from adaptive kernel density estimates and regressions." Structure, 2011.

>**New visualization results (Reviewer Qtp8)**

We have added visualization results in the attached pdf file. To explore pre-training insights, we visualize UMAP representations of 4 random AlphaFold DB proteins in Fig. 4. 

Several interesting phenomena can be observed:
1. *Randomly initiated representations* in Fig. 4(A) 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.**
2. 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. 4(A).
3. 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. 4(B), which may imply some forgetting behavior in the second-stage pre-training. 

>**Paper presentation (Reviewer MBi3, Qtp8, E714)** 

We acknowledge the concerns raised by the reviewers regarding the paper's presentation. Due to the extensive methodology and experimental contribution, it was challenging to fit everything within the 9-page limit. In the submitted version, we dedicated a lengthy section to explaining the fundamental concepts, making it more accessible to readers unfamiliar with diffusion models on proteins. However, we recognize that many intriguing observations and experiments were relegated to the appendix. We believe this problem can be addressed in the camera ready version with an additional page limit. We plan to reorganize our paper based on your suggestions, including the shortening of Sec. 3 and the addition of more details in Sec. 4, merging Sec. 4.4 into the Related Work section, and providing necessary descriptions about the tasks and experimental settings.


Meta-review: The paper propose to pretrain neural models for predicting properties of proteins from sequences by maximizing mutual information between denoising trajectories of pairs of related conformers. The authors lay out the diffusion model first, which constituted DiffPreT, and then propose the trajectory based maximum mutual information objective, called SimDiff. They show this objective can be optimized by optimizing a lower bound on it.

The reviewers found the paper quite compelling. Some reviewers raise the point that maybe in some case the results were not SOTA, but they concurred with the authors' contention that goal of the paper was to show that the approach results in a useful pre-training strategy, that produces strong gains in several cases.

I felt that the authors' did a really good job of answering the reviewer's criticisms. As a result several of the reviewers raised their original assessment and were quite positive on the paper.
Decision: Accept (spotlight)