[Paper Review] Pre-Training Protein Encoder via Siamese Sequence-Structure Diffusion Trajectory Prediction

Zuobai Zhang, Minghao Xu+4 more|arXiv (Cornell University)|Jan 28, 2023

Machine Learning in Bioinformatics11 citations

TL;DR

The paper introduces DiffPreT to pre-train protein encoders via joint diffusion on sequence and structure, and SiamDiff to capture conformer correlations with siamese diffusion trajectories, achieving state-of-the-art mean ranks on downstream tasks.

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. Our implementation is available at https://github.com/DeepGraphLearning/SiamDiff.

Motivation & Objective

• Model the joint distribution of protein sequences and structures using diffusion to learn informative representations.
• Incorporate conformational heterogeneity with Siamese diffusion trajectory prediction to capture correlations between conformers.
• Evaluate the pre-trained encoders on atom- and residue-level structure-based tasks to assess generalization and robustness.
• Demonstrate that the proposed methods are competitive or superior to existing pre-training approaches across diverse downstream tasks.

Proposed method

• DiffPreT performs joint diffusion on protein sequences and structures and trains a noise-prediction network parameterized by a protein encoder.
• The forward diffusion adds Gaussian noise to 3D coordinates and applies masked token diffusion to sequences; the reverse uses SE(3)-invariant/equivariant representations.
• The objective is an ELBO decomposed into structure and sequence losses, simplified to a regression loss on noise for structure and cross-entropy for sequence.
• SiamDiff generates pairs of correlated conformers via torsional perturbations and maximizes mutual information between their diffusion trajectories by mutual denoising.
• Two-stage noise scheduling (large-noise first, then small-noise) is employed to capture coarse- and fine-grained features.
• The final objective combines four loss terms, enforcing cross-trajectory denoising between both conformers and both modalities.

Experimental results

Research questions

• RQ1 Can a joint sequence-structure diffusion model learn representations that capture both co-evolutionary sequence information and detailed structural characteristics?
• RQ2 Does incorporating conformer-aware SiamDiff improve downstream protein understanding tasks beyond DiffPreT and existing baselines?
• RQ3 How does a two-stage noise schedule affect learning of multi-modal protein representations?
• RQ4 What is the impact of siamese diffusion trajectories on capturing conformational correlations across tasks (atom- and residue-level)?

Key findings

• DiffPreT delivers competitive performance across atom- and residue-level tasks.
• SiamDiff achieves new state-of-the-art mean ranks across evaluated tasks.
• Atom-level: SiamDiff Mean Rank = 1.2 vs DiffPreT = 2.4.
• Residue-level: SiamDiff Mean Rank = 1.2 vs DiffPreT = 3.2.
• Overall, SiamDiff improves mean ranks and demonstrates strong cross-task generalization.