Rigidity-Aware Geometric Pretraining for Protein Design and Conformational Ensembles
Organizations: University of Illinois Urbana-Champaign · University of Washington · MPI for Intelligent Systems, Tübingen · National Research Council of Canada · University of Ottawa · The Chinese University of Hong Kong
Abstract
Generative models have recently advanced protein design by learning the statistical regularities of natural structures. However, current approaches face three key limitations: (1) Existing methods cannot jointly learn protein geometry and design tasks, where pretraining can be a solution; (2) Current pretraining methods mostly rely on local, non-rigid atomic representations for property prediction downstream tasks, limiting global geometric understanding for protein generation tasks; and (3) Existing approaches have yet to effectively model the rich dynamic and conformational information of protein structures. To overcome these issues, we introduce (), a geometric pretraining framework that front-loads geometry learning prior to generative finetuning. Phase I (RigidSSL-Perturb) learns geometric priors from 432K structures from the AlphaFold Protein Structure Database with simulated perturbations. Phase II (RigidSSL-MD) refines these representations on 1.3K molecular dynamics trajectories to capture physically realistic transitions. Underpinning both phases is a bi-directional, rigidity-aware flow matching objective that jointly optimizes translational and rotational dynamics to maximize mutual information between conformations. Empirically, RigidSSL variants improve designability by up to 43% while enhancing novelty and diversity in unconditional generation. Furthermore, RigidSSL-Perturb improves the success rate by 5.8% in zero-shot motif scaffolding and RigidSSL-MD captures more biophysically realistic conformational ensembles in G protein-coupled receptor modeling.
Figures & tables
| Model | Pretraining | Designability | Novelty | Diversity | |
| Fraction ( ) | avg. max TM ( ) | pairwise TM ( ) | MaxCluster ( ) | ||
| FrameDiff | None | 0.033 | |||
| FrameDiff | GeoSSL-EBM-NCE | 0.033 | |||
| FrameDiff | GeoSSL-InfoNCE | 0.033 | |||
| FrameDiff | GeoSSL-RR | ||||
| FrameDiff | RigidSSL-Perturb (ours) | 0.033 | |||
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Stage | Data Source | Data Size | Length | Objective | Compute & Time |
| GeoSSL-InfoNCE | Pretraining | AFDB (UniProtKB) | 432,194 | 60 512 | Contrastive (InfoNCE) | 4 days (4 A100) |
| GeoSSL-EBM-NCE | Pretraining | AFDB (UniProtKB) | 432,194 | 60 512 | Contrastive (EBM-NCE) | 4 days (4 A100) |
| GeoSSL-RR | Pretraining | AFDB (UniProtKB) | 432,194 | 60 512 | Generative (Representation Reconstruction) | 3 days (4 A100) |
| RigidSSL-Perturb (Ours) | Pretraining (Phase I) | AFDB (UniProtKB) | 432,194 | 60 512 | Generative (RigidSSL) | 2.75 days (1 H100) |
| RigidSSL-MD (Ours) | Pretraining (Phase II) | ATLAS | 1,390 traj. | 60 512 | Generative (RigidSSL) | 1.88 days (1 H100) |
| FrameDiff | Finetuning | PDB | 20,312 | 60 512 | SE(3) Diffusion | 7 days (4 H100) |
| Model | Pretraining | FPSD vs. | fS | fJSD vs. | ||
| PDB ( ) | AFDB ( ) | (C/A/T) ( ) | PDB (C/A/T) ( ) | AFDB (C/A/T) ( ) | ||
| FrameDiff | None | 908.743 | 840.682 | 1.052 /1.827/4.249 | 3.174 /3.902/5.454 | 2.209 /3.119/4.807 |
| FrameDiff | GeoSSL-EBM-NCE | 827.919 | 770.448 | 1.108 / 2.243 /6.025 | 3.017 /3.593/5.079 | 2.068 /2.807/4.408 |
| FrameDiff | GeoSSL-InfoNCE | 858.332 | 806.346 | 1.002/ 2.060 / 6.862 | 3.360/3.817/5.195 | 2.369/2.965/4.419 |
| FrameDiff | GeoSSL-RR | 892.657 | 846.222 | 1.002/1.851/ 6.559 | 3.364/ 3.474 / 5.037 | 2.373/ 2.452 / 4.357 |
| FrameDiff | RigidSSL-Perturb (ours) | 901.424 | 843.407 | 1.003/1.735/4.105 | 3.359/3.605/5.417 | 2.370/2.632/4.687 |
| Metric | Length | Unpretrained | GeoSSL-EBM-NCE | GeoSSL-RR | GeoSSL-InfoNCE | RigidSSL-Perturb | RigidSSL-MD |
| Clashscore ( ) | 700 | 152.53 | 45.11 | 96.38 | 55.02 | 21.36 | 37.16 |
| 800 | 172.36 | 61.70 | 123.16 | 74.98 | 26.42 | 39.35 | |
| MolProbity score ( ) | 700 | 2.64 | 2.22 | 2.51 | 2.21 | 1.82 | 2.05 |
| 800 | 2.69 | 2.38 | 2.61 | 2.34 | 1.91 | 2.08 |
| Target | Unpretrained | GeoSSL-EBM-NCE | GeoSSL-InfoNCE | GeoSSL-RR | RigidSSL-Perturb | RigidSSL-MD |
| 1BCF | 11 | 2 | 29 | 0 | 47 | 15 |
| 1PRW | 5 | 1 | 11 | 1 | 11 | 7 |
| 1QJG | 11 | 9 | 17 | 7 | 8 | 10 |
| 1YCR | 9 | 10 | 8 | 5 | 19 | 12 |
| 2KL8 | 15 | 5 | 23 | 14 | 17 | 34 |
| 3IXT | 25 | 22 | 23 | 8 | 25 | 19 |
| Metric | Unpretrained | GeoSSL-EBM-NCE | GeoSSL-InfoNCE | GeoSSL-RR | RigidSSL-Perturb (ours) | RigidSSL-MD (ours) | |
| Predicting flexibility | Pairwise RMSD (=1.55) | 2.37 | 2.32 | 2.31 | 2.48 | 2.20 | 2.66 |
| All-atom RMSF (=1.0) | 1.30 | 1.19 | 1.22 | 1.17 | 1.08 | 1.28 | |
| Distributional accuracy | Root mean -dist. | 18.36 | 18.46 | 17.79 | 17.83 | 17.90 | 18.10 |
| Translation contrib. | 18.26 | 18.31 | 17.75 | 17.76 | 17.83 | 18.05 | |
| Variance contrib. | 1.45 | 1.50 | 1.42 | 1.52 | 1.44 | 1.60 | |
| MD PCA -dist. | 3.15 | 2.97 | 3.16 | 3.14 | 3.10 | 3.22 |
| Designability | Novelty | Diversity | |||||
| Model | Downstream Dataset | Pretraining | Steps | Fraction ( ) | avg. max TM ( ) | pairwise TM ( ) | MaxCluster ( ) |
| FoldFlow-2 | PDB + AFDB | None | 500k | ||||
| FoldFlow-2 | PDB | RigidSSL-Perturb | 400k | ||||
| Designability | Novelty | Diversity | |||
| Translation Noise | Rotation Noise | Length | Fraction ( ) | avg. max TM ( ) | pairwise TM ( ) |
| 0.01 | 0.5 | 100–600 | |||
| 0.03 | 0.5 | 100–600 | |||
| 0.05 | 0.5 | 100–600 | |||
| 0.5 | 0.5 | 100–600 | |||
| 0.5 | 0.75 | 100–600 |