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

Notation. A protein with $n_{r}$ residues (amino acids) and $n_{a}$ atoms can be represented as a sequence-structure tuple ${\mathcal{P}}=({\mathcal{S}},{\mathcal{R}})$. We use ${\mathcal{S}}=[s_{1},s_{2},\cdots,s_{n_{r}}]$ to denote its sequence with $s_{i}$ as the type of the $i$-th residue, while ${\mathcal{R}}=[{\bm{r}}_{1},{\bm{r}}_{2}...,{\bm{r}}_{n_{a}}]\in\mathbb{R}^{n_{a}\times 3}$ denotes its structure with ${\bm{r}}_{i}$ as the Cartesian coordinates of the $i$-th atom. To model the structure, we construct a graph for each protein with edges connecting atoms whose Euclidean distance below a certain threshold.

Equivariance. Equivariance is ubiquitous in machine learning for modeling the symmetry in physical systems [66, 74] and is shown to be critical for successful design and better generalization of 3D networks [45]. Formally, a function ${\mathcal{F}}:{\mathcal{X}}\rightarrow{\mathcal{Y}}$ is equivariant w.r.t. a group $G$ if ${\mathcal{F}}\circ\rho_{{\mathcal{X}}}(x)=\rho_{{\mathcal{Y}}}\circ{\mathcal{F}}(x)$, where $\rho_{{\mathcal{X}}}$ and $\rho_{{\mathcal{Y}}}$ are transformations corresponding to an element $g\in G$ acting on the space ${\mathcal{X}}$ and ${\mathcal{Y}}$, respectively. The function is invariant w.r.t $G$ if the transformations $\rho_{{\mathcal{Y}}}$ is identity. In this paper, we consider SE(3) group, i.e., rotations and translations in 3D space.

Problem Definition. Given a set of unlabeled proteins ${\mathcal{D}}=\{{\mathcal{P}}_{1},{\mathcal{P}}_{2},...\}$, our goal is to train a protein encoder $\phi({\mathcal{S}},{\mathcal{R}})$ to extract informative $d$-dimensional residue representations ${\bm{h}}\in\mathbb{R}^{n_{r}\times d}$ and atom representations ${\bm{a}}\in\mathbb{R}^{n_{a}\times d}$ that are SE(3)-invariant w.r.t. protein structures ${\mathcal{R}}$.

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

where $q({\mathcal{P}}^{t}|{\mathcal{P}}^{t-1})$ defines the forward process at step $t$ and $p_{\theta}({\mathcal{P}}^{t-1}|{\mathcal{P}}^{t})$ with learnable parameters $\theta$ defines the reverse process at step $t$. We decompose the forward process into diffusion on protein structures and sequences, respectively. The decomposition is represented as follows:

where the reverse processes use representations ${\bm{a}}^{t}$ and ${\bm{h}}^{t}$ learned by the protein encoder $\phi_{\theta}({\mathcal{S}}^{t},{\mathcal{R}}^{t})$.

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

Here, $\beta_{1},...,\beta_{T}$ and $\rho_{1},...,\rho_{T}$ are a series of fixed variances and masking ratios, respectively. $\text{random\_mask}({\mathcal{S}}^{t-1},\rho_{t})$ denotes the random masking operation, where each residue in ${\mathcal{S}}^{t-1}$ is masked with a probability of $\rho_{t}$ at time step $t$.

Reverse process on structures $p_{\theta}({\mathcal{R}}^{t-1}|{\mathcal{P}}^{t})$. The reverse process on structures is parameterized as a Gaussian with a learnable mean $\mu_{\theta}({\mathcal{P}}^{t},t)$ and user-defined variance $\sigma_{t}$. Given the availability of ${\mathcal{R}}^{t}$ as an input, we reparameterize the mean $\mu_{\theta}({\mathcal{P}}^{t},t)$ following Ho et al. [36]:

where $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$ and the network $\epsilon_{\theta}(\cdot)$ learns to decorrupt the data and should be translation-invariant and rotation-equivariant w.r.t. the structure ${\mathcal{R}}^{t}$.

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

where ${\mathcal{N}}^{t}(i)$ denotes the neighbors of the atom $i$ in the corresponding graph of ${\mathcal{P}}^{t}$. Note that $\epsilon_{\theta}({\mathcal{P}}^{t},t)$ achieves the equivariance requirement, as $m_{i,j}$ is SE(3)-invariant w.r.t. ${\mathcal{R}}^{t}$ while ${\bm{r}}_{i}^{t}-{\bm{r}}_{j}^{t}$ is translation-invariant and rotation-equivariant w.r.t. ${\mathcal{R}}^{t}$.

Reverse process on sequences $p_{\theta}({\mathcal{S}}^{t-1}|{\mathcal{P}}^{t})$. For the reverse process $p_{\theta}({\mathcal{S}}^{t-1}|{\mathcal{P}}^{t})$, we adopt the parameterization proposed in [4]. The diffusion trajectory is characterized by the probability $q({\mathcal{S}}^{t-1}|{\mathcal{S}}^{t},\tilde{{\mathcal{S}}}^{0})$, and we employ a network $\tilde{p}\theta$ to predict the probability of ${\mathcal{S}}^{0}$:

We define the predictor $\tilde{p}_{\theta}$ with residue representations ${\bm{h}}^{t}$. For each masked residue $i$ in ${\mathcal{S}}^{t}$, we feed its representation ${\bm{h}}_{i}^{t}$ to an MLP and predict the type of the corresponding residue type $s^{0}_{i}$ in ${\mathcal{S}}^{0}$:

where the softmax function is applied over all residue types.

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

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

Both loss functions can be simplified as follows.

Structure loss ${\mathcal{L}}^{({\mathcal{R}})}$. It has been shown in Ho et al. [36] that the loss function can be simplified under our parameterization by calculating KL divergence between Gaussians as weighted L2 distances between means $\epsilon_{\theta}$ and $\epsilon$ (see details in App. C.3):

where the coefficients $\gamma_{t}$ are determined by the variances $\beta_{1},...,\beta_{t}$. In practice, we follow Ho et al. [36] to set all weights $\gamma_{t}=1$ for the simplified loss ${\mathcal{L}}_{\text{simple}}^{({\mathcal{R}})}$.

Since $\epsilon_{\theta}$ is designed to be rotation-equivariant w.r.t. ${\mathcal{R}}^{t}$, to make the loss function invariant w.r.t. ${\mathcal{R}}^{t}$, the supervision $\epsilon$ is also supposed to achieve such equivariance. Therefore, we adopt the chain-rule approach proposed in Xu et al. [81], which decomposes the noise on pairwise distances to obtain the modified noise vector $\hat{\epsilon}$ as supervision. We refer readers to Xu et al. [81] for more details.

Sequence loss ${\mathcal{L}}^{({\mathcal{S}})}$. Since we parameterize $p_{\theta}({\mathcal{S}}^{t-1}|{\mathcal{P}}^{t})$ with $\tilde{p}_{\theta}(\tilde{{\mathcal{S}}}^{0}|{\mathcal{P}}^{t})$ and $q({\mathcal{S}}^{t-1}|{\mathcal{S}}^{t},\tilde{{\mathcal{S}}}^{0})$ as in (6), it can be proven that the $t$-th KL divergence term in ${\mathcal{L}}^{({\mathcal{S}})}$ reaches zero when $\tilde{p}_{\theta}(\tilde{{\mathcal{S}}}^{0}|{\mathcal{P}}^{t})$ assigns all mass on the ground truth ${\mathcal{S}}^{0}$ (see proof in App. C.4). Therefore, for pre-training, we can simplify the KL divergence to the cross-entropy between the correct residue type $s_{i}^{0}$ and the prediction:

where CE$(\cdot,\cdot)$ denotes the cross-entropy loss.

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

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

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

Our method begins by generating different conformers of a given protein. However, direct sampling requires an accurate characterization of the energy landscape of protein conformations, which can be difficult and time-consuming. To address this issue, we adopt a commonly used scheme for sampling randomly simulated conformers by adding torsional perturbations to the side chains [35].

Specifically, given the original protein ${\mathcal{P}}=({\mathcal{S}},{\mathcal{R}})$, we consider it as the native state ${\mathcal{P}}_{1}=({\mathcal{S}}_{1},{\mathcal{R}}_{1})$ and generate a correlated conformer ${\mathcal{P}}_{2}=({\mathcal{S}}_{2},{\mathcal{R}}_{2})$ by randomly perturbing the protein structure. That is, we set ${\mathcal{S}}_{2}$ to be the same as ${\mathcal{S}}_{1}$, and ${\mathcal{R}}_{2}$ is obtained by applying a perturbation function $\text{perturb}({\mathcal{R}}_{1},\epsilon)$ to the original residue structure, where $\epsilon\in[0,2\pi)^{n_{r}\times 4}$ is a noise vector drawn from a wrapped normal distribution [15]. The $\text{perturb}(\cdot,\cdot)$ function rotates the side-chain of each residue according to the sampled torsional noises. To avoid atom clashes, we adjust the variance of the added noise and regenerate any conformers that violate physical constraints. It should be noted that the scheme can be adapted flexibly when considering different granularities of structures. For example, instead of considering side chain rotation on a fixed backbone, we can also consider a flexible backbone by rotating backbone angles, thereby generating approximate conformers.

Although the current torsional perturbation scheme is effective, there is potential for improvement. Future research can explore enhancements such as incorporating available rotamer libraries [61] and introducing backbone flexibility using empirical force fields [71].

To maintain the information-rich joint diffusion trajectories from DiffPreT, we generate trajectories for pairs of conformers, a.k.a., siamese trajectories. We first sample the diffusion trajectories ${\mathcal{P}}^{0:T}_{1}$ and ${\mathcal{P}}^{0:T}_{2}$ for conformers ${\mathcal{P}}_{1}$ and ${\mathcal{P}}_{2}$, respectively, using the joint diffusion process outlined in Sec. 3. For example, starting from ${\mathcal{P}}_{1}^{0}={\mathcal{P}}_{1}=({\mathcal{S}}_{1},{\mathcal{R}}_{1})$, we use the joint diffusion on structures and sequences to define the diffusion process $q({\mathcal{P}}_{1}^{1:T}|{\mathcal{P}}_{1}^{0})=q({\mathcal{R}}_{1}^{1:T}|{\mathcal{R}}_{1}^{0})q({\mathcal{S}}_{1}^{1:T}|{\mathcal{S}}_{1}^{0})$. We derive trajectories on structures ${\mathcal{R}}_{1}^{1:T}$ using the Gaussian noise in (3) and derive the sequence diffusion process ${\mathcal{S}}^{1:T}$ using the random masking in (6). In this way, we define the trajectory ${\mathcal{P}}^{0:T}_{1}=\{({\mathcal{S}}^{t}_{1},{\mathcal{R}}^{t}_{1})\}_{t=0}^{T}$ for ${\mathcal{P}}_{1}$ and can derive the siamese trajectory ${\mathcal{P}}_{2}^{0:T}=\{({\mathcal{S}}^{t}_{2},{\mathcal{R}}^{t}_{2})\}_{t=0}^{T}$ similarly.

We aim to maximize the mutual information (MI) between representations of siamese trajectories constructed in a way that reflects the correlation between different conformers of the same protein. Direct optimization of MI is intractable, so we instead maximize a lower bound. In App. C.1, we show that this problem can be transformed into the minimization of a loss function for mutual denoising between two trajectories: ${\mathcal{L}}=\frac{1}{2}({\mathcal{L}}^{(2\rightarrow 1)}+{\mathcal{L}}^{(1\rightarrow 2)})$, where

with $b\rightarrow a$ being either $2\rightarrow 1$ or $1\rightarrow 2$ and ${\bm{\mathcal{P}}}_{b}^{0:T}$ being representations of the trajectory ${\mathcal{P}}_{b}^{0:T}$.

The two terms share the similar formula as the ELBO loss in (8). Take ${\mathcal{L}}^{(2\rightarrow 1)}$ for example. Here $q({\mathcal{P}}_{1}^{t-1}|{\mathcal{P}}_{1}^{t},{\mathcal{P}}_{1}^{0})$ is a posterior analytically tractable with our definition of each diffusion step $q({\mathcal{P}}_{1}^{t}|{\mathcal{P}}_{1}^{t-1})$ in (3) and (6). The reverse process is learnt to generate a less noisy state ${\mathcal{P}}_{1}^{t-1}$ given the current state ${\mathcal{P}}_{1}^{t}$ and representations of the siamese trajectory ${\bm{\mathcal{P}}}_{2}^{0:T}$, which are extracted by the protein encoder to be pre-trained. The parameterization of the reverse process is similar as in Sec. 3.2, with the representations replaced by those of ${\bm{\mathcal{P}}}_{2}^{0:T}$ (see App. B for details).

Our approach involves mutual prediction between two siamese trajectories, which is similar to the idea of mutual representation reconstruction in [27, 14]. However, since ${\mathcal{P}}_{1}$ and ${\mathcal{P}}_{2}$ share information about the same protein, the whole trajectory of ${\mathcal{P}}_{2}$ could provide too many clues for denoising towards ${\mathcal{P}}_{1}^{t-1}$, making the pre-training task trivial. To address this issue, we parameterize $p({\mathcal{P}}_{1}^{t-1}|{\mathcal{P}}_{1}^{t},{\bm{\mathcal{P}}}_{2}^{0:T})$ with $p_{\theta}({\mathcal{P}}_{1}^{t-1}|{\mathcal{P}}_{1}^{t},{\bm{\mathcal{P}}}_{2}^{t})$. For diffusion on sequences, we further guarantee that the same set of residues are masked in ${\mathcal{S}}_{1}^{t}$ and ${\mathcal{S}}_{2}^{t}$ to avoid leakage of ground-truth residue types across correlated trajectories.

Final pre-training objective. Given the similarity between the one-side objective and the ELBO loss in (8), we can use a similar way to decompose the objective into structure and sequence losses and then derive simplified loss functions for each side. To summarize, the ultimate training objective for our method is

where ${\mathcal{L}}_{\text{simple}}^{(\cdot,b\rightarrow a)}$ is the loss term defined by predicting ${\mathcal{P}}_{a}^{0:T}$ from ${\mathcal{P}}_{b}^{0:T}$ (see App. B for derivation).

Now we discuss the relationship between our method and previous works.

Advantages of joint denoising. Compared with previous diffusion models focusing on either protein sequences [82] or structures [28] or cross-modal contrastive learning [12, 83], in this work, we perform joint diffusion on both modalities. Note that given a sequence ${\mathcal{S}}$ and a structure ${\mathcal{R}}$ that exist in the nature with high probability, the sequence-structure tuple ${\mathcal{P}}=({\mathcal{S}},{\mathcal{R}})$ may not be a valid state of this protein. Consequently, instead of modeling the marginal or conditional distribution, we model the joint distribution of protein sequences and structures.

Connection with diffusion models. Diffusion models excel in image and text generation [16, 48] and have been applied to unsupervised representation learning [1]. Previous works explored denoising objectives [22, 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 [51, 28], which can violate the triangular inequality and produce negative values, our approach directly adds noise at atom coordinates, as demonstrated in Xu et al. [81]. This distinction allows us to address the limitations associated with denoising distance matching algorithms used in molecule and protein generation and pre-training.

Comparison with other deep generative models. Self-supervised learning essentially learns an Energy-Based Model (EBM) for modeling data distribution [47], making VAE [44], GAN [25], and normalizing flow [59] applicable for pre-training. However, these models limit flexibility or fail to acquire high sampling quality and diversity compared to diffusion models [53]. Therefore, we focus on using diffusion models for pre-training and leave other generative models for future work.

Pre-training datasets. Following Zhang et al. [86], we pre-train our models with the AlphaFold protein structure database v1 [43, 69], including 365K proteome-wide predicted structures.

Downstream benchmark tasks. In our evaluation, we assess EC prediction task [23] for catalysis behavior of proteins and four ATOM3D tasks [67]. The EC task involves 538 binary classification problems for Enzyme Commission (EC) numbers. We use dataset splits from Gligorijević et al. [23] with a 95% sequence identity cutoff. The ATOM3D tasks include Protein Interface Prediction (PIP), Mutation Stability Prediction (MSP), Residue Identity (RES), and Protein Structure Ranking (PSR) with different dataset splits based on sequence identity or competition year. Details are in App. D.

Baseline methods. In our evaluation, we utilize GearNet-Edge as the underlying model for both atom- and residue-level structures. GearNet-Edge incorporates various types of edges and edge-type-specific convolutions, along with message passing between edges, to model protein structures effectively. We compare our proposed methods with several previous protein structural pre-training algorithms, including multiview contrastive learning [86], denoising score matching [28], and four self-prediction methods (residue type, distance, angle, and dihedral prediction) [86]. For residue-level tasks, we include EC, MSP, and PSR in our evaluation, while PIP and RES tasks are specifically designed for atom-level models. Besides, we exclude EC from the atom-level evaluation due to the limited presence of side-chain atoms in the downloaded PDB dataset.

Training and evaluation. We pre-train our model for 50 epochs on the AlphaFold protein structure database following Zhang et al. [86] and fine-tune it for 50 epochs on EC, MSP, and PSR. However, due to time constraints, we only fine-tune the models for 10 epochs on the RES and PIP datasets. Results are reported as mean and standard deviation across three seeds (0, 1, and 2). Evaluation metrics include F_max and AUPR for EC, AUROC for PIP and MSP, Spearman’s $\rho$ for PSR, and micro-averaged accuracy for RES. More details about experimental setup can be found in App. D.

Tables 1 and 2 provide a comprehensive overview of the results obtained by GearNet-Edge on both atom- and residue-level benchmark tasks. The tables clearly demonstrate that both DiffPreT and SiamDiff exhibit significant improvements over GearNet-Edge without pre-training on both levels, underscoring the effectiveness of our pre-training methods.

An interesting observation from the tables is that previous pre-training methods tend to excel in specific tasks while showing limitations in others. For instance, Multiview Contrast, designed for capturing similar functional motifs [86], struggles with structural intricacies and local atomic interactions, resulting in lower performance on tasks like Protein Interface Prediction (PIP), Protein Structure Ranking (PSR), and Residue Identity (RES). Self-prediction methods excel at capturing structural details or residue type dependencies but show limitations in function prediction tasks, such as Enzyme Commission (EC) number prediction and Mutation Stability Prediction (MSP), and do not consistently improve performance on both atom and residue levels.

In contrast, our DiffPreT approach achieves top-3 performance in nearly all considered tasks, showcasing its versatility and effectiveness across different evaluation criteria. Moreover, SiamDiff surpasses all other pre-training methods, achieving the best results in 6 out of 7 tasks, establishing it as the state-of-the-art pre-training approach. These results provide compelling evidence that our joint diffusion pre-training strategy successfully captures the intricate interactions between different proteins (PIP), captures local structural details (RES) and global structural characteristics (PSR), and extracts informative features crucial for accurate function prediction (EC) across various tasks.

To analyze the effect of different components of SiamDiff, we perform ablation study on atom-level tasks and present results in Table 3. We first examine two degenerate settings of joint diffusion, i.e., "w/o sequence diffusion" and "w/o structure diffusion". These settings lead to a deterioration in performance across all benchmark tasks, highlighting the importance of both sequence diffusion for residue type identification in RES and structure diffusion for capturing the structural stability of mutation effects in MSP. Next, we compare SiamDiff with DiffPreT, which lacks mutual information maximization between correlated conformers. The consistent improvements observed across all tasks indicate the robustness and effectiveness of our proposed mutual information maximization scheme.

Besides, we compare our method to baselines with fixed small ($T=1$) and large ($T=100$) noise levels to demonstrate the benefits of multi-scale denoising in diffusion pre-training. Interestingly, we observe that denoising with large noise enhances performance on MSP by capturing significant structural changes that lead to structural instability, while denoising with small noise improves performance in PSR by capturing fine-grained details for protein structure assessment. By incorporating multi-scale noise, we eliminate the need for manual tuning of the noise level as a hyperparameter and leverage the advantages of both large- and small-scale noise, as evidenced in the table.

Protein language models (PLMs) have recently become a standard method for extracting representations from protein sequences, such as ESM [49]. However, these methods are unable to directly handle structure-related tasks in Atom3D without using protein structures as input. A recent solution addresses this by feeding residue representations outputted by ESM into the protein structure encoder GearNet [85]. To showcase the potential of SiamDiff on PLM-based encoders, we pre-trained the ESM-GearNet encoder using SiamDiff and evaluated its performance on residue-level tasks. Considering the model capacity and computational budget, we selected ESM-2-650M as the base PLM. The results in Table 4 demonstrate the performance improvements obtained by introducing the PLM component in ESM-GearNet. Furthermore, after pre-training with SiamDiff, ESM-GearNet achieves even better performance on all tasks, especially on PSR where ESM-only representations are not indicative for structure ranking. This highlights the benefits of our method for PLM-based encoders.

In addition, we provide experimental results about pre-training datasets in App. E, multimodal baselines in App. F, different diffusion strategies in App. G, and different backbone models in App. H.